//! The `Tensor` type: shape/strides/dtype over reference-counted [`Storage`],
//! with host↔device transfer and one elementwise op (`scale`) wired end-to-end
//! through a CUDA kernel.

use crate::dtype::{DType, TensorDType};
use crate::shape::{self, Dims};
use crate::storage::{Device, Storage};

/// Multi-dimensional array backed by CPU or GPU storage.
///
/// Strides are in elements (row-major). T2 tensors created here are always
/// contiguous; the `strides`/`offset` fields exist so later phases can add
/// zero-copy views without changing this type's shape.
#[derive(Clone)]
pub struct Tensor {
    storage: Storage,
    shape: Dims,
    strides: Dims,
    offset: usize,
    dtype: DType,
}

impl Tensor {
    // --- Creation ---

    /// Build a contiguous CPU tensor from a typed host slice.
    pub fn from_slice<T: TensorDType>(data: &[T], shape: &[usize]) -> Self {
        let numel = shape::num_elements(shape);
        assert_eq!(
            data.len(),
            numel,
            "data length {} != shape numel {numel}",
            data.len()
        );
        let bytes = unsafe {
            std::slice::from_raw_parts(data.as_ptr() as *const u8, numel * T::DTYPE.size_bytes())
        };
        Self {
            storage: Storage::cpu(bytes.to_vec()),
            shape: Dims::from_slice(shape),
            strides: shape::contiguous_strides(shape),
            offset: 0,
            dtype: T::DTYPE,
        }
    }

    /// Zero-filled contiguous tensor on the given device.
    pub fn zeros(shape: &[usize], dtype: DType, device: Device) -> Self {
        let len_bytes = shape::num_elements(shape) * dtype.size_bytes();
        let storage = Storage::zeros(len_bytes, device).expect("zeros alloc failed");
        Self {
            storage,
            shape: Dims::from_slice(shape),
            strides: shape::contiguous_strides(shape),
            offset: 0,
            dtype,
        }
    }

    // --- Properties ---

    pub fn shape(&self) -> &[usize] {
        &self.shape
    }
    pub fn strides(&self) -> &[usize] {
        &self.strides
    }
    pub fn dtype(&self) -> DType {
        self.dtype
    }
    pub fn ndim(&self) -> usize {
        self.shape.len()
    }
    pub fn numel(&self) -> usize {
        shape::num_elements(&self.shape)
    }
    pub fn offset(&self) -> usize {
        self.offset
    }
    pub fn device(&self) -> Device {
        self.storage.device()
    }
    pub fn is_contiguous(&self) -> bool {
        shape::is_contiguous(&self.shape, &self.strides)
    }
    pub fn storage(&self) -> &Storage {
        &self.storage
    }

    // --- Device transfer ---

    /// Move (copy) the tensor to `device`. Returns a cheap clone if already there.
    pub fn to_device(&self, device: Device) -> Self {
        if self.device() == device {
            return self.clone();
        }
        let storage = self
            .storage
            .to_device(device)
            .expect("device transfer failed");
        Self {
            storage,
            shape: self.shape.clone(),
            strides: self.strides.clone(),
            offset: self.offset,
            dtype: self.dtype,
        }
    }

    // --- dtype cast (Phase T12, bf16 mixed precision) ---

    /// Cast between F32 and BF16 (the AMP bridge: fp32 master ↔ bf16 compute).
    /// Same dtype returns a cheap clone. Requires a contiguous CUDA tensor.
    /// I32 is not castable here (only used for token-id targets).
    #[cfg(not(no_cuda))]
    pub fn to_dtype(&self, target: DType) -> Self {
        if self.dtype == target {
            return self.clone();
        }
        assert!(
            matches!(self.device(), Device::Cuda(_)),
            "to_dtype requires a CUDA tensor"
        );
        assert!(self.is_contiguous(), "to_dtype requires contiguous tensor");
        let n = self.numel() as i32;
        let out = Tensor::zeros(&self.shape, target, self.device());
        match (self.dtype, target) {
            (DType::F32, DType::BF16) => unsafe {
                xtrain_cuda::ffi::launch_cast_f32_to_bf16(
                    self.data_ptr() as *const f32,
                    out.data_ptr() as *mut std::ffi::c_void,
                    n,
                    std::ptr::null_mut(),
                );
            },
            (DType::BF16, DType::F32) => unsafe {
                xtrain_cuda::ffi::launch_cast_bf16_to_f32(
                    self.data_ptr() as *const std::ffi::c_void,
                    out.data_ptr() as *mut f32,
                    n,
                    std::ptr::null_mut(),
                );
            },
            (a, b) => panic!("unsupported cast {a} -> {b}"),
        }
        out
    }

    // --- Host data access (CPU only) ---

    /// Typed read-only view of the data. Requires a contiguous CPU tensor.
    pub fn as_slice<T: TensorDType>(&self) -> &[T] {
        assert_eq!(T::DTYPE, self.dtype, "dtype mismatch");
        assert_eq!(self.device(), Device::Cpu, "as_slice requires CPU tensor");
        assert!(self.is_contiguous(), "as_slice requires contiguous tensor");
        let bytes = self.storage.as_cpu_bytes();
        let start = self.offset * self.dtype.size_bytes();
        unsafe { std::slice::from_raw_parts(bytes[start..].as_ptr() as *const T, self.numel()) }
    }

    /// Raw element pointer at the tensor's offset (for kernel launches).
    pub fn data_ptr(&self) -> *const u8 {
        let byte_off = self.offset * self.dtype.size_bytes();
        match self.device() {
            Device::Cpu => unsafe { self.storage.as_cpu_bytes().as_ptr().add(byte_off) },
            Device::Cuda(_) => unsafe { self.storage.gpu_buffer().as_ptr().add(byte_off) },
        }
    }

    // --- Elementwise op (the T2 end-to-end kernel) ---

