max, min, argmax, argmin break on WGPU for floats > 2^15 - 1

[TomWyllie]

Describe the bug
The values returned for tensor.min(), tensor.max(), tensor.argmin(), tensor.argmax() on WGPU all break for floating tensors where the true value is outside the range of [-32767, 32767].

To Reproduce

let x: Tensor<B, 1, Float> = Tensor::from_floats([32768., 32767.], &B::Device::default());
let min_or_max: Tensor<B, 1> = x.clone().min();
println!("x:\t{:?}", x.clone().into_data());
println!("min:\t{:?}", min_or_max.clone().into_data());

Works on NdArray, gives wrong values on Wgpu backend.

Expected behavior
The minimum of the above tensor is 32767.0, not 32768.0. This also breaks if using negatives and max instead of min, i.e. from [-32768., -32767.], -32768.0 will be incorrectly returned.

Desktop (please complete the following information):

• OS: Ubuntu 20.04
• Version: 13.0 (also true in 12.1)

[TomWyllie]

I've done some digging and have found the problem. In the previous hand written shaders in 12.1, and in the JIT shaders for 13.0, there's a few places where the tracker for the loop equivalent of what might be f32::MAX or f32::MIN is replaced with 32767 or -32767, which is already a better running min/max than the true value. These are initialised at the extreme values for i16, is there a reason that the min or max has to be within this range?

./burn-jit/src/kernel/reduce/argmax_dim.rs:17:        gpu!(scope, max = cast(-32767.0));
./burn-jit/src/kernel/reduce/argmax_dim.rs:54:        gpu!(scope, max = cast(-32767.0));
./burn-jit/src/kernel/reduce/argmin_dim.rs:17:        gpu!(scope, min = cast(32767.0));
./burn-jit/src/kernel/reduce/argmin_dim.rs:54:        gpu!(scope, min = cast(32767.0));```