autodiff: bf16 mixed-precision path (fp32 master via cast op)

Tensor ops dispatch on dtype: fp32 branch unchanged (bit-identical),
bf16 branch routes matmul/attention through GemmEx and elementwise
through the bf16 kernels. Norm/softmax/RoPE/cross-entropy upcast to
fp32 around the existing fp32 kernels (standard AMP: reductions/loss
fp32, matmuls bf16). Transposes route bf16 through fp32 (pure layout).

New autodiff `cast` op is the AMP bridge: forward downcasts a fp32
master leaf to bf16 for the matmul; backward upcasts the bf16 grad
back to fp32. So the fp32 leaf accumulates an fp32 grad and AdamW /
clip / DDP all-reduce stay fp32 and completely unchanged.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xtrain-autodiff/src/ops.rs

@@ -13,7 +13,27 @@
 #![cfg(not(no_cuda))]
 
 use crate::tape::Var;
-use xtrain_tensor::Tensor;
+use xtrain_tensor::{DType, Tensor};
+
+/// dtype cast as an autograd node (Phase T12 — the AMP bridge between fp32 master
+/// weights / fp32 reductions and the bf16 compute stream). Forward casts `x` to
+/// `target`; **backward casts the upstream grad back to `x`'s dtype**. So a fp32
+/// master-weight leaf fed through `cast(w, BF16)` into a bf16 matmul accumulates
+/// an **fp32** grad — AdamW / clip / DDP all-reduce stay fp32, untouched.
+pub fn cast(x: &Var, target: DType) -> Var {
+    let src = x.value().dtype();
+    if src == target {
+        return x.clone();
+    }
+    let out = x.value().to_dtype(target);
+    Var::from_op(
+        out,
+        vec![x.clone()],
+        Box::new(move |d, parents| {
+            Var::push_grad(&parents[0], d.to_dtype(src));
+        }),
+    )
+}
 
 /// `C = A @ B` (2D). Backward: `dA = dC @ Bᵀ`, `dB = Aᵀ @ dC`.
 pub fn matmul(a: &Var, b: &Var) -> Var {