    /// Out-of-place elementwise scale: returns a new tensor `out[i] = self[i] * alpha`.
    ///
    /// Runs the `scale_f32` CUDA kernel. Requires a contiguous F32 tensor on the
    /// GPU. Available only when CUDA was compiled in (`not(no_cuda)`).
    #[cfg(not(no_cuda))]
    pub fn scale(&self, alpha: f32) -> Self {
        assert!(
            matches!(self.dtype, DType::F32 | DType::BF16),
            "scale supports F32/BF16"
        );
        assert!(self.is_contiguous(), "scale requires contiguous tensor");
        assert!(
            matches!(self.device(), Device::Cuda(_)),
            "scale requires a CUDA tensor"
        );

        let out = Tensor::zeros(&self.shape, self.dtype, self.device());
        let n = self.numel() as i32;
        match self.dtype {
            DType::F32 => unsafe {
                xtrain_cuda::ffi::launch_scale_f32(
                    self.data_ptr() as *const f32,
                    out.data_ptr() as *mut f32,
                    alpha,
                    n,
                    std::ptr::null_mut(), // default stream
                );
            },
            DType::BF16 => unsafe {
                xtrain_cuda::ffi::launch_scale_bf16(
                    self.data_ptr() as *const std::ffi::c_void,
                    out.data_ptr() as *mut std::ffi::c_void,
                    alpha,
                    n,
                    std::ptr::null_mut(),
                );
            },
            _ => unreachable!(),
        }
        out
    }

    // --- GEMM (the T3 kernels) ---

    /// Matrix multiply: `C = self @ other`. `self`:[M,K], `other`:[K,N] → [M,N].
    ///
    /// Routes through cuBLAS `Sgemm` (Phase T7). fp32, so it is the same GEMM as
    /// the T3 tiled kernel up to rounding order. Requires contiguous F32 tensors
    /// on the same GPU. Available only when CUDA is compiled in.
    #[cfg(not(no_cuda))]
    pub fn matmul(&self, other: &Tensor) -> Self {
        assert_eq!(self.dtype, other.dtype, "matmul dtype mismatch");
        assert!(
            matches!(self.dtype, DType::F32 | DType::BF16),
            "matmul supports F32/BF16"
        );
        assert_eq!(self.ndim(), 2, "matmul requires 2D lhs");
        assert_eq!(other.ndim(), 2, "matmul requires 2D rhs");
        assert_eq!(
            self.shape[1], other.shape[0],
            "inner dimension mismatch: {:?} @ {:?}",
            self.shape, other.shape
        );
        assert!(
            self.is_contiguous() && other.is_contiguous(),
            "matmul requires contiguous tensors"
        );
        assert_eq!(self.device(), other.device(), "matmul device mismatch");
        assert!(
            matches!(self.device(), Device::Cuda(_)),
            "matmul requires CUDA tensors"
        );

        let m = self.shape[0];
        let k = self.shape[1];
        let n = other.shape[1];
        let out = Tensor::zeros(&[m, n], self.dtype, self.device());
        match self.dtype {
            // fp32 path — unchanged (bit-identical to T7/T10/T11).
            DType::F32 => xtrain_cuda::cublas::sgemm(
                false,
                false,
                m,
                n,
                k,
                1.0,
                self.data_ptr() as *const f32,
                other.data_ptr() as *const f32,
                0.0,
                out.data_ptr() as *mut f32,
            ),
            // bf16 path — GemmEx, bf16 in/out, fp32 accumulation.
            DType::BF16 => xtrain_cuda::cublas::gemm_ex(
                false,
                false,
                m,
                n,
                k,
                1.0,
                self.data_ptr() as *const std::ffi::c_void,
                other.data_ptr() as *const std::ffi::c_void,
                0.0,
                out.data_ptr() as *mut std::ffi::c_void,
            ),
            _ => unreachable!(),
        }
        out
    }

    /// Out-of-place 2D transpose: returns a new contiguous tensor `out[j,i] =
    /// self[i,j]`. Requires a contiguous F32 CUDA tensor.
    #[cfg(not(no_cuda))]
    pub fn transpose_2d(&self) -> Self {
        assert_eq!(self.ndim(), 2, "transpose_2d requires 2D tensor");
        assert!(self.is_contiguous(), "transpose requires contiguous tensor");
        if self.dtype == DType::BF16 {
            return self
                .to_dtype(DType::F32)
                .transpose_2d()
                .to_dtype(DType::BF16);
        }
        assert_eq!(self.dtype, DType::F32, "transpose supports F32/BF16");
        assert!(
            matches!(self.device(), Device::Cuda(_)),
            "transpose requires a CUDA tensor"
        );

        let rows = self.shape[0];
        let cols = self.shape[1];
        let out = Tensor::zeros(&[cols, rows], DType::F32, self.device());
        unsafe {
            xtrain_cuda::ffi::launch_transpose_f32(
                self.data_ptr() as *const f32,
                out.data_ptr() as *mut f32,
                rows as i32,
                cols as i32,
                std::ptr::null_mut(),
            );
        }
        out
    }

    /// Backward of `C = A @ B` given the upstream gradient `dC` (shape [M,N]).
    /// Returns `(dA, dB)` where `dA = dC @ Bᵀ` ([M,K]) and `dB = Aᵀ @ dC`
    /// ([K,N]). All tensors contiguous F32 on the same GPU.
    ///
    /// Phase T7: cuBLAS applies the transposes internally via its op flags, so we
    /// avoid the two transpose kernels (and their allocations) the T3 version ran.
    #[cfg(not(no_cuda))]
    pub fn matmul_backward(a: &Tensor, b: &Tensor, dc: &Tensor) -> (Tensor, Tensor) {
        assert_eq!(a.ndim(), 2, "matmul_backward requires 2D A");
        assert_eq!(b.ndim(), 2, "matmul_backward requires 2D B");
        assert_eq!(dc.ndim(), 2, "matmul_backward requires 2D dC");
        assert_eq!(a.shape[1], b.shape[0], "A/B inner dim mismatch");
        assert_eq!(dc.shape[0], a.shape[0], "dC rows != A rows (M)");
        assert_eq!(dc.shape[1], b.shape[1], "dC cols != B cols (N)");

        assert_eq!(a.dtype, b.dtype, "matmul_backward dtype mismatch");
        assert_eq!(a.dtype, dc.dtype, "matmul_backward dtype mismatch");
        let (m, k, n) = (a.shape[0], a.shape[1], b.shape[1]);
        let dt = a.dtype;
        // dA[M,K] = dC[M,N] · Bᵀ  (B stored [K,N], transposed by cuBLAS)
        let da = Tensor::zeros(&[m, k], dt, a.device());
        // dB[K,N] = Aᵀ · dC[M,N]  (A stored [M,K], transposed by cuBLAS)
        let db = Tensor::zeros(&[k, n], dt, a.device());
        match dt {
            DType::F32 => {
                xtrain_cuda::cublas::sgemm(
                    false,
                    true,
                    m,
                    k,
                    n,
                    1.0,
                    dc.data_ptr() as *const f32,
                    b.data_ptr() as *const f32,
                    0.0,
                    da.data_ptr() as *mut f32,
                );
                xtrain_cuda::cublas::sgemm(
                    true,
                    false,
                    k,
                    n,
                    m,
                    1.0,
                    a.data_ptr() as *const f32,
                    dc.data_ptr() as *const f32,
                    0.0,
                    db.data_ptr() as *mut f32,
                );
            }
            DType::BF16 => {
                xtrain_cuda::cublas::gemm_ex(
                    false,
                    true,
                    m,
                    k,
                    n,
                    1.0,
                    dc.data_ptr() as *const std::ffi::c_void,
                    b.data_ptr() as *const std::ffi::c_void,
                    0.0,
                    da.data_ptr() as *mut std::ffi::c_void,
                );
                xtrain_cuda::cublas::gemm_ex(
                    true,
                    false,
                    k,
                    n,
                    m,
                    1.0,
                    a.data_ptr() as *const std::ffi::c_void,
                    dc.data_ptr() as *const std::ffi::c_void,
                    0.0,
                    db.data_ptr() as *mut std::ffi::c_void,
                );
            }
            _ => panic!("matmul_backward supports F32/BF16"),
        }
        (da, db)
    }

