ops: embedding/reshape/transpose/split-merge-heads fwd+bwd

Phase T5 structural ops on top of the T4 set, needed to assemble the
tiny transformer:
- embedding: gather rows by I32 ids (CUDA kernel) / scatter-add backward
  (atomic, so repeated ids accumulate). csrc/ops/model.cu + ffi.
- reshape: contiguous metadata-only view (Tensor::reshape), no kernel.
- transpose_3d01: [a,b,c]->[b,a,c] for the multi-head layout (kernel).
- autograd nodes: embedding/reshape/transpose_3d01/transpose_2d, plus
  split_heads (->Vec<Var>) / merge_heads for per-head attention.
- tape: Var::zero_grad + set_value so a hand-written GD step can update
  params and clear grads between steps.

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

crates/xtrain-tensor/src/tensor.rs

@@ -563,6 +563,98 @@ impl Tensor {
         dx
     }
 
+    // --- Structural / model ops (the T5 kernels) ---
+
+    /// Reshape to `new_shape` (must keep `numel`). Pure metadata change on a
+    /// contiguous tensor — no data movement, shares the same storage. The
+    /// multi-head layout `[seq, n_heads*head_dim] <-> [seq, n_heads, head_dim]`
+    /// is exactly this.
+    pub fn reshape(&self, new_shape: &[usize]) -> Self {
+        assert!(self.is_contiguous(), "reshape requires a contiguous tensor");
+        assert_eq!(
+            shape::num_elements(new_shape),
+            self.numel(),
+            "reshape numel mismatch: {:?} -> {:?}",
+            self.shape.as_slice(),
+            new_shape
+        );
+        Self {
+            storage: self.storage.clone(),
+            shape: Dims::from_slice(new_shape),
+            strides: shape::contiguous_strides(new_shape),
+            offset: self.offset,
+            dtype: self.dtype,
+        }
+    }
+
+    /// Embedding gather: `out[s,:] = self[ids[s], :]`. `self`:[vocab,dim] table,
+    /// `ids`:[seq] I32 → out:[seq,dim].
+    #[cfg(not(no_cuda))]
+    pub fn embedding(&self, ids: &Tensor) -> Self {
+        assert_eq!(self.dtype, DType::F32, "embedding table must be F32");
+        assert_eq!(self.ndim(), 2, "embedding table must be [vocab,dim]");
+        assert_eq!(ids.dtype, DType::I32, "embedding ids must be I32");
+        assert_eq!(ids.ndim(), 1, "embedding ids must be 1D");
+        let (seq, dim) = (ids.shape[0], self.shape[1]);
+        let out = Tensor::zeros(&[seq, dim], DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_embedding_fwd_f32(
+                self.data_ptr() as *const f32,
+                ids.data_ptr() as *const i32,
+                out.data_ptr() as *mut f32,
+                seq as i32,
+                dim as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("embedding sync failed");
+        out
+    }
+
+    /// Embedding backward (scatter-add): `dtable[ids[s],:] += dout[s,:]`, where
+    /// `dout`:[seq,dim], `ids`:[seq] I32. `vocab` sizes the output table.
+    #[cfg(not(no_cuda))]
+    pub fn embedding_backward(dout: &Tensor, ids: &Tensor, vocab: usize) -> Self {
+        let (seq, dim) = (dout.shape[0], dout.shape[1]);
+        let dtable = Tensor::zeros(&[vocab, dim], DType::F32, dout.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_embedding_bwd_f32(
+                dout.data_ptr() as *const f32,
+                ids.data_ptr() as *const i32,
+                dtable.data_ptr() as *mut f32,
+                seq as i32,
+                dim as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("embedding_backward sync failed");
+        dtable
+    }
+
+    /// 3D axis-(0,1) transpose: `self`:[a,b,c] → [b,a,c], `out[j,i,k]=self[i,j,k]`.
+    /// Lays out multi-head attention (`[seq,heads,hd] <-> [heads,seq,hd]`). Its
+    /// own backward is the same op (swap a,b).
+    #[cfg(not(no_cuda))]
+    pub fn transpose_3d01(&self) -> Self {
+        assert_eq!(self.dtype, DType::F32, "transpose_3d01 only supports F32");
+        assert_eq!(self.ndim(), 3, "transpose_3d01 requires a 3D tensor");
+        assert!(self.is_contiguous(), "transpose_3d01 requires contiguous");
+        let (a, b, c) = (self.shape[0], self.shape[1], self.shape[2]);
+        let out = Tensor::zeros(&[b, a, c], DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_transpose_3d01_f32(
+                self.data_ptr() as *const f32,
+                out.data_ptr() as *mut f32,
+                a as i32,
+                b as i32,
+                c as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("transpose_3d01 sync failed");
+        out
+    }
+
     // Shared validation for same-shape binary elementwise ops.
     #[cfg(not(no_cuda))]
     fn check_binary(&self, other: &Tensor, op: &str) {