Merge branch 'feat/f64-cpu-ops'

crates/core/rtx-backend-cpu/Cargo.toml

@@ -12,6 +12,7 @@ categories = ["science", "mathematics"]
 [dependencies]
 rtx-backend = { workspace = true }
 thiserror = { workspace = true }
+num-traits = "0.2"
 parking_lot = { workspace = true }
 rand = { workspace = true }
 rand_distr = { workspace = true }

crates/core/rtx-backend-cpu/src/device.rs

@@ -2,7 +2,7 @@
 
 use rtx_backend::{DeviceId, DeviceOps};
 
-use crate::CpuBackend;
+use crate::{CpuBackend, CpuBackendF64};
 
 /// CPU device for the CPU backend.
 ///
@@ -74,7 +74,7 @@ impl DeviceOps<CpuBackend> for CpuDevice {
 
     fn memory_available(&self) -> usize {
         // Estimate available memory
-        self.memory_capacity() / 2
+        (16 * 1024 * 1024 * 1024_usize) / 2
     }
 
     fn compute_capability(&self) -> Option<(u32, u32)> {
@@ -91,6 +91,31 @@ impl DeviceOps<CpuBackend> for CpuDevice {
     }
 }
 
+/// Same device, viewed through the f64 backend (CPU has one logical device).
+impl DeviceOps<CpuBackendF64> for CpuDevice {
+    fn id(&self) -> DeviceId {
+        DeviceId::Cpu
+    }
+
+    fn memory_capacity(&self) -> usize {
+        16 * 1024 * 1024 * 1024 // 16GB default
+    }
+
+    fn memory_available(&self) -> usize {
+        (16 * 1024 * 1024 * 1024_usize) / 2
+    }
+
+    fn compute_capability(&self) -> Option<(u32, u32)> {
+        None
+    }
+
+    fn synchronize(&self) {}
+
+    fn is_available(&self) -> bool {
+        true
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -98,7 +123,8 @@ mod tests {
     #[test]
     fn test_device_creation() {
         let device = CpuDevice::new();
-        assert!(device.is_available());
+        // `is_available` is shared by both DeviceOps impls (f32/f64); pick one.
+        assert!(DeviceOps::<CpuBackend>::is_available(&device));
         assert!(device.num_threads() > 0);
         println!("CPU threads: {}", device.num_threads());
     }

crates/core/rtx-backend-cpu/src/lib.rs

@@ -47,6 +47,16 @@ pub use tensor::CpuTensorPrimitive;
 
 use rtx_backend::{Backend, BoolU8};
 
+/// Element bound for the generic CPU ops.
+///
+/// `num_traits::Float` supplies the arithmetic (`Add`/`Sub`/`Mul`/`Div` via
+/// `Num`), `zero()`/`one()`, transcendentals (`exp`/`ln`/`sqrt`/`sin`/`cos`/
+/// `powf`/`abs`/`max`/`min`), and `from::<f64>()` casts the ops need; `Send +
+/// Sync + 'static` are required by rayon. Both `f32` and `f64` satisfy it, so a
+/// single op implementation serves `CpuBackend` (f32) and `CpuBackendF64` (f64).
+pub trait CpuFloat: num_traits::Float + Send + Sync + 'static {}
+impl<T: num_traits::Float + Send + Sync + 'static> CpuFloat for T {}
+
 /// CPU backend for RustyTorch++.
 ///
 /// This backend provides a pure Rust implementation suitable for:
@@ -413,8 +423,384 @@ impl Backend for CpuBackend {
     }
 }
 
+/// Double-precision (f64) CPU backend for RustyTorch++.
+///
+/// Identical to [`CpuBackend`] but with `FloatElem = f64`, backed by
+/// `CpuTensorPrimitive<D, f64>`. Every method delegates to the same generic
+/// `ops::*` implementations (which infer the element type from the f64 primitive),
+/// so the two backends share one code path. Use this where quantum-precision
+/// gradients / numerics need f64; `CpuBackend` (f32) is unchanged.
+#[derive(Clone, Debug, Default)]
+pub struct CpuBackendF64;
+
+impl Backend for CpuBackendF64 {
+    type TensorPrimitive<const D: usize> = CpuTensorPrimitive<D, f64>;
+    type Device = CpuDevice;
+    type FloatElem = f64;
+    type IntElem = i32;
+    type BoolElem = BoolU8;
+
+    fn name() -> &'static str {
+        "cpu_f64"
+    }
+
+    fn seed(seed: u64) {
+        ops::seed_rng(seed);
+    }
+
+    fn zeros<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D> {
+        ops::creation::zeros(shape, device)
+    }
+
+    fn ones<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D> {
+        ops::creation::ones(shape, device)
+    }
+
+    fn full<const D: usize>(
+        shape: [usize; D],
+        fill_value: Self::FloatElem,
+        device: &Self::Device,
+    ) -> Self::TensorPrimitive<D> {
+        ops::creation::full(shape, fill_value, device)
+    }
+
+    fn rand<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D> {
+        ops::creation::rand(shape, device)
+    }
+
+    fn randn<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D> {
+        ops::creation::randn(shape, device)
+    }
+
+    fn from_data<const D: usize>(
+        data: &[Self::FloatElem],
+        shape: [usize; D],
+        device: &Self::Device,
+    ) -> Self::TensorPrimitive<D> {
+        ops::creation::from_data(data, shape, device)
+    }
+
+    fn add<const D: usize>(
+        lhs: Self::TensorPrimitive<D>,
+        rhs: Self::TensorPrimitive<D>,
+    ) -> Self::TensorPrimitive<D> {
+        ops::basic::add(&lhs, &rhs)
+    }
+
+    fn sub<const D: usize>(
+        lhs: Self::TensorPrimitive<D>,
+        rhs: Self::TensorPrimitive<D>,
+    ) -> Self::TensorPrimitive<D> {
+        ops::basic::sub(&lhs, &rhs)
+    }
+
+    fn mul<const D: usize>(
+        lhs: Self::TensorPrimitive<D>,
+        rhs: Self::TensorPrimitive<D>,
+    ) -> Self::TensorPrimitive<D> {
+        ops::basic::mul(&lhs, &rhs)
+    }
+
+    fn div<const D: usize>(
+        lhs: Self::TensorPrimitive<D>,
+        rhs: Self::TensorPrimitive<D>,
+    ) -> Self::TensorPrimitive<D> {
+        ops::basic::div(&lhs, &rhs)
+    }
+
+    fn matmul(
+        lhs: Self::TensorPrimitive<2>,
+        rhs: Self::TensorPrimitive<2>,
+    ) -> Self::TensorPrimitive<2> {
+        ops::gemm::matmul(&lhs, &rhs)
+    }
+
+    fn bmm(
+        lhs: Self::TensorPrimitive<3>,
+        rhs: Self::TensorPrimitive<3>,
+    ) -> Self::TensorPrimitive<3> {
+        ops::gemm::bmm(&lhs, &rhs)
+    }
+
+    fn neg<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
+        ops::unary::neg(&tensor)
+    }
+
+    fn exp<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
+        ops::unary::exp(&tensor)
+    }
+
+    fn log<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
+        ops::unary::log(&tensor)
+    }
+
+    fn sqrt<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
+        ops::unary::sqrt(&tensor)
+    }
+
+    fn abs<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
+        ops::unary::abs(&tensor)
+    }
+
+    fn sin<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
+        ops::unary::sin(&tensor)
+    }
+
+    fn cos<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
+        ops::unary::cos(&tensor)
+    }
+
+    fn pow<const D: usize>(
+        tensor: Self::TensorPrimitive<D>,
+        exp: Self::FloatElem,
+    ) -> Self::TensorPrimitive<D> {
+        ops::unary::pow(&tensor, exp)
+    }
+
+    fn clamp<const D: usize>(
+        tensor: Self::TensorPrimitive<D>,
+        min: Self::FloatElem,
+        max: Self::FloatElem,
+    ) -> Self::TensorPrimitive<D> {
+        ops::unary::clamp(&tensor, min, max)
+    }
+
+    fn relu<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
+        ops::activation::relu(&tensor)
+    }
+
+    fn sigmoid<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
+        ops::activation::sigmoid(&tensor)
+    }
+
+    fn tanh<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
+        ops::activation::tanh(&tensor)
+    }
+
+    fn sum<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
+        ops::reduction::sum(&tensor)
+    }
+
+    fn sum_dim<const D: usize>(
+        tensor: Self::TensorPrimitive<D>,
+        dim: usize,
+    ) -> Self::TensorPrimitive<D> {
+        ops::reduction::sum_dim(&tensor, dim)
+    }
+
+    fn mean<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
+        ops::reduction::mean(&tensor)
+    }
+
+    fn mean_dim<const D: usize>(
+        tensor: Self::TensorPrimitive<D>,
+        dim: usize,
+    ) -> Self::TensorPrimitive<D> {
+        ops::reduction::mean_dim(&tensor, dim)
+    }
+
+    fn var<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
+        ops::reduction::var(&tensor)
+    }
+
+    fn var_dim<const D: usize>(
+        tensor: Self::TensorPrimitive<D>,
+        dim: usize,
+    ) -> Self::TensorPrimitive<D> {
+        ops::reduction::var_dim(&tensor, dim)
+    }
+
+    fn max<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
+        ops::reduction::max(&tensor)
+    }
+
+    fn min<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
+        ops::reduction::min(&tensor)
+    }
+
+    fn shape<const D: usize>(tensor: &Self::TensorPrimitive<D>) -> [usize; D] {
+        tensor.shape
+    }
+
+    fn reshape<const D1: usize, const D2: usize>(
+        tensor: Self::TensorPrimitive<D1>,
+        shape: [usize; D2],
+    ) -> Self::TensorPrimitive<D2> {
+        ops::shape::reshape(tensor, shape)
+    }
+
+    fn transpose<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
+        ops::shape::transpose(&tensor)
+    }
+
+    fn swap_dims<const D: usize>(
+        tensor: Self::TensorPrimitive<D>,
+        dim1: usize,
+        dim2: usize,
+    ) -> Self::TensorPrimitive<D> {
+        ops::shape::swap_dims(&tensor, dim1, dim2)
+    }
+
+    fn flash_attention(
+        query: Self::TensorPrimitive<4>,
+        key: Self::TensorPrimitive<4>,
+        value: Self::TensorPrimitive<4>,
+        mask: Option<&Self::TensorPrimitive<4>>,
+        scale: Self::FloatElem,
+        causal: bool,
+    ) -> Self::TensorPrimitive<4> {
+        ops::attention::flash_attention(&query, &key, &value, mask, scale, causal)
+    }
+
+    fn softmax<const D: usize>(
+        tensor: Self::TensorPrimitive<D>,
+        dim: usize,
+    ) -> Self::TensorPrimitive<D> {
+        ops::activation::softmax(&tensor, dim)
+    }
+
+    fn layer_norm<const D: usize>(
+        tensor: Self::TensorPrimitive<D>,
+        weight: &Self::TensorPrimitive<1>,
+        bias: Option<&Self::TensorPrimitive<1>>,
+        eps: Self::FloatElem,
+    ) -> Self::TensorPrimitive<D> {
+        ops::normalization::layer_norm(&tensor, weight, bias, eps)
+    }
+
+    fn rms_norm<const D: usize>(
+        tensor: Self::TensorPrimitive<D>,
+        weight: &Self::TensorPrimitive<1>,
+        eps: Self::FloatElem,
+    ) -> Self::TensorPrimitive<D> {
+        ops::normalization::rms_norm(&tensor, weight, eps)
+    }
+
+    fn rope<const D: usize>(
+        tensor: Self::TensorPrimitive<D>,
+        cos: &Self::TensorPrimitive<2>,
+        sin: &Self::TensorPrimitive<2>,
+    ) -> Self::TensorPrimitive<D> {
+        ops::attention::rope(&tensor, cos, sin)
+    }
+
+    fn gelu<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
+        ops::activation::gelu(&tensor)
+    }
+
+    fn silu<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
+        ops::activation::silu(&tensor)
+    }
+
+    fn leaky_relu<const D: usize>(
+        tensor: Self::TensorPrimitive<D>,
+        negative_slope: Self::FloatElem,
+    ) -> Self::TensorPrimitive<D> {
+        ops::activation::leaky_relu(&tensor, negative_slope)
+    }
+
+    fn elu<const D: usize>(
+        tensor: Self::TensorPrimitive<D>,
+        alpha: Self::FloatElem,
+    ) -> Self::TensorPrimitive<D> {
+        ops::activation::elu(&tensor, alpha)
+    }
+
+    fn gt_scalar<const D: usize>(
+        tensor: Self::TensorPrimitive<D>,
+        value: Self::FloatElem,
+    ) -> Self::TensorPrimitive<D> {
+        ops::basic::gt_scalar(&tensor, value)
+    }
+
+    fn conv2d(
+        input: Self::TensorPrimitive<4>,
+        weight: &Self::TensorPrimitive<4>,
+        bias: Option<&Self::TensorPrimitive<1>>,
+        stride: [usize; 2],
+        padding: [usize; 2],
+        dilation: [usize; 2],
+        groups: usize,
+    ) -> Self::TensorPrimitive<4> {
+        ops::conv::conv2d(&input, weight, bias, stride, padding, dilation, groups)
+    }
+
+    fn max_pool2d(
+        input: Self::TensorPrimitive<4>,
+        kernel_size: [usize; 2],
+        stride: [usize; 2],
+        padding: [usize; 2],
+    ) -> Self::TensorPrimitive<4> {
+        ops::pooling::max_pool2d(&input, kernel_size, stride, padding)
+    }
+
+    fn avg_pool2d(
+        input: Self::TensorPrimitive<4>,
+        kernel_size: [usize; 2],
+        stride: [usize; 2],
+        padding: [usize; 2],
+        count_include_pad: bool,
+    ) -> Self::TensorPrimitive<4> {
+        ops::pooling::avg_pool2d(&input, kernel_size, stride, padding, count_include_pad)
+    }
+
+    fn device<const D: usize>(tensor: &Self::TensorPrimitive<D>) -> Self::Device {
+        tensor.device.clone()
+    }
+
+    fn to_device<const D: usize>(
+        tensor: Self::TensorPrimitive<D>,
+        _device: &Self::Device,
+    ) -> Self::TensorPrimitive<D> {
+        tensor
+    }
+
+    fn to_data<const D: usize>(tensor: &Self::TensorPrimitive<D>) -> Vec<Self::FloatElem> {
+        tensor.data.clone()
+    }
+
+    fn sync(_device: &Self::Device) {}
+}
+
 /// Type alias for training with CPU + autodiff.
 pub type CpuTraining = CpuBackend;
 
 /// Type alias for inference with CPU (no autodiff overhead).
 pub type CpuInference = CpuBackend;