    // --- Transformer / autograd op primitives (the T4 kernels) ---
    //
    // Each is a thin, contiguous-F32-on-GPU wrapper over a kernel in
    // csrc/ops/nn.cu. The autograd `Var` layer (xtrain-autodiff) builds nodes on
    // top of these; the analytic backwards are derived in docs/03-autograd-engine.md.

    /// Elementwise `out = self + other` (same shape).
    #[cfg(not(no_cuda))]
    pub fn add(&self, other: &Tensor) -> Self {
        self.check_binary(other, "add");
        let out = Tensor::zeros(&self.shape, self.dtype, self.device());
        let n = self.numel() as i32;
        match self.dtype {
            DType::F32 => unsafe {
                xtrain_cuda::ffi::launch_add_f32(
                    self.data_ptr() as *const f32,
                    other.data_ptr() as *const f32,
                    out.data_ptr() as *mut f32,
                    n,
                    std::ptr::null_mut(),
                );
            },
            DType::BF16 => unsafe {
                xtrain_cuda::ffi::launch_add_bf16(
                    self.data_ptr() as *const std::ffi::c_void,
                    other.data_ptr() as *const std::ffi::c_void,
                    out.data_ptr() as *mut std::ffi::c_void,
                    n,
                    std::ptr::null_mut(),
                );
            },
            _ => unreachable!(),
        }
        out
    }

    /// Elementwise `out = self * other` (same shape, Hadamard product).
    #[cfg(not(no_cuda))]
    pub fn mul(&self, other: &Tensor) -> Self {
        self.check_binary(other, "mul");
        let out = Tensor::zeros(&self.shape, self.dtype, self.device());
        let n = self.numel() as i32;
        match self.dtype {
            DType::F32 => unsafe {
                xtrain_cuda::ffi::launch_mul_f32(
                    self.data_ptr() as *const f32,
                    other.data_ptr() as *const f32,
                    out.data_ptr() as *mut f32,
                    n,
                    std::ptr::null_mut(),
                );
            },
            DType::BF16 => unsafe {
                xtrain_cuda::ffi::launch_mul_bf16(
                    self.data_ptr() as *const std::ffi::c_void,
                    other.data_ptr() as *const std::ffi::c_void,
                    out.data_ptr() as *mut std::ffi::c_void,
                    n,
                    std::ptr::null_mut(),
                );
            },
            _ => unreachable!(),
        }
        out
    }

    /// Broadcast bias add: `out[r,c] = self[r,c] + bias[c]`.
    /// `self`:[rows,cols], `bias`:[cols].
    #[cfg(not(no_cuda))]
    pub fn add_bias(&self, bias: &Tensor) -> Self {
        assert_eq!(self.ndim(), 2, "add_bias requires 2D input");
        assert_eq!(bias.ndim(), 1, "bias must be 1D");
        assert_eq!(self.shape[1], bias.shape[0], "bias len != cols");
        assert_eq!(self.dtype, bias.dtype, "add_bias dtype mismatch");
        let (rows, cols) = (self.shape[0], self.shape[1]);
        let out = Tensor::zeros(&self.shape, self.dtype, self.device());
        match self.dtype {
            DType::F32 => unsafe {
                xtrain_cuda::ffi::launch_add_bias_f32(
                    self.data_ptr() as *const f32,
                    bias.data_ptr() as *const f32,
                    out.data_ptr() as *mut f32,
                    rows as i32,
                    cols as i32,
                    std::ptr::null_mut(),
                );
            },
            DType::BF16 => unsafe {
                xtrain_cuda::ffi::launch_add_bias_bf16(
                    self.data_ptr() as *const std::ffi::c_void,
                    bias.data_ptr() as *const std::ffi::c_void,
                    out.data_ptr() as *mut std::ffi::c_void,
                    rows as i32,
                    cols as i32,
                    std::ptr::null_mut(),
                );
            },
            _ => panic!("add_bias supports F32/BF16"),
        }
        out
    }

    /// Column-sum over rows: `out[c] = sum_r self[r,c]`. This is the bias
    /// backward (sum the upstream grad over the broadcast dim). `self`:[rows,cols]
    /// → [cols].
    #[cfg(not(no_cuda))]
    pub fn sum_rows(&self) -> Self {
        assert_eq!(self.ndim(), 2, "sum_rows requires 2D input");
        let (rows, cols) = (self.shape[0], self.shape[1]);
        let out = Tensor::zeros(&[cols], self.dtype, self.device());
        match self.dtype {
            DType::F32 => unsafe {
                xtrain_cuda::ffi::launch_sum_rows_f32(
                    self.data_ptr() as *const f32,
                    out.data_ptr() as *mut f32,
                    rows as i32,
                    cols as i32,
                    std::ptr::null_mut(),
                );
            },
            DType::BF16 => unsafe {
                xtrain_cuda::ffi::launch_sum_rows_bf16(
                    self.data_ptr() as *const std::ffi::c_void,
                    out.data_ptr() as *mut std::ffi::c_void,
                    rows as i32,
                    cols as i32,
                    std::ptr::null_mut(),
                );
            },
            _ => panic!("sum_rows supports F32/BF16"),
        }
        out
    }

