burn-ndarray: relu and clamp drop NaN once a tensor is long enough to vectorize

Describe the bug

On burn-ndarray, relu, clamp_min, clamp_max and clamp replace NaN with the bound as soon as a tensor is long enough to take the SIMD path. Short tensors fall through to the scalar comparison below it and keep NaN, so the result depends on tensor length.

This is the same class of bug as #5609 and #5659, but it reaches further: relu is affected, not just clamp, and any tensor of a realistic size is on the broken path.

To Reproduce

One NaN at index 0, every other element 2.0, --features ndarray:

LEN     3  clamp_min(0)=NaN  clamp_max(1)=NaN  clamp(0,1)=NaN  relu=NaN
LEN     8  clamp_min(0)=NaN  clamp_max(1)=NaN  clamp(0,1)=NaN  relu=NaN
LEN    16  clamp_min(0)=NaN  clamp_max(1)=NaN  clamp(0,1)=NaN  relu=NaN
LEN    64  clamp_min(0)=0.0  clamp_max(1)=1.0  clamp(0,1)=1.0  relu=0.0
LEN  1024  clamp_min(0)=0.0  clamp_max(1)=1.0  clamp(0,1)=1.0  relu=0.0

The switch is wherever should_use_simd starts returning true, between 16 and 64 f32 elements here.

Analysis

dispatch_binary_scalar_simd! returns early when a SIMD implementation exists (Ok(out) => return out), so the mapv_inplace comparison underneath it never runs for f32 or f64. macerator lowers VecMax, VecMin and VecClamp to the hardware max, min and clamp, which return the non-NaN operand. relu goes through NdArrayMathOps::clamp_min, so it is hit as well.

crates/burn-ndarray/src/ops/base.rs:1357 (clamp_min), :1383 (clamp_max), :1409 (clamp), crates/burn-ndarray/src/ops/simd/binary_elemwise.rs:119 onwards.

Worth correcting: the backend table in #5658 lists ndarray as passing for relu, clamp_min and clamp_max. That measurement used three-element tensors and only exercised the scalar path. The assertions added there have the same limitation.

Expected behavior

NaN survives regardless of tensor length, matching PyTorch and the scalar path.

Candidate fixes

1. Put the bound first in the vector ops. The hardware instruction returns its second operand when either is NaN, so splat.max(data) keeps the element where data.max(splat) does not. Verified against macerator directly on x86:

   复制

   data.max(splat) lane0 = 0.0   nan_kept=false
   splat.max(data) lane0 = NaN   nan_kept=true

   The scalar tail still needs its own comparison, since f32::max ignores NaN in either order. Free, but it depends on behaviour macerator does not promise: VOrd::vmax is documented as Ord::max, which says nothing about NaN for floats. ARM FMAX and wasm f32x4.max propagate NaN in both orders, so the swap is neutral there, but only x86 has been checked here.

2. Drop f32 and f64 from the SIMD dispatch list so floats take the comparison path. Depends on nothing beyond the comparison operators, but it costs real throughput.

3. Ask macerator for a NaN-propagating min/max, or a select for floats so a blend can be written here. Cleanest, since the semantics would then be documented rather than inferred, but it needs an upstream release and a version bump.

Option 1 is implemented in https://github.com/Liberxue/burn/tree/fix/ndarray-vectorized-clamp-nan; option 2 was measured and discarded. Happy to switch to whichever the maintainers prefer, or to open the macerator issue for option 3.

Measurements

4M f32, best of 7 runs of 20 iterations, Ubuntu 26.04, Xeon E5-2686 v4. Baseline drifted under 1% between the first and last group.

Option 1 sits inside the run-to-run spread of the baseline.

Desktop

• Behaviour: macOS 15.7.9, x86_64, Core i7-8559U, rustc 1.96.0
• Measurements: Ubuntu 26.04 LTS, x86_64, Xeon E5-2686 v4, rustc 1.97.1
• burn: 52c101f45

Additional context

Any regression test for this has to use a tensor long enough to vectorize. Three elements pass on the broken code.

Tensor-tensor min/max in ops/simd/binary.rs look like they have the same split between the vector path and the scalar tail. Not investigated, and left alone by the fix above, which adds separate NaN-propagating methods rather than changing MinMax::min and MinMax::max.