+
+#[cfg(test)]
+mod f64_tests {
+    use super::*;
+
+    #[test]
+    fn cpu_backend_f64_add_and_matmul() {
+        let dev = CpuDevice::new();
+
+        // Element-wise add in f64.
+        let a = CpuBackendF64::from_data(&[1.0_f64, 2.0, 3.0, 4.0], [2, 2], &dev);
+        let b = CpuBackendF64::from_data(&[10.0_f64, 20.0, 30.0, 40.0], [2, 2], &dev);
+        let sum = CpuBackendF64::add(a, b);
+        assert_eq!(sum.to_vec(), vec![11.0_f64, 22.0, 33.0, 44.0]);
+
+        // 2x2 matmul in f64: [[1,2],[3,4]] @ [[5,6],[7,8]] = [[19,22],[43,50]].
+        let x = CpuBackendF64::from_data(&[1.0_f64, 2.0, 3.0, 4.0], [2, 2], &dev);
+        let y = CpuBackendF64::from_data(&[5.0_f64, 6.0, 7.0, 8.0], [2, 2], &dev);
+        let z = CpuBackendF64::matmul(x, y);
+        assert_eq!(z.to_vec(), vec![19.0_f64, 22.0, 43.0, 50.0]);
+    }
+
+    #[test]
+    fn cpu_backend_f64_exceeds_f32_precision() {
+        // A value that f32 cannot represent but f64 can: 1 + 2^-30.
+        // The f64 backend must preserve it; the f32 backend would round to 1.0.
+        let dev = CpuDevice::new();
+        let eps = 2.0_f64.powi(-30); // ~9.3e-10, below f32 epsilon (~1.2e-7)
+
+        let t = CpuBackendF64::from_data(&[1.0_f64 + eps], [1], &dev);
+        let one = CpuBackendF64::from_data(&[1.0_f64], [1], &dev);
+        let diff = CpuBackendF64::sub(t, one).to_vec()[0];
+        // f64 retains the tiny difference exactly; f32 would yield 0.0.
+        assert!(diff > 0.0, "f64 backend lost sub-f32 precision: diff={diff:e}");
+        assert!((diff - eps).abs() < 1e-18, "diff={diff:e} expected {eps:e}");
+    }
+}

crates/core/rtx-backend-cpu/src/ops/activation.rs

@@ -1,96 +1,121 @@
 //! Activation function operations.
 
-use crate::CpuTensorPrimitive;
+use crate::{CpuFloat, CpuTensorPrimitive};
 use rayon::prelude::*;
 
 /// GELU activation function.