    /// RMSNorm forward: `y[r,c] = x[r,c] * inv_rms[r] * gamma[c]` with
    /// `inv_rms = rsqrt(mean(x²) + eps)`. `self`:[rows,cols], `gamma`:[cols].
    /// Returns `(y, inv_rms)`; `inv_rms`:[rows] is cached for backward.
    #[cfg(not(no_cuda))]
    pub fn rms_norm(&self, gamma: &Tensor, eps: f32) -> (Tensor, Tensor) {
        assert_eq!(self.ndim(), 2, "rms_norm requires 2D input");
        assert_eq!(gamma.ndim(), 1, "gamma must be 1D");
        assert_eq!(self.shape[1], gamma.shape[0], "gamma len != cols");
        // bf16: compute the reduction in fp32 (standard AMP), downcast y back to
        // bf16. inv_rms stays fp32 (the cache the fp32 backward kernel consumes).
        if self.dtype == DType::BF16 {
            let (y, inv_rms) = self
                .to_dtype(DType::F32)
                .rms_norm(&gamma.to_dtype(DType::F32), eps);
            return (y.to_dtype(DType::BF16), inv_rms);
        }
        let (rows, cols) = (self.shape[0], self.shape[1]);
        let y = Tensor::zeros(&self.shape, DType::F32, self.device());
        let inv_rms = Tensor::zeros(&[rows], DType::F32, self.device());
        unsafe {
            xtrain_cuda::ffi::launch_rms_norm_f32(
                self.data_ptr() as *const f32,
                gamma.data_ptr() as *const f32,
                y.data_ptr() as *mut f32,
                inv_rms.data_ptr() as *mut f32,
                rows as i32,
                cols as i32,
                eps,
                std::ptr::null_mut(),
            );
        }
        (y, inv_rms)
    }

    /// RMSNorm backward. Inputs are the forward `x`, `gamma`, upstream `dy`, and
    /// the cached `inv_rms`. Returns `(dx, dgamma)`.
    #[cfg(not(no_cuda))]
    pub fn rms_norm_backward(
        x: &Tensor,
        gamma: &Tensor,
        dy: &Tensor,
        inv_rms: &Tensor,
    ) -> (Tensor, Tensor) {
        // bf16: upcast (x, gamma, dy) to fp32, run the fp32 backward, downcast the
        // grads back to bf16 (inv_rms is already the fp32 cache).
        if x.dtype == DType::BF16 {
            let (dx, dgamma) = Tensor::rms_norm_backward(
                &x.to_dtype(DType::F32),
                &gamma.to_dtype(DType::F32),
                &dy.to_dtype(DType::F32),
                inv_rms,
            );
            return (dx.to_dtype(DType::BF16), dgamma.to_dtype(DType::BF16));
        }
        let (rows, cols) = (x.shape[0], x.shape[1]);
        let dx = Tensor::zeros(&[rows, cols], DType::F32, x.device());
        let dgamma = Tensor::zeros(&[cols], DType::F32, x.device());
        unsafe {
            xtrain_cuda::ffi::launch_rms_norm_dx_f32(
                x.data_ptr() as *const f32,
                gamma.data_ptr() as *const f32,
                dy.data_ptr() as *const f32,
                inv_rms.data_ptr() as *const f32,
                dx.data_ptr() as *mut f32,
                rows as i32,
                cols as i32,
                std::ptr::null_mut(),
            );
            xtrain_cuda::ffi::launch_rms_norm_dgamma_f32(
                x.data_ptr() as *const f32,
                dy.data_ptr() as *const f32,
                inv_rms.data_ptr() as *const f32,
                dgamma.data_ptr() as *mut f32,
                rows as i32,
                cols as i32,
                std::ptr::null_mut(),
            );
        }
        (dx, dgamma)
    }

    /// SiLU forward: `y = x * sigmoid(x)`, elementwise.
    #[cfg(not(no_cuda))]
    pub fn silu(&self) -> Self {
        assert!(
            matches!(self.dtype, DType::F32 | DType::BF16),
            "silu supports F32/BF16"
        );
        let out = Tensor::zeros(&self.shape, self.dtype, self.device());
        let n = self.numel() as i32;
        match self.dtype {
            DType::F32 => unsafe {
                xtrain_cuda::ffi::launch_silu_f32(
                    self.data_ptr() as *const f32,
                    out.data_ptr() as *mut f32,
                    n,
                    std::ptr::null_mut(),
                );
            },
            DType::BF16 => unsafe {
                xtrain_cuda::ffi::launch_silu_bf16(
                    self.data_ptr() as *const std::ffi::c_void,
                    out.data_ptr() as *mut std::ffi::c_void,
                    n,
                    std::ptr::null_mut(),
                );
            },
            _ => unreachable!(),
        }
        out
    }

    /// SiLU backward: `dx = dy * (sig + x*sig*(1-sig))`, `sig = sigmoid(x)`.
    /// Inputs are the forward `x` and upstream `dy`.
    #[cfg(not(no_cuda))]
    pub fn silu_backward(x: &Tensor, dy: &Tensor) -> Self {
        let dx = Tensor::zeros(&x.shape, x.dtype, x.device());
        let n = x.numel() as i32;
        match x.dtype {
            DType::F32 => unsafe {
                xtrain_cuda::ffi::launch_silu_dx_f32(
                    x.data_ptr() as *const f32,
                    dy.data_ptr() as *const f32,
                    dx.data_ptr() as *mut f32,
                    n,
                    std::ptr::null_mut(),
                );
            },
            DType::BF16 => unsafe {
                xtrain_cuda::ffi::launch_silu_dx_bf16(
                    x.data_ptr() as *const std::ffi::c_void,
                    dy.data_ptr() as *const std::ffi::c_void,
                    dx.data_ptr() as *mut std::ffi::c_void,
                    n,
                    std::ptr::null_mut(),
                );
            },
            _ => panic!("silu_backward supports F32/BF16"),
        }
        dx
    }

