Checklist

• Confirmed that run-checks all script has been executed.
• Made sure the book is up to date with changes in this PR.

Related Issues/PRs

Follow-up of #2483

Changes

Added float_cast implementation for burn-jit, burn-fusion and burn-router.

• Removed elem type generic from JitTensor
  • JitTensor now contains the dtype field
  • Refactored all calls with generics to have the correct generic types propagation (new execute_with_dtype macro to call cubecl kernels)
• Added Primitive trait for backend tensor primitives
  • Motivation: all primitives previously had the same trait requirements, and now we need to report the correct dtype in the Display implementation for Tensor (K::elem_type_name() would report the default floating point dtype even for tensors cast to another precision)
• Added dtype checks for binary float ops (no automatic type promotion for now, so we panic for different types)

Known issue
Trying to concatenate tensors cast to a different floating point dtype will fail with the current cat_with_slice_assign implementation because the output tensor is initialized with the default floating point element precision in K::empty(...).

let tensor1: Tensor<B, 2> = Tensor::from_floats([[2., 3.], [4., 5.]], &device).cast(DType::F16);
let tensor2 = Tensor::ones_like(&tensor1).cast(DType::F16);
let tensor3 = Tensor::zeros_like(&tensor1).cast(DType::F16);

// Panics with incompatible dtype (for output)
let output = Tensor::<B, 2>::cat(vec![tensor1, tensor2, tensor3], 0);

TODO

• Implement float_cast for ndarray
• Remove redundant precision bridge

Testing

Added new tests for primitive trait.