rustyverse/rustytorch
2
0
Fork 0
Code Issues Pull Requests Actions 46 Packages Projects Releases Wiki Activity
Files
af19860c04e064dfdf86cfe92c82b96c61a2f341
rustytorch/crates/core/rtx-backend-cpu/src/lib.rs
T
Omar SobhandClaude Fable 5 9969d8a661 feat(backend): add differentiable index_select / index_add row ops 
Add two row-indexing ops along dim 0 to the `Backend` trait so gather /
scatter-add message passing (GNNs, segment softmax, bias tiling) can be
trained through `Autodiff<B>`:

- `index_select(tensor, indices)` — out[i, ..] = tensor[indices[i], ..]
- `index_add(tensor, indices, num_rows)` — out = zeros; out[idx[i], ..] += tensor[i, ..]

They are each other's adjoint, which is what the backward passes use.

Both trait methods have default bodies (host round-trip via to_data /
from_data) so every existing backend keeps compiling and is correct;
backends override with native kernels:

- rtx-backend-cpu: new ops/index.rs (rayon-parallel gather over output
  rows above a size threshold, sequential deterministic scatter-add),
  wired into CpuBackend and CpuBackendF64, with unit tests for D=1/2/3,
  duplicates, untouched rows, empty inputs, bounds panics and adjointness.
- rtx-autograd: Autodiff<B> overrides both ops and records
  IndexSelectBackward / IndexAddBackward (new ops/index.rs); finite-
  difference gradchecks on the real CpuBackend cover repeated-index
  accumulation, untouched-row zero grads, bias tiling via index_select
  of a [1,F] row, and a full per-segment softmax.
- rtx-fusion: forward both ops to the inner backend.

Co-Authored-By: Claude Fable 5 <[email protected]>
2026-08-20 11:05:42 -07:00

863 lines
26 KiB
Rust
Raw Blame History

//! # RustyTorch++ CPU Backend
//!
//! Pure Rust CPU backend implementation for portability and fallback.
//!
//! ## Features
//!
//! - **Zero Dependencies**: No GPU drivers or special hardware required
//! - **Parallel Execution**: Uses Rayon for multi-threaded operations
//! - **SIMD Ready**: Optional SIMD optimizations with feature flag
//! - **Reference Implementation**: Useful for testing and validation
//!
//! ## Architecture
//!
//! ```text
//! CpuBackend
//! ├── CpuTensorPrimitive - Vec<f32> storage
//! ├── CpuDevice - CPU device abstraction
//! └── Ops
//! ├── Basic - Add, mul, etc. (parallel)
//! ├── GEMM - Matrix multiply (cache-blocked)
//! └── Attention - Reference implementation
//! ```
//!
//! ## Example
//!
//! ```rust,ignore
//! use rtx_backend_cpu::{CpuBackend, CpuDevice};
//! use rtx_backend::Backend;
//!
//! let device = CpuDevice::default();
//! let a = CpuBackend::zeros([1024, 1024], &device);
//! let b = CpuBackend::randn([1024, 1024], &device);
//! let c = CpuBackend::matmul(&a, &b);
//! ```
#![warn(missing_docs)]
mod device;
mod error;
mod ops;
pub mod simd;
mod tensor;
pub use device::CpuDevice;
pub use error::{CpuBackendError, CpuBackendResult};
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:
/// - Systems without GPU support
/// - Testing and validation against GPU implementations
/// - Development and debugging
/// - Small models where GPU overhead isn't justified
#[derive(Clone, Debug, Default)]
pub struct CpuBackend;
impl Backend for CpuBackend {
type TensorPrimitive<const D: usize> = CpuTensorPrimitive<D>;
type Device = CpuDevice;
type FloatElem = f32;
type IntElem = i32;
type BoolElem = BoolU8;
fn name() -> &'static str {
"cpu"
}
fn seed(seed: u64) {
ops::seed_rng(seed);
}
// ==================== Tensor Creation ====================
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)
}
// ==================== Basic Operations ====================
// NOTE: Operations take owned tensors for ownership-based fusion.
// Internally we borrow for the actual computation.
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)
}
// ==================== Unary Operations ====================
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)
}
// ==================== Activation Functions ====================
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)
}
// ==================== Reduction Operations ====================
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)
}
// ==================== Shape Operations ====================
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)
}
// ==================== Row Indexing ====================
fn index_select<const D: usize>(
tensor: Self::TensorPrimitive<D>,
indices: &[usize],
) -> Self::TensorPrimitive<D> {
ops::index::index_select(&tensor, indices)
}
fn index_add<const D: usize>(
tensor: Self::TensorPrimitive<D>,
indices: &[usize],
num_rows: usize,
) -> Self::TensorPrimitive<D> {
ops::index::index_add(&tensor, indices, num_rows)
}
// ==================== LLM-Specific Operations ====================
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)
}
// ==================== Comparison Operations ====================
fn gt_scalar<const D: usize>(
tensor: Self::TensorPrimitive<D>,
value: Self::FloatElem,
) -> Self::TensorPrimitive<D> {
ops::basic::gt_scalar(&tensor, value)
}
// ==================== Convolution Operations ====================
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)
}
// ==================== Pooling Operations ====================
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)
}
// ==================== Device Management ====================
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> {
// CPU backend only has one device, so this is a no-op
tensor
}
fn to_data<const D: usize>(tensor: &Self::TensorPrimitive<D>) -> Vec<Self::FloatElem> {
tensor.data.clone()
}
fn sync(_device: &Self::Device) {
// CPU operations are synchronous, nothing to do
}
}
/// 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)
}
// ==================== Row Indexing ====================
fn index_select<const D: usize>(
tensor: Self::TensorPrimitive<D>,
indices: &[usize],
) -> Self::TensorPrimitive<D> {
ops::index::index_select(&tensor, indices)
}
fn index_add<const D: usize>(
tensor: Self::TensorPrimitive<D>,
indices: &[usize],
num_rows: usize,
) -> Self::TensorPrimitive<D> {
ops::index::index_add(&tensor, indices, num_rows)
}
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 backend_trait_index_ops_f32_and_f64() {
let dev = CpuDevice::new();
let x = CpuBackend::from_data(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0], [3, 2], &dev);
let g = <CpuBackend as Backend>::index_select(x, &[2, 2, 0]);
assert_eq!(CpuBackend::shape(&g), [3, 2]);
assert_eq!(g.to_vec(), vec![5.0, 6.0, 5.0, 6.0, 1.0, 2.0]);
let s = <CpuBackend as Backend>::index_add(g, &[1, 1, 3], 4);
assert_eq!(CpuBackend::shape(&s), [4, 2]);
assert_eq!(s.to_vec(), vec![0.0, 0.0, 10.0, 12.0, 0.0, 0.0, 1.0, 2.0]);
let x = CpuBackendF64::from_data(&[1.0_f64, 2.0, 3.0, 4.0], [2, 2], &dev);
let g = <CpuBackendF64 as Backend>::index_select(x, &[1, 0, 1]);
assert_eq!(g.to_vec(), vec![3.0, 4.0, 1.0, 2.0, 3.0, 4.0]);
let s = <CpuBackendF64 as Backend>::index_add(g, &[0, 0, 2], 3);
assert_eq!(s.to_vec(), vec![4.0, 6.0, 0.0, 0.0, 3.0, 4.0]);
}
#[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}");
}
}
Reference in New Issue View Git Blame Copy Permalink