    /// Dropout forward (Phase T18). Returns `(out, mask)` where, for each element
    /// `i`, a counter-based RNG draws `u = hash(seed, i) ∈ [0,1)` and keeps the
    /// element iff `u >= p`; kept elements are scaled by `1/(1-p)` (inverted
    /// dropout, so `E[out] == x`). `mask[i]` stores that per-element factor
    /// (`1/(1-p)` if kept, else `0`) for the backward to reuse — the same mask, so
    /// the op is a fixed elementwise scale w.r.t. `x` (and finite-diff-checkable).
    ///
    /// The mask depends only on `(seed, i)`, NOT on `self`'s values, so a re-run
    /// with the same `seed` reproduces the same mask (T13 recompute stays exact).
    /// `mask` is always fp32 (the uniform is computed in fp32, dtype-independent);
    /// `out` matches `self`'s dtype. Requires `0 <= p < 1`.
    #[cfg(not(no_cuda))]
    pub fn dropout(&self, p: f32, seed: u64) -> (Self, Self) {
        assert!(
            matches!(self.dtype, DType::F32 | DType::BF16),
            "dropout supports F32/BF16"
        );
        assert!((0.0..1.0).contains(&p), "dropout p must be in [0,1)");
        assert!(self.is_contiguous(), "dropout requires contiguous tensor");
        let scale = 1.0 / (1.0 - p);
        let out = Tensor::zeros(&self.shape, self.dtype, self.device());
        let mask = Tensor::zeros(&self.shape, DType::F32, self.device());
        let n = self.numel() as i32;
        match self.dtype {
            DType::F32 => unsafe {
                xtrain_cuda::ffi::launch_dropout_fwd_f32(
                    self.data_ptr() as *const f32,
                    out.data_ptr() as *mut f32,
                    mask.data_ptr() as *mut f32,
                    p,
                    scale,
                    seed,
                    n,
                    std::ptr::null_mut(),
                );
            },
            DType::BF16 => unsafe {
                xtrain_cuda::ffi::launch_dropout_fwd_bf16(
                    self.data_ptr() as *const std::ffi::c_void,
                    out.data_ptr() as *mut std::ffi::c_void,
                    mask.data_ptr() as *mut f32,
                    p,
                    scale,
                    seed,
                    n,
                    std::ptr::null_mut(),
                );
            },
            _ => unreachable!(),
        }
        (out, mask)
    }

    /// Dropout backward: `dx = d ⊙ mask` (the SAME `mask` the forward cached).
    /// `d` is the upstream grad (activation dtype); `mask` is the fp32 factor
    /// tensor from [`Self::dropout`]. Output matches `d`'s dtype.
    #[cfg(not(no_cuda))]
    pub fn dropout_backward(d: &Tensor, mask: &Tensor) -> Self {
        assert_eq!(d.numel(), mask.numel(), "dropout_backward shape mismatch");
        assert_eq!(mask.dtype, DType::F32, "dropout mask must be F32");
        let dx = Tensor::zeros(&d.shape, d.dtype, d.device());
        let n = d.numel() as i32;
        match d.dtype {
            DType::F32 => unsafe {
                xtrain_cuda::ffi::launch_dropout_bwd_f32(
                    d.data_ptr() as *const f32,
                    mask.data_ptr() as *const f32,
                    dx.data_ptr() as *mut f32,
                    n,
                    std::ptr::null_mut(),
                );
            },
            DType::BF16 => unsafe {
                xtrain_cuda::ffi::launch_dropout_bwd_bf16(
                    d.data_ptr() as *const std::ffi::c_void,
                    mask.data_ptr() as *const f32,
                    dx.data_ptr() as *mut std::ffi::c_void,
                    n,
                    std::ptr::null_mut(),
                );
            },
            _ => panic!("dropout_backward supports F32/BF16"),
        }
        dx
    }

    /// RoPE forward (rotate_half). `self`:[tokens,heads,head_dim]; each token's
    /// position is `row % period`. `period` = sequence length, so a flattened
    /// batch `[B*S,heads,head_dim]` gets per-sequence positions (pass `period=S`);
    /// pass `period=tokens` for a single sequence (position = row). Returns the
    /// rotated tensor.
    #[cfg(not(no_cuda))]
    pub fn rope(&self, theta: f32, period: usize) -> Self {
        assert_eq!(self.ndim(), 3, "rope requires [tokens,heads,head_dim]");
        let (tokens, heads, head_dim) = (self.shape[0], self.shape[1], self.shape[2]);
        assert_eq!(head_dim % 2, 0, "head_dim must be even");
        assert!(
            period > 0 && tokens % period == 0,
            "tokens must be a multiple of period"
        );
        if self.dtype == DType::BF16 {
            return self
                .to_dtype(DType::F32)
                .rope(theta, period)
                .to_dtype(DType::BF16);
        }
        let out = Tensor::zeros(&self.shape, DType::F32, self.device());
        unsafe {
            xtrain_cuda::ffi::launch_rope_f32(
                self.data_ptr() as *const f32,
                out.data_ptr() as *mut f32,
                tokens as i32,
                heads as i32,
                head_dim as i32,
                theta,
                period as i32,
                std::ptr::null_mut(),
            );
        }
        out
    }

    /// RoPE backward: apply the inverse (transpose) rotation to `dy`. RoPE is an
    /// orthogonal map, so it needs no cached forward values, only `theta`/`period`.
    #[cfg(not(no_cuda))]
    pub fn rope_backward(dy: &Tensor, theta: f32, period: usize) -> Self {
        if dy.dtype == DType::BF16 {
            return Tensor::rope_backward(&dy.to_dtype(DType::F32), theta, period)
                .to_dtype(DType::BF16);
        }
        let (tokens, heads, head_dim) = (dy.shape[0], dy.shape[1], dy.shape[2]);
        let dx = Tensor::zeros(&dy.shape, DType::F32, dy.device());
        unsafe {
            xtrain_cuda::ffi::launch_rope_dx_f32(
                dy.data_ptr() as *const f32,
                dx.data_ptr() as *mut f32,
                tokens as i32,
                heads as i32,
                head_dim as i32,
                theta,
                period as i32,
                std::ptr::null_mut(),
            );
        }
        dx
    }

    /// Row-wise safe softmax over the last dim. `self`:[rows,cols].
    #[cfg(not(no_cuda))]
    pub fn softmax(&self) -> Self {
        assert_eq!(self.ndim(), 2, "softmax requires 2D input");
        if self.dtype == DType::BF16 {
            return self.to_dtype(DType::F32).softmax().to_dtype(DType::BF16);
        }
        let (rows, cols) = (self.shape[0], self.shape[1]);
        let out = Tensor::zeros(&self.shape, DType::F32, self.device());
        unsafe {
            xtrain_cuda::ffi::launch_softmax_f32(
                self.data_ptr() as *const f32,
                out.data_ptr() as *mut f32,
                rows as i32,
                cols as i32,
                std::ptr::null_mut(),
            );
        }
        out
    }

