Subtraction producing NaNs (f16-related)

**Describe the bug** A subtraction where both the lhs and rhs contain no INFs nor NaNs result in a NaN. **To Reproduce** ```toml # Cargo.toml [package] name = "burn-testing" version = "0.1.0" edition = "2024" [dependencies.burn] version = "0.21.0" default-features = false features = ["cuda", "std"] ``` ```rs // main.rs use burn::prelude::*; type CudaF16 = burn::backend::Cuda<burn::tensor::f16, i32>; type CudaF32 = burn::backend::Cuda<f32, i32>; fn main() { println!("infer cuda f16 (test 1) - smaller range and bigger shape"); infer::<CudaF16>(-13.46875 , -0.0056266785, [12, 13, 4, 64, 64]); // OK println!("infer cuda f16 (test 2) - bigger range and smaller shape"); infer::<CudaF16>(-59.03125, -0.001540184, [1, 13, 4, 64, 64]); // OK println!("infer cuda f32 (test 3) - bigger range and bigger shape, at f32"); infer::<CudaF32>(-59.03125, -0.001540184, [12, 13, 4, 64, 64]); // NaN @ res println!("infer cuda f16 (test 4) - bigger range and bigger shape, at f16"); infer::<CudaF16>(-59.03125, -0.001540184, [12, 13, 4, 64, 64]); // NaN @ res println!("finish (OK)"); } pub fn infer<B: Backend>(start: f32, end: f32, shape: [usize; 5]) { let device = <B::Device>::default(); let range = Tensor::<B, 1, Int>::arange(0..64, &device).float(); // 0 ~ 63 let range = range / (63); // 0.0 ~ 1.0 let base = Tensor::<B, 1>::ones_like(&range) * (end / start); let base = base.powf(range) * (-start); // start..end (logspace) sanity(&base, "base"); let base: Tensor<B, 4> = base.unsqueeze_dims(&[0, 1, 2]); // [1, 1, 1, 64] let base = base.expand([shape[0], shape[1], shape[2], shape[3]]); // e.g. [12, 13, 4, 64] let lhs: Tensor<B, 5> = base.clone().unsqueeze_dim(4); // e.g. [12, 13, 4, 64, 1] let rhs: Tensor<B, 5> = base.unsqueeze_dim(3); // e.g. [12, 13, 4, 1, 64] let lhs = lhs.expand(shape); let rhs = rhs.expand(shape); sanity(&lhs, "lhs"); sanity(&rhs, "rhs"); let res = lhs - rhs; sanity(&res, "res"); } pub const DENY_NAN: bool = true; pub const DENY_INF: bool = true; pub fn sanity<B: Backend, const D: usize>(t: &Tensor<B, D>, target: &str) { let mut has_nan = false; let mut has_inf = false; if DENY_NAN { has_nan = t.clone().contains_nan().into_scalar().to_bool(); if has_nan { eprintln!("got a NaN"); } } if DENY_INF { has_inf = t.clone().is_inf().any().into_scalar().to_bool(); if has_inf { eprintln!("got a INF"); } } if has_nan || has_inf { panic!("sanity check failed for {target}"); } } ``` **Expected behavior** No INFs nor NaNs (no panics). **Desktop (please complete the following information):** - OS: Linux 6.2. - GPU: RTX 2060, CUDA 12.2. **Additional context** - Some experiments notes: - From the default setup, `test 3` (F32) panics (NaN at `res`). If I comment `test 3` out and let `test 4` run (which is same as `test 3` but in F16), then `test 4` it also panics (NaN at `res`). - From the default setup, if I change `test 1` and `test 2` to run in F32, then all tests pass, including `test 3` and `test 4`. This raises a suspicion that the problem relates to UB from previous F16 operations. - From the default setup, if I repeat `test 1` and `test 2` 4 times each, all of those tests passes, only `test 3` (and `test 4`) fail. So I _think_ the value ranges and the used shapes *are* important. - If I make lhs and rhs come from completely different tensors (one completely new and separate range and a base for each), the error still persist. After doing this, also enforcing each side to be contiguous (as shown below) right before they're used for the subtraction also has no effect. ```rs let lhs = lhs.into_data(); let lhs: Tensor<B, 5> = Tensor::from_data(lhs, &device); let rhs = rhs.into_data(); let rhs: Tensor<B, 5> = Tensor::from_data(rhs, &device); ```