-pub fn gelu<const D: usize>(tensor: &CpuTensorPrimitive<D>) -> CpuTensorPrimitive<D> {
+pub fn gelu<const D: usize, E: CpuFloat>(
-    const SQRT_2_OVER_PI: f32 = 0.797_884_6;
+    tensor: &CpuTensorPrimitive<D, E>,
+) -> CpuTensorPrimitive<D, E> {
+    let half = E::from(0.5).unwrap();
+    let one = E::one();
+    let c = E::from(0.044715).unwrap();
+    let sqrt_2_over_pi = E::from(0.797_884_6_f64).unwrap();
 
-    let result: Vec<f32> = tensor
+    let result: Vec<E> = tensor
         .data
         .par_iter()
-        .map(|&x| 0.5 * x * (1.0 + (SQRT_2_OVER_PI * (x + 0.044715 * x * x * x)).tanh()))
+        .map(|&x| half * x * (one + (sqrt_2_over_pi * (x + c * x * x * x)).tanh()))
         .collect();
 
     CpuTensorPrimitive::new(result, tensor.shape, tensor.device.clone())
 }
 
 /// SiLU (Swish) activation function.
-pub fn silu<const D: usize>(tensor: &CpuTensorPrimitive<D>) -> CpuTensorPrimitive<D> {
+pub fn silu<const D: usize, E: CpuFloat>(
-    let result: Vec<f32> = tensor
+    tensor: &CpuTensorPrimitive<D, E>,
+) -> CpuTensorPrimitive<D, E> {
+    let one = E::one();
+    let result: Vec<E> = tensor
         .data
         .par_iter()
-        .map(|&x| x * (1.0 / (1.0 + (-x).exp())))
+        .map(|&x| x * (one / (one + (-x).exp())))
         .collect();
 
     CpuTensorPrimitive::new(result, tensor.shape, tensor.device.clone())
 }
 
 /// ReLU activation function.
-pub fn relu<const D: usize>(tensor: &CpuTensorPrimitive<D>) -> CpuTensorPrimitive<D> {
+pub fn relu<const D: usize, E: CpuFloat>(
-    let result: Vec<f32> = tensor.data.par_iter().map(|&x| x.max(0.0)).collect();
+    tensor: &CpuTensorPrimitive<D, E>,
+) -> CpuTensorPrimitive<D, E> {
+    let zero = E::zero();
+    let result: Vec<E> = tensor.data.par_iter().map(|&x| x.max(zero)).collect();
     CpuTensorPrimitive::new(result, tensor.shape, tensor.device.clone())
 }
 
 /// Leaky ReLU activation function.
-pub fn leaky_relu<const D: usize>(
+pub fn leaky_relu<const D: usize, E: CpuFloat>(
-    tensor: &CpuTensorPrimitive<D>,
+    tensor: &CpuTensorPrimitive<D, E>,
-    negative_slope: f32,
+    negative_slope: E,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
-    let result: Vec<f32> = tensor
+    let zero = E::zero();
+    let result: Vec<E> = tensor
         .data
         .par_iter()
-        .map(|&x| if x >= 0.0 { x } else { negative_slope * x })
+        .map(|&x| if x >= zero { x } else { negative_slope * x })
         .collect();
     CpuTensorPrimitive::new(result, tensor.shape, tensor.device.clone())
 }
 
 /// ELU activation function.
-pub fn elu<const D: usize>(tensor: &CpuTensorPrimitive<D>, alpha: f32) -> CpuTensorPrimitive<D> {
+pub fn elu<const D: usize, E: CpuFloat>(
-    let result: Vec<f32> = tensor
+    tensor: &CpuTensorPrimitive<D, E>,
+    alpha: E,
+) -> CpuTensorPrimitive<D, E> {
+    let zero = E::zero();
+    let one = E::one();
+    let result: Vec<E> = tensor
         .data
         .par_iter()
-        .map(|&x| if x > 0.0 { x } else { alpha * (x.exp() - 1.0) })
+        .map(|&x| if x > zero { x } else { alpha * (x.exp() - one) })
         .collect();
     CpuTensorPrimitive::new(result, tensor.shape, tensor.device.clone())
 }
 
 /// Sigmoid activation function.
-pub fn sigmoid<const D: usize>(tensor: &CpuTensorPrimitive<D>) -> CpuTensorPrimitive<D> {
+pub fn sigmoid<const D: usize, E: CpuFloat>(
-    let result: Vec<f32> = tensor
+    tensor: &CpuTensorPrimitive<D, E>,
+) -> CpuTensorPrimitive<D, E> {
+    let one = E::one();
+    let result: Vec<E> = tensor
         .data
         .par_iter()
-        .map(|&x| 1.0 / (1.0 + (-x).exp()))
+        .map(|&x| one / (one + (-x).exp()))
         .collect();
     CpuTensorPrimitive::new(result, tensor.shape, tensor.device.clone())
 }
 
 /// Tanh activation function.
-pub fn tanh<const D: usize>(tensor: &CpuTensorPrimitive<D>) -> CpuTensorPrimitive<D> {
+pub fn tanh<const D: usize, E: CpuFloat>(
-    let result: Vec<f32> = tensor.data.par_iter().map(|&x| x.tanh()).collect();
+    tensor: &CpuTensorPrimitive<D, E>,
+) -> CpuTensorPrimitive<D, E> {
+    let result: Vec<E> = tensor.data.par_iter().map(|&x| x.tanh()).collect();
     CpuTensorPrimitive::new(result, tensor.shape, tensor.device.clone())
 }
 
 /// Softmax along a dimension.
-pub fn softmax<const D: usize>(
+pub fn softmax<const D: usize, E: CpuFloat>(
-    tensor: &CpuTensorPrimitive<D>,
+    tensor: &CpuTensorPrimitive<D, E>,
     dim: usize,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
     assert!(dim < D, "Dimension out of range");
 
     let dim_size = tensor.shape[D - 1];
     let batch_size = tensor.numel() / dim_size;
 
-    let mut result = vec![0.0f32; tensor.numel()];
+    let mut result = vec![E::zero(); tensor.numel()];
 
     for b in 0..batch_size {
         let offset = b * dim_size;
         let slice = &tensor.data[offset..offset + dim_size];
 
-        let max_val = slice.iter().copied().fold(f32::NEG_INFINITY, f32::max);
+        let max_val = slice
-        let exp_vals: Vec<f32> = slice.iter().map(|&x| (x - max_val).exp()).collect();
+            .iter()
-        let sum: f32 = exp_vals.iter().sum();
+            .copied()
+            .fold(E::neg_infinity(), |a, b| a.max(b));
+        let exp_vals: Vec<E> = slice.iter().map(|&x| (x - max_val).exp()).collect();
+        let sum: E = exp_vals.iter().copied().fold(E::zero(), |a, b| a + b);
 
         for (i, &exp_val) in exp_vals.iter().enumerate() {
             result[offset + i] = exp_val / sum;
@@ -101,23 +126,30 @@ pub fn softmax<const D: usize>(
 }
 
 /// Log-softmax along a dimension.
-pub fn log_softmax<const D: usize>(
+pub fn log_softmax<const D: usize, E: CpuFloat>(
-    tensor: &CpuTensorPrimitive<D>,
+    tensor: &CpuTensorPrimitive<D, E>,
     dim: usize,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
     assert!(dim < D, "Dimension out of range");
 
     let dim_size = tensor.shape[D - 1];
     let batch_size = tensor.numel() / dim_size;
 
-    let mut result = vec![0.0f32; tensor.numel()];
+    let mut result = vec![E::zero(); tensor.numel()];
 
     for b in 0..batch_size {
         let offset = b * dim_size;
         let slice = &tensor.data[offset..offset + dim_size];
 
-        let max_val = slice.iter().copied().fold(f32::NEG_INFINITY, f32::max);
+        let max_val = slice
-        let log_sum_exp: f32 = slice.iter().map(|&x| (x - max_val).exp()).sum::<f32>().ln();
+            .iter()
+            .copied()
+            .fold(E::neg_infinity(), |a, b| a.max(b));
+        let log_sum_exp: E = slice
+            .iter()
+            .map(|&x| (x - max_val).exp())
+            .fold(E::zero(), |a, b| a + b)
+            .ln();
 
         for (i, &x) in slice.iter().enumerate() {
             result[offset + i] = x - max_val - log_sum_exp;

crates/core/rtx-backend-cpu/src/ops/attention.rs

@@ -1,37 +1,37 @@
 //! Attention operations.
 
-use crate::CpuTensorPrimitive;
+use crate::{CpuFloat, CpuTensorPrimitive};
 
 /// Flash Attention implementation for CPU.
 ///
 /// Reference implementation of scaled dot-product attention.
-pub fn flash_attention(
+pub fn flash_attention<E: CpuFloat>(
-    query: &CpuTensorPrimitive<4>, // [batch, heads, seq_len, head_dim]
+    query: &CpuTensorPrimitive<4, E>, // [batch, heads, seq_len, head_dim]
-    key: &CpuTensorPrimitive<4>,
+    key: &CpuTensorPrimitive<4, E>,
-    value: &CpuTensorPrimitive<4>,
+    value: &CpuTensorPrimitive<4, E>,
-    mask: Option<&CpuTensorPrimitive<4>>,
+    mask: Option<&CpuTensorPrimitive<4, E>>,
-    scale: f32,
+    scale: E,
     causal: bool,
-) -> CpuTensorPrimitive<4> {
+) -> CpuTensorPrimitive<4, E> {
     let [batch, heads, seq_q, head_dim] = query.shape;
     let [_, _, seq_k, _] = key.shape;
 
-    let mut output = vec![0.0f32; batch * heads * seq_q * head_dim];
+    let mut output = vec![E::zero(); batch * heads * seq_q * head_dim];
 
     for b in 0..batch {
         for h in 0..heads {
             // Compute attention scores: Q @ K^T
-            let mut scores = vec![0.0f32; seq_q * seq_k];
+            let mut scores = vec![E::zero(); seq_q * seq_k];
 
             for i in 0..seq_q {
                 for j in 0..seq_k {
-                    let mut dot = 0.0f32;
+                    let mut dot = E::zero();
                     for d in 0..head_dim {
                         let q_idx =
                             b * heads * seq_q * head_dim + h * seq_q * head_dim + i * head_dim + d;
                         let k_idx =
                             b * heads * seq_k * head_dim + h * seq_k * head_dim + j * head_dim + d;
-                        dot += query.data[q_idx] * key.data[k_idx];
+                        dot = dot + query.data[q_idx] * key.data[k_idx];
                     }
                     scores[i * seq_k + j] = dot * scale;
                 }
@@ -42,7 +42,7 @@ pub fn flash_attention(
                 for i in 0..seq_q {
                     for j in 0..seq_k {
                         if j > i {
-                            scores[i * seq_k + j] = f32::NEG_INFINITY;
+                            scores[i * seq_k + j] = E::neg_infinity();
                         }
                     }
                 }
@@ -55,7 +55,8 @@ pub fn flash_attention(
                         let mask_idx =
                             b * heads * seq_q * seq_k + h * seq_q * seq_k + i * seq_k + j;
                         if mask.data.len() > mask_idx {
-                            scores[i * seq_k + j] += mask.data[mask_idx];
+                            let s = i * seq_k + j;
+                            scores[s] = scores[s] + mask.data[mask_idx];
                         }
                     }
                 }
@@ -66,9 +67,9 @@ pub fn flash_attention(
                 let row_start = i * seq_k;
                 let row = &mut scores[row_start..row_start + seq_k];
 
-                let max_val = row.iter().copied().fold(f32::NEG_INFINITY, f32::max);
+                let max_val = row.iter().copied().fold(E::neg_infinity(), |a, b| a.max(b));
-                let exp_vals: Vec<f32> = row.iter().map(|&x| (x - max_val).exp()).collect();
+                let exp_vals: Vec<E> = row.iter().map(|&x| (x - max_val).exp()).collect();
-                let sum: f32 = exp_vals.iter().sum();
+                let sum: E = exp_vals.iter().copied().fold(E::zero(), |a, b| a + b);
 
                 for (j, &exp_val) in exp_vals.iter().enumerate() {
                     row[j] = exp_val / sum;
@@ -78,11 +79,11 @@ pub fn flash_attention(
             // Compute output: attention_weights @ V
             for i in 0..seq_q {
                 for d in 0..head_dim {
-                    let mut sum = 0.0f32;
+                    let mut sum = E::zero();
                     for j in 0..seq_k {
                         let v_idx =
                             b * heads * seq_k * head_dim + h * seq_k * head_dim + j * head_dim + d;
-                        sum += scores[i * seq_k + j] * value.data[v_idx];
+                        sum = sum + scores[i * seq_k + j] * value.data[v_idx];
                     }
                     let out_idx =
                         b * heads * seq_q * head_dim + h * seq_q * head_dim + i * head_dim + d;
@@ -100,13 +101,13 @@ pub fn flash_attention(
 }
 
 /// Rotary Position Embedding (RoPE).
-pub fn rope<const D: usize>(
+pub fn rope<const D: usize, E: CpuFloat>(
-    tensor: &CpuTensorPrimitive<D>,
+    tensor: &CpuTensorPrimitive<D, E>,
-    cos: &CpuTensorPrimitive<2>,
+    cos: &CpuTensorPrimitive<2, E>,
-    sin: &CpuTensorPrimitive<2>,
+    sin: &CpuTensorPrimitive<2, E>,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
     let numel = tensor.numel();
-    let mut result = vec![0.0f32; numel];
+    let mut result = vec![E::zero(); numel];
 
     // cos/sin shape is [seq_len, head_dim]
     let [seq_len, head_dim] = cos.shape;
@@ -136,13 +137,16 @@ pub fn rope<const D: usize>(
 }
 
 /// Create causal attention mask.
-pub fn causal_mask(device: &crate::CpuDevice, seq_len: usize) -> CpuTensorPrimitive<2> {
+pub fn causal_mask<E: CpuFloat>(
-    let mut mask = vec![0.0f32; seq_len * seq_len];
+    device: &crate::CpuDevice,
+    seq_len: usize,
+) -> CpuTensorPrimitive<2, E> {
+    let mut mask = vec![E::zero(); seq_len * seq_len];
 
     for i in 0..seq_len {
         for j in 0..seq_len {
             if j > i {
-                mask[i * seq_len + j] = f32::NEG_INFINITY;
+                mask[i * seq_len + j] = E::neg_infinity();
             }
         }
     }

crates/core/rtx-backend-cpu/src/ops/basic.rs

@@ -1,16 +1,16 @@
 //! Basic element-wise operations.
 
-use crate::CpuTensorPrimitive;
+use crate::{CpuFloat, CpuTensorPrimitive};
 use rayon::prelude::*;
 
 /// Element-wise addition.
-pub fn add<const D: usize>(
+pub fn add<const D: usize, E: CpuFloat>(
-    lhs: &CpuTensorPrimitive<D>,
+    lhs: &CpuTensorPrimitive<D, E>,
-    rhs: &CpuTensorPrimitive<D>,
+    rhs: &CpuTensorPrimitive<D, E>,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
     assert_eq!(lhs.shape, rhs.shape, "Shapes must match for addition");
 
-    let result: Vec<f32> = lhs
+    let result: Vec<E> = lhs
         .data
         .par_iter()
         .zip(rhs.data.par_iter())
@@ -21,13 +21,13 @@ pub fn add<const D: usize>(
 }
 
 /// Element-wise subtraction.
-pub fn sub<const D: usize>(
+pub fn sub<const D: usize, E: CpuFloat>(
-    lhs: &CpuTensorPrimitive<D>,
+    lhs: &CpuTensorPrimitive<D, E>,
-    rhs: &CpuTensorPrimitive<D>,
+    rhs: &CpuTensorPrimitive<D, E>,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
     assert_eq!(lhs.shape, rhs.shape, "Shapes must match for subtraction");
 
-    let result: Vec<f32> = lhs
+    let result: Vec<E> = lhs
         .data
         .par_iter()
         .zip(rhs.data.par_iter())
@@ -38,13 +38,13 @@ pub fn sub<const D: usize>(
 }
 
 /// Element-wise multiplication.
-pub fn mul<const D: usize>(
+pub fn mul<const D: usize, E: CpuFloat>(
-    lhs: &CpuTensorPrimitive<D>,
+    lhs: &CpuTensorPrimitive<D, E>,
-    rhs: &CpuTensorPrimitive<D>,
+    rhs: &CpuTensorPrimitive<D, E>,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
     assert_eq!(lhs.shape, rhs.shape, "Shapes must match for multiplication");
 
-    let result: Vec<f32> = lhs
+    let result: Vec<E> = lhs
         .data
         .par_iter()
         .zip(rhs.data.par_iter())
@@ -55,13 +55,13 @@ pub fn mul<const D: usize>(
 }
 
 /// Element-wise division.
-pub fn div<const D: usize>(
+pub fn div<const D: usize, E: CpuFloat>(
-    lhs: &CpuTensorPrimitive<D>,
+    lhs: &CpuTensorPrimitive<D, E>,
-    rhs: &CpuTensorPrimitive<D>,
+    rhs: &CpuTensorPrimitive<D, E>,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
     assert_eq!(lhs.shape, rhs.shape, "Shapes must match for division");
 
-    let result: Vec<f32> = lhs
+    let result: Vec<E> = lhs
         .data
         .par_iter()
         .zip(rhs.data.par_iter())
@@ -72,35 +72,35 @@ pub fn div<const D: usize>(
 }
 
 /// Scalar addition.
-pub fn add_scalar<const D: usize>(
+pub fn add_scalar<const D: usize, E: CpuFloat>(
-    tensor: &CpuTensorPrimitive<D>,
+    tensor: &CpuTensorPrimitive<D, E>,
-    scalar: f32,
+    scalar: E,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
-    let result: Vec<f32> = tensor.data.par_iter().map(|&x| x + scalar).collect();
+    let result: Vec<E> = tensor.data.par_iter().map(|&x| x + scalar).collect();
 
     CpuTensorPrimitive::new(result, tensor.shape, tensor.device.clone())
 }
 
 /// Scalar multiplication.
-pub fn mul_scalar<const D: usize>(
+pub fn mul_scalar<const D: usize, E: CpuFloat>(
-    tensor: &CpuTensorPrimitive<D>,
+    tensor: &CpuTensorPrimitive<D, E>,
-    scalar: f32,
+    scalar: E,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
-    let result: Vec<f32> = tensor.data.par_iter().map(|&x| x * scalar).collect();
+    let result: Vec<E> = tensor.data.par_iter().map(|&x| x * scalar).collect();
 
     CpuTensorPrimitive::new(result, tensor.shape, tensor.device.clone())
 }
 
 /// Greater than scalar comparison.
 /// Returns 1.0 where element > value, 0.0 otherwise.
-pub fn gt_scalar<const D: usize>(
+pub fn gt_scalar<const D: usize, E: CpuFloat>(
-    tensor: &CpuTensorPrimitive<D>,
+    tensor: &CpuTensorPrimitive<D, E>,
-    value: f32,
+    value: E,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
-    let result: Vec<f32> = tensor
+    let result: Vec<E> = tensor
         .data
         .par_iter()
-        .map(|&x| if x > value { 1.0 } else { 0.0 })
+        .map(|&x| if x > value { E::one() } else { E::zero() })
         .collect();
 
     CpuTensorPrimitive::new(result, tensor.shape, tensor.device.clone())

crates/core/rtx-backend-cpu/src/ops/conv.rs

@@ -2,7 +2,7 @@
 //!
 //! Implements 2D convolution using the im2col approach for efficiency.
 
-use crate::CpuTensorPrimitive;
+use crate::{CpuFloat, CpuTensorPrimitive};
 use rayon::prelude::*;
 
 /// 2D Convolution using im2col + GEMM approach.
@@ -15,15 +15,15 @@ use rayon::prelude::*;
 /// * `padding` - [pad_h, pad_w]
 /// * `dilation` - [dilation_h, dilation_w]
 /// * `groups` - Number of groups
-pub fn conv2d(
+pub fn conv2d<E: CpuFloat>(
-    input: &CpuTensorPrimitive<4>,
+    input: &CpuTensorPrimitive<4, E>,
-    weight: &CpuTensorPrimitive<4>,
+    weight: &CpuTensorPrimitive<4, E>,
-    bias: Option<&CpuTensorPrimitive<1>>,
+    bias: Option<&CpuTensorPrimitive<1, E>>,
     stride: [usize; 2],
     padding: [usize; 2],
     dilation: [usize; 2],
     groups: usize,
-) -> CpuTensorPrimitive<4> {
+) -> CpuTensorPrimitive<4, E> {
     let [batch, in_channels, in_h, in_w] = input.shape;
     let [out_channels, in_channels_per_group, kernel_h, kernel_w] = weight.shape;
 
@@ -46,7 +46,7 @@ pub fn conv2d(
 
     // Output shape: [batch, out_channels, out_h, out_w]
     let output_size = batch * out_channels * out_h * out_w;
-    let mut output = vec![0.0f32; output_size];
+    let mut output = vec![E::zero(); output_size];
 
     // Process each batch in parallel
     output
@@ -68,7 +68,7 @@ pub fn conv2d(
                     // For each output position
                     for oh in 0..out_h {
                         for ow in 0..out_w {
-                            let mut sum = 0.0f32;
+                            let mut sum = E::zero();
 
                             // Convolve over input channels and kernel
                             for ic in 0..in_channels_per_group {
@@ -102,7 +102,7 @@ pub fn conv2d(
                                                 + kh * kernel_w
                                                 + kw;
 
-                                            sum += input.data[in_idx] * weight.data[w_idx];
+                                            sum = sum + input.data[in_idx] * weight.data[w_idx];
                                         }
                                     }
                                 }
@@ -110,7 +110,7 @@ pub fn conv2d(
 
                             // Add bias if present
                             if let Some(bias) = bias {
-                                sum += bias.data[global_oc];
+                                sum = sum + bias.data[global_oc];
                             }
 
                             // Store result

crates/core/rtx-backend-cpu/src/ops/creation.rs

@@ -1,58 +1,75 @@
 //! Tensor creation operations.
 
-use crate::{CpuDevice, CpuTensorPrimitive};
+use crate::{CpuDevice, CpuFloat, CpuTensorPrimitive};
 use rand::Rng;
 use rand_distr::{Distribution, StandardNormal};
 
 /// Create a tensor filled with zeros.
-pub fn zeros<const D: usize>(shape: [usize; D], device: &CpuDevice) -> CpuTensorPrimitive<D> {
+pub fn zeros<const D: usize, E: CpuFloat>(
+    shape: [usize; D],
+    device: &CpuDevice,
+) -> CpuTensorPrimitive<D, E> {
     let numel: usize = shape.iter().product();
-    let data = vec![0.0f32; numel];
+    let data = vec![E::zero(); numel];
     CpuTensorPrimitive::new(data, shape, device.clone())
 }
 
 /// Create a tensor filled with ones.
-pub fn ones<const D: usize>(shape: [usize; D], device: &CpuDevice) -> CpuTensorPrimitive<D> {
+pub fn ones<const D: usize, E: CpuFloat>(
+    shape: [usize; D],
+    device: &CpuDevice,
+) -> CpuTensorPrimitive<D, E> {
     let numel: usize = shape.iter().product();
-    let data = vec![1.0f32; numel];
+    let data = vec![E::one(); numel];
     CpuTensorPrimitive::new(data, shape, device.clone())
 }
 
 /// Create a tensor filled with a specific value.
-pub fn full<const D: usize>(
+pub fn full<const D: usize, E: CpuFloat>(
     shape: [usize; D],
-    fill_value: f32,
+    fill_value: E,
     device: &CpuDevice,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
     let numel: usize = shape.iter().product();
     let data = vec![fill_value; numel];
     CpuTensorPrimitive::new(data, shape, device.clone())
 }
 
 /// Create a tensor with uniform random values in [0, 1).
-pub fn rand<const D: usize>(shape: [usize; D], device: &CpuDevice) -> CpuTensorPrimitive<D> {
+pub fn rand<const D: usize, E: CpuFloat>(
+    shape: [usize; D],
+    device: &CpuDevice,
+) -> CpuTensorPrimitive<D, E> {
     let numel: usize = shape.iter().product();
     let mut rng = super::get_rng();
-    let data: Vec<f32> = (0..numel).map(|_| rng.r#gen::<f32>()).collect();
+    let data: Vec<E> = (0..numel)
+        .map(|_| E::from(rng.r#gen::<f64>()).unwrap())
+        .collect();
     CpuTensorPrimitive::new(data, shape, device.clone())
 }
 
 /// Create a tensor with normal distributed random values (mean=0, std=1).
-pub fn randn<const D: usize>(shape: [usize; D], device: &CpuDevice) -> CpuTensorPrimitive<D> {
+pub fn randn<const D: usize, E: CpuFloat>(
+    shape: [usize; D],
+    device: &CpuDevice,
+) -> CpuTensorPrimitive<D, E> {
     let numel: usize = shape.iter().product();
     let mut rng = super::get_rng();
-    let data: Vec<f32> = (0..numel)
+    let data: Vec<E> = (0..numel)
-        .map(|_| StandardNormal.sample(&mut rng))
+        .map(|_| {
+            let s: f64 = StandardNormal.sample(&mut rng);
+            E::from(s).unwrap()
+        })
         .collect();
     CpuTensorPrimitive::new(data, shape, device.clone())
 }
 
 /// Create a tensor from existing data.
-pub fn from_data<const D: usize>(
+pub fn from_data<const D: usize, E: CpuFloat>(
-    data: &[f32],
+    data: &[E],
     shape: [usize; D],
     device: &CpuDevice,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
     let numel: usize = shape.iter().product();
     assert_eq!(data.len(), numel, "Data length must match shape");
     CpuTensorPrimitive::new(data.to_vec(), shape, device.clone())

crates/core/rtx-backend-cpu/src/ops/gemm.rs

@@ -1,18 +1,21 @@
 //! Matrix multiplication operations.
 
-use crate::CpuTensorPrimitive;
+use crate::{CpuFloat, CpuTensorPrimitive};
 use rayon::prelude::*;
 
 /// Matrix multiplication with cache blocking.
 ///
 /// Uses a blocked algorithm for better cache utilization.
-pub fn matmul(lhs: &CpuTensorPrimitive<2>, rhs: &CpuTensorPrimitive<2>) -> CpuTensorPrimitive<2> {
+pub fn matmul<E: CpuFloat>(
+    lhs: &CpuTensorPrimitive<2, E>,
+    rhs: &CpuTensorPrimitive<2, E>,
+) -> CpuTensorPrimitive<2, E> {
     let [m, k1] = lhs.shape;
     let [k2, n] = rhs.shape;
     assert_eq!(k1, k2, "Inner dimensions must match: {k1} != {k2}");
 
     let k = k1;
-    let mut result = vec![0.0f32; m * n];
+    let mut result = vec![E::zero(); m * n];
 
     // Block size for cache efficiency
     const BLOCK_SIZE: usize = 64;
@@ -29,7 +32,7 @@ pub fn matmul(lhs: &CpuTensorPrimitive<2>, rhs: &CpuTensorPrimitive<2>) -> CpuTe
                     for j in j_block..j_end {
                         let mut sum = result[i * n + j];
                         for kk in k_block..k_end {
-                            sum += lhs.data[i * k + kk] * rhs.data[kk * n + j];
+                            sum = sum + lhs.data[i * k + kk] * rhs.data[kk * n + j];
                         }
                         result[i * n + j] = sum;
                     }
@@ -42,14 +45,17 @@ pub fn matmul(lhs: &CpuTensorPrimitive<2>, rhs: &CpuTensorPrimitive<2>) -> CpuTe
 }
 
 /// Batched matrix multiplication.
-pub fn bmm(lhs: &CpuTensorPrimitive<3>, rhs: &CpuTensorPrimitive<3>) -> CpuTensorPrimitive<3> {
+pub fn bmm<E: CpuFloat>(
+    lhs: &CpuTensorPrimitive<3, E>,
+    rhs: &CpuTensorPrimitive<3, E>,
+) -> CpuTensorPrimitive<3, E> {
     let [batch, m, k1] = lhs.shape;
     let [batch2, k2, n] = rhs.shape;
     assert_eq!(batch, batch2, "Batch sizes must match");
     assert_eq!(k1, k2, "Inner dimensions must match");
 
     let k = k1;
-    let mut result = vec![0.0f32; batch * m * n];
+    let mut result = vec![E::zero(); batch * m * n];
 
     // Parallelize over batch dimension
     result
@@ -61,10 +67,10 @@ pub fn bmm(lhs: &CpuTensorPrimitive<3>, rhs: &CpuTensorPrimitive<3>) -> CpuTenso
 
             for i in 0..m {
                 for j in 0..n {
-                    let mut sum = 0.0f32;
+                    let mut sum = E::zero();
                     for kk in 0..k {
-                        sum +=
+                        sum = sum
-                            lhs.data[lhs_offset + i * k + kk] * rhs.data[rhs_offset + kk * n + j];
+                            + lhs.data[lhs_offset + i * k + kk] * rhs.data[rhs_offset + kk * n + j];
                     }
                     res_batch[i * n + j] = sum;
                 }

crates/core/rtx-backend-cpu/src/ops/normalization.rs

@@ -1,31 +1,35 @@
 //! Normalization operations.
 
-use crate::CpuTensorPrimitive;
+use crate::{CpuFloat, CpuTensorPrimitive};
 
 /// Layer normalization.
-pub fn layer_norm<const D: usize>(
+pub fn layer_norm<const D: usize, E: CpuFloat>(
-    tensor: &CpuTensorPrimitive<D>,
+    tensor: &CpuTensorPrimitive<D, E>,
-    weight: &CpuTensorPrimitive<1>,
+    weight: &CpuTensorPrimitive<1, E>,
-    bias: Option<&CpuTensorPrimitive<1>>,
+    bias: Option<&CpuTensorPrimitive<1, E>>,
-    eps: f32,
+    eps: E,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
     let norm_size = weight.numel();
     let batch_size = tensor.numel() / norm_size;
+    let norm_size_e = E::from(norm_size).unwrap();
 
-    let mut result = vec![0.0f32; tensor.numel()];
+    let mut result = vec![E::zero(); tensor.numel()];
 
     for b in 0..batch_size {
         let offset = b * norm_size;
         let slice = &tensor.data[offset..offset + norm_size];
 
         // Compute mean
-        let mean: f32 = slice.iter().sum::<f32>() / norm_size as f32;
+        let mean: E = slice.iter().copied().fold(E::zero(), |a, b| a + b) / norm_size_e;
 
         // Compute variance
-        let variance: f32 =
+        let variance: E = slice
-            slice.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / norm_size as f32;
+            .iter()
+            .map(|&x| (x - mean).powi(2))
+            .fold(E::zero(), |a, b| a + b)
+            / norm_size_e;
 
-        let inv_std = 1.0 / (variance + eps).sqrt();
+        let inv_std = E::one() / (variance + eps).sqrt();
 
         // Normalize and apply affine transform
         for i in 0..norm_size {
@@ -43,25 +47,31 @@ pub fn layer_norm<const D: usize>(
 }
 
 /// RMS normalization (used in LLaMA and other models).
-pub fn rms_norm<const D: usize>(
+pub fn rms_norm<const D: usize, E: CpuFloat>(
-    tensor: &CpuTensorPrimitive<D>,
+    tensor: &CpuTensorPrimitive<D, E>,
-    weight: &CpuTensorPrimitive<1>,
+    weight: &CpuTensorPrimitive<1, E>,
-    eps: f32,
+    eps: E,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
     let norm_size = weight.numel();
     let batch_size = tensor.numel() / norm_size;
+    let norm_size_e = E::from(norm_size).unwrap();
 
-    let mut result = vec![0.0f32; tensor.numel()];
+    let mut result = vec![E::zero(); tensor.numel()];
 
     for b in 0..batch_size {
         let offset = b * norm_size;
         let slice = &tensor.data[offset..offset + norm_size];
 
         // Compute RMS (root mean square)
-        let rms: f32 =
+        let rms: E = (slice
-            (slice.iter().map(|x| x.powi(2)).sum::<f32>() / norm_size as f32 + eps).sqrt();
+            .iter()
+            .map(|&x| x.powi(2))
+            .fold(E::zero(), |a, b| a + b)
+            / norm_size_e
+            + eps)
+            .sqrt();
 
-        let inv_rms = 1.0 / rms;
+        let inv_rms = E::one() / rms;
 
         // Normalize and apply scale
         for i in 0..norm_size {
@@ -75,30 +85,30 @@ pub fn rms_norm<const D: usize>(
 }
 
 /// Batch normalization (inference mode).
-pub fn batch_norm<const D: usize>(
+pub fn batch_norm<const D: usize, E: CpuFloat>(
-    tensor: &CpuTensorPrimitive<D>,
+    tensor: &CpuTensorPrimitive<D, E>,
-    running_mean: &CpuTensorPrimitive<1>,
+    running_mean: &CpuTensorPrimitive<1, E>,
-    running_var: &CpuTensorPrimitive<1>,
+    running_var: &CpuTensorPrimitive<1, E>,
-    weight: Option<&CpuTensorPrimitive<1>>,
+    weight: Option<&CpuTensorPrimitive<1, E>>,
-    bias: Option<&CpuTensorPrimitive<1>>,
+    bias: Option<&CpuTensorPrimitive<1, E>>,
-    eps: f32,
+    eps: E,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
     assert!(D >= 2, "BatchNorm requires at least 2D tensor");
 
     let num_channels = tensor.shape[1];
     let spatial_size: usize = tensor.shape[2..].iter().product();
     let batch_size = tensor.shape[0];
 
-    let mut result = vec![0.0f32; tensor.numel()];
+    let mut result = vec![E::zero(); tensor.numel()];
 
     for n in 0..batch_size {
         for c in 0..num_channels {
             let mean = running_mean.data[c];
             let var = running_var.data[c];
-            let inv_std = 1.0 / (var + eps).sqrt();
+            let inv_std = E::one() / (var + eps).sqrt();
 
-            let gamma = weight.map_or(1.0, |w| w.data[c]);
+            let gamma = weight.map_or(E::one(), |w| w.data[c]);
-            let beta = bias.map_or(0.0, |b| b.data[c]);
+            let beta = bias.map_or(E::zero(), |b| b.data[c]);
 
             for s in 0..spatial_size {
                 let idx = n * num_channels * spatial_size + c * spatial_size + s;

crates/core/rtx-backend-cpu/src/ops/pooling.rs

@@ -1,6 +1,6 @@
 //! Pooling operations for CPU backend.
 
-use crate::CpuTensorPrimitive;
+use crate::{CpuFloat, CpuTensorPrimitive};
 use rayon::prelude::*;
 
 /// Compute contiguous strides for a given shape.
@@ -18,12 +18,12 @@ fn compute_strides<const D: usize>(shape: [usize; D]) -> [usize; D] {
 ///
 /// Input: [batch, channels, height, width]
 /// Output: [batch, channels, out_height, out_width]
-pub fn max_pool2d(
+pub fn max_pool2d<E: CpuFloat>(
-    input: &CpuTensorPrimitive<4>,
+    input: &CpuTensorPrimitive<4, E>,
     kernel_size: [usize; 2],
     stride: [usize; 2],
     padding: [usize; 2],
-) -> CpuTensorPrimitive<4> {
+) -> CpuTensorPrimitive<4, E> {
     let [batch, channels, in_h, in_w] = input.shape;
     let [kh, kw] = kernel_size;
     let [sh, sw] = stride;
@@ -37,7 +37,7 @@ pub fn max_pool2d(
     let output_size = batch * channels * out_h * out_w;
 
     // Process each output position in parallel
-    let output_data: Vec<f32> = (0..output_size)
+    let output_data: Vec<E> = (0..output_size)
         .into_par_iter()
         .map(|idx| {
             // Calculate position in output
@@ -50,7 +50,7 @@ pub fn max_pool2d(
             let ih_start = oh * sh;
             let iw_start = ow * sw;
 
-            let mut max_val = f32::NEG_INFINITY;
+            let mut max_val = E::neg_infinity();
 
             // Find max in kernel window
             for ki in 0..kh {
@@ -88,13 +88,13 @@ pub fn max_pool2d(
 ///
 /// Input: [batch, channels, height, width]
 /// Output: [batch, channels, out_height, out_width]
-pub fn avg_pool2d(
+pub fn avg_pool2d<E: CpuFloat>(
-    input: &CpuTensorPrimitive<4>,
+    input: &CpuTensorPrimitive<4, E>,
     kernel_size: [usize; 2],
     stride: [usize; 2],
     padding: [usize; 2],
     count_include_pad: bool,
-) -> CpuTensorPrimitive<4> {
+) -> CpuTensorPrimitive<4, E> {
     let [batch, channels, in_h, in_w] = input.shape;
     let [kh, kw] = kernel_size;
     let [sh, sw] = stride;
@@ -108,7 +108,7 @@ pub fn avg_pool2d(
     let output_size = batch * channels * out_h * out_w;
 
     // Process each output position in parallel
-    let output_data: Vec<f32> = (0..output_size)
+    let output_data: Vec<E> = (0..output_size)
         .into_par_iter()
         .map(|idx| {
             // Calculate position in output
@@ -121,7 +121,7 @@ pub fn avg_pool2d(
             let ih_start = oh * sh;
             let iw_start = ow * sw;
 
-            let mut sum = 0.0f32;
+            let mut sum = E::zero();
             let mut count = 0usize;
 
             // Sum values in kernel window
@@ -138,7 +138,7 @@ pub fn avg_pool2d(
                             + c * (in_h * in_w)
                             + actual_ih * in_w
                             + actual_iw;
-                        sum += input.data[input_idx];
+                        sum = sum + input.data[input_idx];
                         count += 1;
                     } else if count_include_pad {
                         // Padding contributes 0 to sum but counts toward divisor
@@ -151,9 +151,9 @@ pub fn avg_pool2d(
             let divisor = if count_include_pad { kh * kw } else { count };
 
             if divisor > 0 {
-                sum / divisor as f32
+                sum / E::from(divisor).unwrap()
             } else {
-                0.0
+                E::zero()
             }
         })
         .collect();
@@ -171,7 +171,11 @@ mod tests {
     use super::*;
     use crate::CpuDevice;
 
-    fn make_tensor(data: Vec<f32>, shape: [usize; 4], device: CpuDevice) -> CpuTensorPrimitive<4> {
+    fn make_tensor(
+        data: Vec<f32>,
+        shape: [usize; 4],
+        device: CpuDevice,
+    ) -> CpuTensorPrimitive<4, f32> {
         CpuTensorPrimitive {
             data,
             shape,

crates/core/rtx-backend-cpu/src/ops/reduction.rs

@@ -1,19 +1,21 @@
 //! Reduction operations.
 
-use crate::CpuTensorPrimitive;
+use crate::{CpuFloat, CpuTensorPrimitive};
 use rayon::prelude::*;
 
 /// Sum all elements.
-pub fn sum<const D: usize>(tensor: &CpuTensorPrimitive<D>) -> CpuTensorPrimitive<1> {
+pub fn sum<const D: usize, E: CpuFloat>(
-    let sum: f32 = tensor.data.par_iter().sum();
+    tensor: &CpuTensorPrimitive<D, E>,
+) -> CpuTensorPrimitive<1, E> {
+    let sum = tensor.data.par_iter().copied().reduce(E::zero, |a, b| a + b);
     CpuTensorPrimitive::new(vec![sum], [1], tensor.device.clone())
 }
 
 /// Sum along a dimension.
-pub fn sum_dim<const D: usize>(
+pub fn sum_dim<const D: usize, E: CpuFloat>(
-    tensor: &CpuTensorPrimitive<D>,
+    tensor: &CpuTensorPrimitive<D, E>,
     dim: usize,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
     assert!(dim < D, "Dimension out of range");
 
     let mut out_shape = tensor.shape;
@@ -24,14 +26,14 @@ pub fn sum_dim<const D: usize>(
     let outer_size: usize = tensor.shape[..dim].iter().product();
     let dim_size = tensor.shape[dim];
 
-    let mut result = vec![0.0f32; out_numel];
+    let mut result = vec![E::zero(); out_numel];
 
     for outer in 0..outer_size {
         for inner in 0..inner_size {
-            let mut sum = 0.0f32;
+            let mut sum = E::zero();
             for d in 0..dim_size {
                 let idx = outer * dim_size * inner_size + d * inner_size + inner;
-                sum += tensor.data[idx];
+                sum = sum + tensor.data[idx];
             }
             let out_idx = outer * inner_size + inner;
             result[out_idx] = sum;
@@ -42,68 +44,79 @@ pub fn sum_dim<const D: usize>(
 }
 
 /// Mean of all elements.
-pub fn mean<const D: usize>(tensor: &CpuTensorPrimitive<D>) -> CpuTensorPrimitive<1> {
+pub fn mean<const D: usize, E: CpuFloat>(
-    let sum: f32 = tensor.data.par_iter().sum();
+    tensor: &CpuTensorPrimitive<D, E>,
-    let mean = sum / tensor.numel() as f32;
+) -> CpuTensorPrimitive<1, E> {
+    let sum = tensor.data.par_iter().copied().reduce(E::zero, |a, b| a + b);
+    let mean = sum / E::from(tensor.numel()).unwrap();
     CpuTensorPrimitive::new(vec![mean], [1], tensor.device.clone())
 }
 
 /// Mean along a dimension.
-pub fn mean_dim<const D: usize>(
+pub fn mean_dim<const D: usize, E: CpuFloat>(
-    tensor: &CpuTensorPrimitive<D>,
+    tensor: &CpuTensorPrimitive<D, E>,
     dim: usize,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
     let sum_result = sum_dim(tensor, dim);
-    let dim_size = tensor.shape[dim] as f32;
+    let dim_size = E::from(tensor.shape[dim]).unwrap();
 
-    let result: Vec<f32> = sum_result.data.iter().map(|x| x / dim_size).collect();
+    let result: Vec<E> = sum_result.data.iter().map(|&x| x / dim_size).collect();
     CpuTensorPrimitive::new(result, sum_result.shape, tensor.device.clone())
 }
 
 /// Maximum element.
-pub fn max<const D: usize>(tensor: &CpuTensorPrimitive<D>) -> CpuTensorPrimitive<1> {
+pub fn max<const D: usize, E: CpuFloat>(
+    tensor: &CpuTensorPrimitive<D, E>,
+) -> CpuTensorPrimitive<1, E> {
     let max = tensor
         .data
         .par_iter()
-        .cloned()
+        .copied()
-        .reduce(|| f32::NEG_INFINITY, f32::max);
+        .reduce(E::neg_infinity, |a, b| a.max(b));
     CpuTensorPrimitive::new(vec![max], [1], tensor.device.clone())
 }
 
 /// Minimum element.
-pub fn min<const D: usize>(tensor: &CpuTensorPrimitive<D>) -> CpuTensorPrimitive<1> {
+pub fn min<const D: usize, E: CpuFloat>(
+    tensor: &CpuTensorPrimitive<D, E>,
+) -> CpuTensorPrimitive<1, E> {
     let min = tensor
         .data
         .par_iter()
-        .cloned()
+        .copied()
-        .reduce(|| f32::INFINITY, f32::min);
+        .reduce(E::infinity, |a, b| a.min(b));
     CpuTensorPrimitive::new(vec![min], [1], tensor.device.clone())
 }
 
 /// Variance of all elements.
-pub fn var<const D: usize>(tensor: &CpuTensorPrimitive<D>) -> CpuTensorPrimitive<1> {
+pub fn var<const D: usize, E: CpuFloat>(
-    let mean: f32 = tensor.data.par_iter().sum::<f32>() / tensor.numel() as f32;
+    tensor: &CpuTensorPrimitive<D, E>,
-    let variance: f32 = tensor
+) -> CpuTensorPrimitive<1, E> {
+    let n = E::from(tensor.numel()).unwrap();
+    let mean = tensor.data.par_iter().copied().reduce(E::zero, |a, b| a + b) / n;
+    let variance = tensor
         .data
         .par_iter()
         .map(|&x| (x - mean).powi(2))
-        .sum::<f32>()
+        .reduce(E::zero, |a, b| a + b)
-        / tensor.numel() as f32;
+        / n;
     CpuTensorPrimitive::new(vec![variance], [1], tensor.device.clone())
 }
 
 /// Standard deviation of all elements.
-pub fn std<const D: usize>(tensor: &CpuTensorPrimitive<D>) -> CpuTensorPrimitive<1> {
+pub fn std<const D: usize, E: CpuFloat>(
+    tensor: &CpuTensorPrimitive<D, E>,
+) -> CpuTensorPrimitive<1, E> {
     let var_result = var(tensor);
     let std = var_result.data[0].sqrt();
     CpuTensorPrimitive::new(vec![std], [1], tensor.device.clone())
 }
 
 /// Variance along a dimension.
-pub fn var_dim<const D: usize>(
+pub fn var_dim<const D: usize, E: CpuFloat>(
-    tensor: &CpuTensorPrimitive<D>,
+    tensor: &CpuTensorPrimitive<D, E>,
     dim: usize,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
     assert!(dim < D, "Dimension out of range");
 
     // First compute mean along dimension
@@ -116,22 +129,23 @@ pub fn var_dim<const D: usize>(
     let inner_size: usize = tensor.shape[dim + 1..].iter().product();
     let outer_size: usize = tensor.shape[..dim].iter().product();
     let dim_size = tensor.shape[dim];
+    let dim_size_e = E::from(dim_size).unwrap();
 
-    let mut result = vec![0.0f32; out_numel];
+    let mut result = vec![E::zero(); out_numel];
 
     for outer in 0..outer_size {
         for inner in 0..inner_size {
             let mean_idx = outer * inner_size + inner;
             let mean_val = mean_tensor.data[mean_idx];
 
-            let mut var_sum = 0.0f32;
+            let mut var_sum = E::zero();
             for d in 0..dim_size {
                 let idx = outer * dim_size * inner_size + d * inner_size + inner;
-                var_sum += (tensor.data[idx] - mean_val).powi(2);
+                var_sum = var_sum + (tensor.data[idx] - mean_val).powi(2);
             }
 
             let out_idx = outer * inner_size + inner;
-            result[out_idx] = var_sum / dim_size as f32;
+            result[out_idx] = var_sum / dim_size_e;
         }
     }
 

crates/core/rtx-backend-cpu/src/ops/shape.rs

@@ -1,12 +1,12 @@
 //! Shape manipulation operations.
 
-use crate::CpuTensorPrimitive;
+use crate::{CpuFloat, CpuTensorPrimitive};
 
 /// Reshape a tensor.
-pub fn reshape<const D1: usize, const D2: usize>(
+pub fn reshape<const D1: usize, const D2: usize, E: CpuFloat>(
-    tensor: CpuTensorPrimitive<D1>,
+    tensor: CpuTensorPrimitive<D1, E>,
     shape: [usize; D2],
-) -> CpuTensorPrimitive<D2> {
+) -> CpuTensorPrimitive<D2, E> {
     let old_numel: usize = tensor.shape.iter().product();
     let new_numel: usize = shape.iter().product();
     assert_eq!(old_numel, new_numel, "Total elements must remain the same");
@@ -15,7 +15,9 @@ pub fn reshape<const D1: usize, const D2: usize>(
 }
 
 /// Transpose the last two dimensions.
-pub fn transpose<const D: usize>(tensor: &CpuTensorPrimitive<D>) -> CpuTensorPrimitive<D> {
+pub fn transpose<const D: usize, E: CpuFloat>(
+    tensor: &CpuTensorPrimitive<D, E>,
+) -> CpuTensorPrimitive<D, E> {
     if D < 2 {
         return tensor.clone();
     }
@@ -23,11 +25,11 @@ pub fn transpose<const D: usize>(tensor: &CpuTensorPrimitive<D>) -> CpuTensorPri
 }
 
 /// Swap two dimensions.
-pub fn swap_dims<const D: usize>(
+pub fn swap_dims<const D: usize, E: CpuFloat>(
-    tensor: &CpuTensorPrimitive<D>,
+    tensor: &CpuTensorPrimitive<D, E>,
     dim1: usize,
     dim2: usize,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
     assert!(dim1 < D && dim2 < D, "Dimensions out of range");
 
     if dim1 == dim2 {
@@ -48,7 +50,7 @@ pub fn swap_dims<const D: usize>(
     }
 
     let numel = tensor.numel();
-    let mut result = vec![0.0f32; numel];
+    let mut result = vec![E::zero(); numel];
 
     // For each element in the source tensor
     for i in 0..numel {
@@ -77,16 +79,18 @@ pub fn swap_dims<const D: usize>(
 }
 
 /// Flatten tensor to 1D.
-pub fn flatten<const D: usize>(tensor: CpuTensorPrimitive<D>) -> CpuTensorPrimitive<1> {
+pub fn flatten<const D: usize, E: CpuFloat>(
+    tensor: CpuTensorPrimitive<D, E>,
+) -> CpuTensorPrimitive<1, E> {
     let numel = tensor.numel();
     CpuTensorPrimitive::new(tensor.data, [numel], tensor.device)
 }
 
 /// Squeeze dimension (remove dimension of size 1).
-pub fn squeeze<const D: usize>(
+pub fn squeeze<const D: usize, E: CpuFloat>(
-    tensor: &CpuTensorPrimitive<D>,
+    tensor: &CpuTensorPrimitive<D, E>,
     dim: usize,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
     assert!(dim < D, "Dimension out of range");
     assert_eq!(
         tensor.shape[dim], 1,
@@ -99,10 +103,10 @@ pub fn squeeze<const D: usize>(
 }
 
 /// Unsqueeze (add dimension of size 1).
-pub fn unsqueeze<const D: usize>(
+pub fn unsqueeze<const D: usize, E: CpuFloat>(
-    tensor: &CpuTensorPrimitive<D>,
+    tensor: &CpuTensorPrimitive<D, E>,
     _dim: usize,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
     // Note: In a real implementation, this would return a tensor with more dimensions
     // For now, we just return a clone since const generics prevent changing D
     tensor.clone()

crates/core/rtx-backend-cpu/src/ops/unary.rs

@@ -1,62 +1,85 @@
 //! Unary operations.
 
-use crate::CpuTensorPrimitive;
+use crate::{CpuFloat, CpuTensorPrimitive};
 use rayon::prelude::*;
 
 /// Element-wise negation.
-pub fn neg<const D: usize>(tensor: &CpuTensorPrimitive<D>) -> CpuTensorPrimitive<D> {
+pub fn neg<const D: usize, E: CpuFloat>(
-    let result: Vec<f32> = tensor.data.par_iter().map(|&x| -x).collect();
+    tensor: &CpuTensorPrimitive<D, E>,
+) -> CpuTensorPrimitive<D, E> {
+    let result: Vec<E> = tensor.data.par_iter().map(|&x| -x).collect();
     CpuTensorPrimitive::new(result, tensor.shape, tensor.device.clone())
 }
 
 /// Element-wise exponential.
-pub fn exp<const D: usize>(tensor: &CpuTensorPrimitive<D>) -> CpuTensorPrimitive<D> {
+pub fn exp<const D: usize, E: CpuFloat>(
-    let result: Vec<f32> = tensor.data.par_iter().map(|&x| x.exp()).collect();
+    tensor: &CpuTensorPrimitive<D, E>,
+) -> CpuTensorPrimitive<D, E> {
+    let result: Vec<E> = tensor.data.par_iter().map(|&x| x.exp()).collect();
     CpuTensorPrimitive::new(result, tensor.shape, tensor.device.clone())
 }
 
 /// Element-wise natural logarithm.
-pub fn log<const D: usize>(tensor: &CpuTensorPrimitive<D>) -> CpuTensorPrimitive<D> {
+pub fn log<const D: usize, E: CpuFloat>(
-    let result: Vec<f32> = tensor.data.par_iter().map(|&x| x.ln()).collect();
+    tensor: &CpuTensorPrimitive<D, E>,
+) -> CpuTensorPrimitive<D, E> {
+    let result: Vec<E> = tensor.data.par_iter().map(|&x| x.ln()).collect();
     CpuTensorPrimitive::new(result, tensor.shape, tensor.device.clone())
 }
 
 /// Element-wise square root.
-pub fn sqrt<const D: usize>(tensor: &CpuTensorPrimitive<D>) -> CpuTensorPrimitive<D> {
+pub fn sqrt<const D: usize, E: CpuFloat>(
-    let result: Vec<f32> = tensor.data.par_iter().map(|&x| x.sqrt()).collect();
+    tensor: &CpuTensorPrimitive<D, E>,
+) -> CpuTensorPrimitive<D, E> {
+    let result: Vec<E> = tensor.data.par_iter().map(|&x| x.sqrt()).collect();
     CpuTensorPrimitive::new(result, tensor.shape, tensor.device.clone())
 }
 
 /// Element-wise absolute value.
-pub fn abs<const D: usize>(tensor: &CpuTensorPrimitive<D>) -> CpuTensorPrimitive<D> {
+pub fn abs<const D: usize, E: CpuFloat>(
-    let result: Vec<f32> = tensor.data.par_iter().map(|&x| x.abs()).collect();
+    tensor: &CpuTensorPrimitive<D, E>,
+) -> CpuTensorPrimitive<D, E> {
+    let result: Vec<E> = tensor.data.par_iter().map(|&x| x.abs()).collect();
     CpuTensorPrimitive::new(result, tensor.shape, tensor.device.clone())
 }
 
 /// Element-wise sine.
-pub fn sin<const D: usize>(tensor: &CpuTensorPrimitive<D>) -> CpuTensorPrimitive<D> {
+pub fn sin<const D: usize, E: CpuFloat>(
-    let result: Vec<f32> = tensor.data.par_iter().map(|&x| x.sin()).collect();
+    tensor: &CpuTensorPrimitive<D, E>,
+) -> CpuTensorPrimitive<D, E> {
+    let result: Vec<E> = tensor.data.par_iter().map(|&x| x.sin()).collect();
     CpuTensorPrimitive::new(result, tensor.shape, tensor.device.clone())
 }
 
 /// Element-wise cosine.
-pub fn cos<const D: usize>(tensor: &CpuTensorPrimitive<D>) -> CpuTensorPrimitive<D> {
+pub fn cos<const D: usize, E: CpuFloat>(
-    let result: Vec<f32> = tensor.data.par_iter().map(|&x| x.cos()).collect();
+    tensor: &CpuTensorPrimitive<D, E>,
+) -> CpuTensorPrimitive<D, E> {
+    let result: Vec<E> = tensor.data.par_iter().map(|&x| x.cos()).collect();
     CpuTensorPrimitive::new(result, tensor.shape, tensor.device.clone())
 }
 
 /// Element-wise power.
-pub fn pow<const D: usize>(tensor: &CpuTensorPrimitive<D>, exp: f32) -> CpuTensorPrimitive<D> {
+pub fn pow<const D: usize, E: CpuFloat>(
-    let result: Vec<f32> = tensor.data.par_iter().map(|&x| x.powf(exp)).collect();
+    tensor: &CpuTensorPrimitive<D, E>,
+    exp: E,
+) -> CpuTensorPrimitive<D, E> {
+    let result: Vec<E> = tensor.data.par_iter().map(|&x| x.powf(exp)).collect();
     CpuTensorPrimitive::new(result, tensor.shape, tensor.device.clone())
 }
 
 /// Clamp values to a range.
-pub fn clamp<const D: usize>(
+pub fn clamp<const D: usize, E: CpuFloat>(
-    tensor: &CpuTensorPrimitive<D>,
+    tensor: &CpuTensorPrimitive<D, E>,
-    min: f32,
+    min: E,
-    max: f32,
+    max: E,
-) -> CpuTensorPrimitive<D> {
+) -> CpuTensorPrimitive<D, E> {
-    let result: Vec<f32> = tensor.data.par_iter().map(|&x| x.clamp(min, max)).collect();
+    // num_traits::Float has no `clamp`; compose max/min (matches f32::clamp for
+    // non-NaN bounds, which is the invariant callers rely on).
+    let result: Vec<E> = tensor
+        .data
+        .par_iter()
+        .map(|&x| x.max(min).min(max))
+        .collect();
     CpuTensorPrimitive::new(result, tensor.shape, tensor.device.clone())
 }