    /// Softmax backward (Jacobian): `dx[r,c] = y[r,c]*(dy[r,c] - sum_c'(dy*y))`.
    /// Inputs are the forward output `y` and upstream `dy`.
    #[cfg(not(no_cuda))]
    pub fn softmax_backward(y: &Tensor, dy: &Tensor) -> Self {
        if y.dtype == DType::BF16 {
            return Tensor::softmax_backward(&y.to_dtype(DType::F32), &dy.to_dtype(DType::F32))
                .to_dtype(DType::BF16);
        }
        let (rows, cols) = (y.shape[0], y.shape[1]);
        let dx = Tensor::zeros(&y.shape, DType::F32, y.device());
        unsafe {
            xtrain_cuda::ffi::launch_softmax_dx_f32(
                y.data_ptr() as *const f32,
                dy.data_ptr() as *const f32,
                dx.data_ptr() as *mut f32,
                rows as i32,
                cols as i32,
                std::ptr::null_mut(),
            );
        }
        dx
    }

    /// Cross-entropy forward over logits `self`:[rows,cols] with one I32 target
    /// per row. Returns `(probs, loss)` where `probs`:[rows,cols] is the softmax
    /// (cached for backward) and `loss`:[rows] is the per-row negative log-likelihood.
    #[cfg(not(no_cuda))]
    pub fn cross_entropy(&self, target: &Tensor) -> (Tensor, Tensor) {
        assert_eq!(self.ndim(), 2, "cross_entropy requires 2D logits");
        assert_eq!(target.dtype, DType::I32, "target must be I32");
        assert_eq!(target.numel(), self.shape[0], "one target per row");
        // CE math (log-sum-exp) is fp32 (probs/loss cached fp32). The model casts
        // logits→fp32 before CE; this guard keeps the op robust to bf16 logits.
        if self.dtype == DType::BF16 {
            return self.to_dtype(DType::F32).cross_entropy(target);
        }
        let (rows, cols) = (self.shape[0], self.shape[1]);
        let probs = Tensor::zeros(&self.shape, DType::F32, self.device());
        let loss = Tensor::zeros(&[rows], DType::F32, self.device());
        unsafe {
            xtrain_cuda::ffi::launch_cross_entropy_fwd_f32(
                self.data_ptr() as *const f32,
                target.data_ptr() as *const i32,
                probs.data_ptr() as *mut f32,
                loss.data_ptr() as *mut f32,
                rows as i32,
                cols as i32,
                std::ptr::null_mut(),
            );
        }
        (probs, loss)
    }

