Trait ComputeBackend 

pub trait ComputeBackend: Send + Sync {
    type Tensor: TensorOps;
    type Device: DeviceOps;

    // Required methods
    fn device(&self) -> &Self::Device;
    fn allocate(&self, shape: &[usize], dtype: DType) -> Result<Self::Tensor>;
    fn from_slice(&self, data: &[f32], shape: &[usize]) -> Result<Self::Tensor>;
    fn matmul(&self, a: &Self::Tensor, b: &Self::Tensor) -> Result<Self::Tensor>;
    fn batch_matmul(
        &self,
        a: &Self::Tensor,
        b: &Self::Tensor,
    ) -> Result<Self::Tensor>;
    fn attention(
        &self,
        q: &Self::Tensor,
        k: &Self::Tensor,
        v: &Self::Tensor,
        mask: Option<&Self::Tensor>,
        scale: Option<f32>,
    ) -> Result<Self::Tensor>;
    fn rms_norm(
        &self,
        x: &Self::Tensor,
        weight: &Self::Tensor,
        eps: f32,
    ) -> Result<Self::Tensor>;
    fn layer_norm(
        &self,
        x: &Self::Tensor,
        weight: &Self::Tensor,
        bias: Option<&Self::Tensor>,
        eps: f32,
    ) -> Result<Self::Tensor>;
    fn silu(&self, x: &Self::Tensor) -> Result<Self::Tensor>;
    fn gelu(&self, x: &Self::Tensor) -> Result<Self::Tensor>;
    fn relu(&self, x: &Self::Tensor) -> Result<Self::Tensor>;
    fn softmax(&self, x: &Self::Tensor, dim: i32) -> Result<Self::Tensor>;
    fn add(&self, a: &Self::Tensor, b: &Self::Tensor) -> Result<Self::Tensor>;
    fn mul(&self, a: &Self::Tensor, b: &Self::Tensor) -> Result<Self::Tensor>;
    fn transpose(&self, x: &Self::Tensor) -> Result<Self::Tensor>;
    fn reshape(&self, x: &Self::Tensor, shape: &[usize]) -> Result<Self::Tensor>;
    fn to_device(&self, tensor: &Self::Tensor) -> Result<Self::Tensor>;
    fn to_cpu(&self, tensor: &Self::Tensor) -> Result<Vec<f32>>;
}

Trait defining a compute backend.

Required Associated Types

type Tensor: TensorOps

The tensor type for this backend.

type Device: DeviceOps

The device type for this backend.

Required Methods

fn device(&self) -> &Self::Device

Returns the device.

fn allocate(&self, shape: &[usize], dtype: DType) -> Result<Self::Tensor>

Allocates a new tensor filled with zeros.

fn from_slice(&self, data: &[f32], shape: &[usize]) -> Result<Self::Tensor>

Creates a tensor from raw data.

fn matmul(&self, a: &Self::Tensor, b: &Self::Tensor) -> Result<Self::Tensor>

Performs matrix multiplication: C = A @ B.

fn batch_matmul( &self, a: &Self::Tensor, b: &Self::Tensor, ) -> Result<Self::Tensor>

Performs batched matrix multiplication.

fn attention( &self, q: &Self::Tensor, k: &Self::Tensor, v: &Self::Tensor, mask: Option<&Self::Tensor>, scale: Option<f32>, ) -> Result<Self::Tensor>

Performs scaled dot-product attention. Returns softmax(Q @ K^T / sqrt(d_k)) @ V with optional causal masking.

fn rms_norm( &self, x: &Self::Tensor, weight: &Self::Tensor, eps: f32, ) -> Result<Self::Tensor>

Applies RMS normalization: x * weight / sqrt(mean(x^2) + eps).

fn layer_norm( &self, x: &Self::Tensor, weight: &Self::Tensor, bias: Option<&Self::Tensor>, eps: f32, ) -> Result<Self::Tensor>

Applies Layer normalization.

fn silu(&self, x: &Self::Tensor) -> Result<Self::Tensor>

Applies SiLU (Swish) activation: x * sigmoid(x).

fn gelu(&self, x: &Self::Tensor) -> Result<Self::Tensor>

Applies GELU activation.

fn relu(&self, x: &Self::Tensor) -> Result<Self::Tensor>

Applies ReLU activation.

fn softmax(&self, x: &Self::Tensor, dim: i32) -> Result<Self::Tensor>

Applies softmax along the specified dimension.

fn add(&self, a: &Self::Tensor, b: &Self::Tensor) -> Result<Self::Tensor>

Element-wise addition.

fn mul(&self, a: &Self::Tensor, b: &Self::Tensor) -> Result<Self::Tensor>

Element-wise multiplication.

fn transpose(&self, x: &Self::Tensor) -> Result<Self::Tensor>

Transposes the last two dimensions.

fn reshape(&self, x: &Self::Tensor, shape: &[usize]) -> Result<Self::Tensor>

Reshapes a tensor.

fn to_device(&self, tensor: &Self::Tensor) -> Result<Self::Tensor>

Copies tensor to device.

fn to_cpu(&self, tensor: &Self::Tensor) -> Result<Vec<f32>>

Copies tensor data to CPU as f32 vector.

Dyn Compatibility

This trait is dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors

impl ComputeBackend for CpuBackend

type Tensor = CpuTensor

type Device = CpuDevice