    /// Cross-entropy backward: `dx = scale * (probs - onehot(target))`. With
    /// `scale = upstream / rows`, this is the gradient of the mean per-row loss.
    #[cfg(not(no_cuda))]
    pub fn cross_entropy_backward(probs: &Tensor, target: &Tensor, scale: f32) -> Self {
        let (rows, cols) = (probs.shape[0], probs.shape[1]);
        let dx = Tensor::zeros(&probs.shape, DType::F32, probs.device());
        unsafe {
            xtrain_cuda::ffi::launch_cross_entropy_dx_f32(
                probs.data_ptr() as *const f32,
                target.data_ptr() as *const i32,
                dx.data_ptr() as *mut f32,
                rows as i32,
                cols as i32,
                scale,
                std::ptr::null_mut(),
            );
        }
        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(),
            );
        }
        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(),
            );
        }
        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.ndim(), 3, "transpose_3d01 requires a 3D tensor");
        assert!(self.is_contiguous(), "transpose_3d01 requires contiguous");
        if self.dtype == DType::BF16 {
            return self
                .to_dtype(DType::F32)
                .transpose_3d01()
                .to_dtype(DType::BF16);
        }
        assert_eq!(self.dtype, DType::F32, "transpose_3d01 supports F32/BF16");
        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(),
            );
        }
        out
    }

    // --- Batched attention (the T10 fused op) ---

    /// Batched causal scaled-dot-product attention. `self`=Q, `k`, `v` are each
    /// `[bh, seq, head_dim]` (bh = batch·n_heads), contiguous F32 on one GPU.
    /// Computes, per batch element, `out = softmax(causal(Q·Kᵀ / √hd)) · V`. The
    /// two GEMMs run as `cublasSgemmStridedBatched` and the softmax+scale+causal
    /// mask is one kernel, so the whole attention is 3 launches regardless of bh.
    /// Returns `(out, probs)` where `probs`:[bh,seq,seq] is cached for backward.
    #[cfg(not(no_cuda))]
    pub fn attention(&self, k: &Tensor, v: &Tensor, scale: f32) -> (Tensor, Tensor) {
        assert_eq!(self.ndim(), 3, "attention Q must be [bh,seq,head_dim]");
        assert_eq!(self.shape(), k.shape(), "Q/K shape mismatch");
        assert_eq!(self.shape(), v.shape(), "Q/V shape mismatch");
        assert_eq!(self.dtype, k.dtype, "Q/K dtype mismatch");
        assert_eq!(self.dtype, v.dtype, "Q/V dtype mismatch");
        let (bh, seq, hd) = (self.shape[0], self.shape[1], self.shape[2]);
        let dev = self.device();
        let dt = self.dtype;

        // scores[bh,seq,seq] = Q[bh,seq,hd] · Kᵀ[bh,hd,seq]  (GEMM in self dtype)
        let scores = Tensor::zeros(&[bh, seq, seq], dt, dev);
        strided_batched_gemm(
            dt,
            false,
            true,
            seq,
            seq,
            hd,
            self.data_ptr(),
            seq * hd,
            k.data_ptr(),
            seq * hd,
            scores.data_ptr(),
            seq * seq,
            bh,
        );
        // probs = softmax(causal(scores · scale)). Softmax math is fp32 (stable);
        // for bf16 we upcast scores → f32 → kernel → downcast probs back to bf16
        // (so the cached probs activation is half-size). One block per [bh·seq] row.
        let scores_f32 = scores.to_dtype(DType::F32);
        let probs_f32 = Tensor::zeros(&[bh, seq, seq], DType::F32, dev);
        unsafe {
            xtrain_cuda::ffi::launch_softmax_causal_f32(
                scores_f32.data_ptr() as *const f32,
                probs_f32.data_ptr() as *mut f32,
                (bh * seq) as i32,
                seq as i32,
                scale,
                std::ptr::null_mut(),
            );
        }
        let probs = probs_f32.to_dtype(dt);
        // out[bh,seq,hd] = probs[bh,seq,seq] · V[bh,seq,hd]
        let out = Tensor::zeros(&[bh, seq, hd], dt, dev);
        strided_batched_gemm(
            dt,
            false,
            false,
            seq,
            hd,
            seq,
            probs.data_ptr(),
            seq * seq,
            v.data_ptr(),
            seq * hd,
            out.data_ptr(),
            seq * hd,
            bh,
        );
        (out, probs)
    }

    /// Backward of [`attention`](Self::attention). Inputs: forward `q`,`k`,`v`,
    /// the cached `probs`, the upstream `dout` (all batched `[bh,seq,*]`), and the
    /// same `scale`. Returns `(dq, dk, dv)`.
    ///
    ///   dP = dOut · Vᵀ ; dV = Pᵀ · dOut
    ///   dScores = softmax_jacobian(P, dP) · scale     (scale folded back in)
    ///   dQ = dScores · K ; dK = dScoresᵀ · Q
    ///
    /// Masked (future) entries of P are 0, so the softmax Jacobian zeros their
    /// gradient — the causal mask needs no special handling here.
    #[cfg(not(no_cuda))]
    pub fn attention_backward(
        q: &Tensor,
        k: &Tensor,
        v: &Tensor,
        probs: &Tensor,
        dout: &Tensor,
        scale: f32,
    ) -> (Tensor, Tensor, Tensor) {
        let (bh, seq, hd) = (q.shape[0], q.shape[1], q.shape[2]);
        let dev = q.device();
        let dt = q.dtype;

        // dP[bh,seq,seq] = dOut[bh,seq,hd] · Vᵀ[bh,hd,seq]
        let dp = Tensor::zeros(&[bh, seq, seq], dt, dev);
        strided_batched_gemm(
            dt,
            false,
            true,
            seq,
            seq,
            hd,
            dout.data_ptr(),
            seq * hd,
            v.data_ptr(),
            seq * hd,
            dp.data_ptr(),
            seq * seq,
            bh,
        );
        // dV[bh,seq,hd] = Pᵀ[bh,seq,seq] · dOut[bh,seq,hd]
        let dv = Tensor::zeros(&[bh, seq, hd], dt, dev);
        strided_batched_gemm(
            dt,
            true,
            false,
            seq,
            hd,
            seq,
            probs.data_ptr(),
            seq * seq,
            dout.data_ptr(),
            seq * hd,
            dv.data_ptr(),
            seq * hd,
            bh,
        );
        // dScores = softmax Jacobian (per row) applied to dP, then ×scale.
        // softmax_backward + scale are dtype-aware (fp32 math inside for bf16).
        let dscores = Tensor::softmax_backward(
            &probs.reshape(&[bh * seq, seq]),
            &dp.reshape(&[bh * seq, seq]),
        )
        .reshape(&[bh, seq, seq]);
        let dscores = dscores.scale(scale);
        // dQ[bh,seq,hd] = dScores[bh,seq,seq] · K[bh,seq,hd]
        let dq = Tensor::zeros(&[bh, seq, hd], dt, dev);
        strided_batched_gemm(
            dt,
            false,
            false,
            seq,
            hd,
            seq,
            dscores.data_ptr(),
            seq * seq,
            k.data_ptr(),
            seq * hd,
            dq.data_ptr(),
            seq * hd,
            bh,
        );
        // dK[bh,seq,hd] = dScoresᵀ[bh,seq,seq] · Q[bh,seq,hd]
        let dk = Tensor::zeros(&[bh, seq, hd], dt, dev);
        strided_batched_gemm(
            dt,
            true,
            false,
            seq,
            hd,
            seq,
            dscores.data_ptr(),
            seq * seq,
            q.data_ptr(),
            seq * hd,
            dk.data_ptr(),
            seq * hd,
            bh,
        );
        (dq, dk, dv)
    }

    // --- Fused flash-attention (the T14 op) ---

    /// Fused flash-attention forward (Phase T14). `self`=Q, `k`, `v` each
    /// `[bh, seq, head_dim]`, contiguous on one GPU. Computes, per batch element,
    /// `out = softmax(causal(Q·Kᵀ·scale))·V` in a SINGLE kernel that streams over
    /// KV tiles with an online softmax — the `[bh,seq,seq]` score matrix is NEVER
    /// materialized. Returns `(out, lse)` where `lse`:[bh,seq] (F32) is the per-row
    /// logsumexp cached for backward (O(N), vs the composed path's O(N²) probs).
    ///
    /// The fused kernel is fp32; for bf16 we upcast Q/K/V → f32 → kernel → downcast
    /// `out` back to bf16 (same fp32-softmax policy as the composed [`attention`]),
    /// so flash and composed produce the same softmax numerics. `lse` stays fp32.
    #[cfg(not(no_cuda))]
    pub fn flash_attention(&self, k: &Tensor, v: &Tensor, scale: f32) -> (Tensor, Tensor) {
        assert_eq!(
            self.ndim(),
            3,
            "flash_attention Q must be [bh,seq,head_dim]"
        );
        assert_eq!(self.shape(), k.shape(), "Q/K shape mismatch");
        assert_eq!(self.shape(), v.shape(), "Q/V shape mismatch");
        assert_eq!(self.dtype, k.dtype, "Q/K dtype mismatch");
        assert_eq!(self.dtype, v.dtype, "Q/V dtype mismatch");
        let (bh, seq, hd) = (self.shape[0], self.shape[1], self.shape[2]);
        let dev = self.device();
        let dt = self.dtype;

        let qf = self.to_dtype(DType::F32);
        let kf = k.to_dtype(DType::F32);
        let vf = v.to_dtype(DType::F32);
        let out_f32 = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
        let lse = Tensor::zeros(&[bh, seq], DType::F32, dev);
        unsafe {
            xtrain_cuda::ffi::launch_flash_attention_fwd_f32(
                qf.data_ptr() as *const f32,
                kf.data_ptr() as *const f32,
                vf.data_ptr() as *const f32,
                out_f32.data_ptr() as *mut f32,
                lse.data_ptr() as *mut f32,
                bh as i32,
                seq as i32,
                hd as i32,
                scale,
                std::ptr::null_mut(),
            );
        }
        (out_f32.to_dtype(dt), lse)
    }

    /// Backward of [`flash_attention`](Self::flash_attention). Inputs: forward
    /// `q`,`k`,`v`, the forward output `out`, the cached `lse`:[bh,seq], the upstream
    /// `dout`, and the same `scale`. Returns `(dq, dk, dv)`.
    ///
    /// flash-style: NO cached probs. Recomputes scores from Q/K/V + `lse`, uses
    /// `D[i]=Σ dOᵢ·Oᵢ` to collapse the softmax Jacobian, streams KV in tiles. dQ is
    /// owned per query row; dK/dV are accumulated across rows (atomicAdd). Same
    /// fp32 kernel; bf16 callers get fp32 grads which the autograd `cast` op casts.
    #[cfg(not(no_cuda))]
    pub fn flash_attention_backward(
        q: &Tensor,
        k: &Tensor,
        v: &Tensor,
        out: &Tensor,
        lse: &Tensor,
        dout: &Tensor,
        scale: f32,
    ) -> (Tensor, Tensor, Tensor) {
        let (bh, seq, hd) = (q.shape[0], q.shape[1], q.shape[2]);
        let dev = q.device();
        let dt = q.dtype;

        let qf = q.to_dtype(DType::F32);
        let kf = k.to_dtype(DType::F32);
        let vf = v.to_dtype(DType::F32);
        let of = out.to_dtype(DType::F32);
        let dof = dout.to_dtype(DType::F32);
        // D[i] = Σ_d dO[i,d]·O[i,d] (one scalar per query row, O(N)).
        let d = Tensor::zeros(&[bh, seq], DType::F32, dev);
        unsafe {
            xtrain_cuda::ffi::launch_flash_attention_rowdot_f32(
                dof.data_ptr() as *const f32,
                of.data_ptr() as *const f32,
                d.data_ptr() as *mut f32,
                (bh * seq) as i32,
                hd as i32,
                std::ptr::null_mut(),
            );
        }
        // dq/dk/dv pre-zeroed (Tensor::zeros memsets); dk/dv accumulate via atomicAdd.
        let dq = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
        let dk = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
        let dv = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
        unsafe {
            xtrain_cuda::ffi::launch_flash_attention_bwd_f32(
                qf.data_ptr() as *const f32,
                kf.data_ptr() as *const f32,
                vf.data_ptr() as *const f32,
                dof.data_ptr() as *const f32,
                lse.data_ptr() as *const f32,
                d.data_ptr() as *mut f32,
                dq.data_ptr() as *mut f32,
                dk.data_ptr() as *mut f32,
                dv.data_ptr() as *mut f32,
                bh as i32,
                seq as i32,
                hd as i32,
                scale,
                std::ptr::null_mut(),
            );
        }
        (dq.to_dtype(dt), dk.to_dtype(dt), dv.to_dtype(dt))
    }

    /// 4D axis-(1,2) transpose: `self`:[a,b,c,d] → [a,c,b,d],
    /// `out[i,k,j,l]=self[i,j,k,l]`. Lays out batched multi-head attention
    /// (`[B,S,nh,hd] <-> [B,nh,S,hd]`). Its own backward is the same op (swap b,c).
    #[cfg(not(no_cuda))]
    pub fn transpose_4d12(&self) -> Self {
        assert_eq!(self.ndim(), 4, "transpose_4d12 requires a 4D tensor");
        assert!(self.is_contiguous(), "transpose_4d12 requires contiguous");
        if self.dtype == DType::BF16 {
            return self
                .to_dtype(DType::F32)
                .transpose_4d12()
                .to_dtype(DType::BF16);
        }
        assert_eq!(self.dtype, DType::F32, "transpose_4d12 supports F32/BF16");
        let (a, b, c, d) = (self.shape[0], self.shape[1], self.shape[2], self.shape[3]);
        let out = Tensor::zeros(&[a, c, b, d], DType::F32, self.device());
        unsafe {
            xtrain_cuda::ffi::launch_transpose_4d12_f32(
                self.data_ptr() as *const f32,
                out.data_ptr() as *mut f32,
                a as i32,
                b as i32,
                c as i32,
                d as i32,
                std::ptr::null_mut(),
            );
        }
        out
    }

    // Shared validation for same-shape binary elementwise ops.
    #[cfg(not(no_cuda))]
    fn check_binary(&self, other: &Tensor, op: &str) {
        assert!(
            matches!(self.dtype, DType::F32 | DType::BF16),
            "{op} supports F32/BF16"
        );
        assert_eq!(self.dtype, other.dtype, "{op} dtype mismatch");
        assert_eq!(self.shape(), other.shape(), "{op} shape mismatch");
        assert_eq!(self.device(), other.device(), "{op} device mismatch");
        assert!(
            self.is_contiguous() && other.is_contiguous(),
            "{op} requires contiguous tensors"
        );
    }
}

/// Dispatch a strided-batched GEMM on `dt`: fp32 → `sgemm_strided_batched`,
/// bf16 → `gemm_ex_strided_batched` (bf16 in/out, fp32 accum). Pointers are the
/// raw `data_ptr()` bytes of contiguous same-dtype tensors. `alpha=1, beta=0`.
/// The fp32 path is bit-identical to the inlined T10 call it replaces.
#[cfg(not(no_cuda))]
#[allow(clippy::too_many_arguments)]
fn strided_batched_gemm(
    dt: DType,
    trans_a: bool,
    trans_b: bool,
    m: usize,
    n: usize,
    k: usize,
    a: *const u8,
    stride_a: usize,
    b: *const u8,
    stride_b: usize,
    c: *const u8,
    stride_c: usize,
    batch: usize,
) {
    match dt {
        DType::F32 => xtrain_cuda::cublas::sgemm_strided_batched(
            trans_a,
            trans_b,
            m,
            n,
            k,
            1.0,
            a as *const f32,
            stride_a,
            b as *const f32,
            stride_b,
            0.0,
            c as *mut f32,
            stride_c,
            batch,
        ),
        DType::BF16 => xtrain_cuda::cublas::gemm_ex_strided_batched(
            trans_a,
            trans_b,
            m,
            n,
            k,
            1.0,
            a as *const std::ffi::c_void,
            stride_a,
            b as *const std::ffi::c_void,
            stride_b,
            0.0,
            c as *mut std::ffi::c_void,
            stride_c,
            batch,
        ),
        _ => panic!("strided_batched_gemm supports F32/BF16"),
    }
}

impl std::fmt::Debug for Tensor {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Tensor(shape={:?}, dtype={}, device={}, contiguous={})",
            self.shape.as_slice(),
            self.dtype,
            self.device(),
            self.is_contiguous()
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn from_slice_shape_and_data() {
        let t = Tensor::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]);
        assert_eq!(t.shape(), &[2, 3]);
        assert_eq!(t.strides(), &[3, 1]);
        assert_eq!(t.numel(), 6);
        assert_eq!(t.device(), Device::Cpu);
        assert!(t.is_contiguous());
        assert_eq!(t.as_slice::<f32>(), &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
    }

    #[test]
    fn zeros_cpu() {
        let t = Tensor::zeros(&[4], DType::F32, Device::Cpu);
        assert_eq!(t.as_slice::<f32>(), &[0.0, 0.0, 0.0, 0.0]);
    }
}