use crate::inference_context::InferenceContext;
use crate::{DType, Device, Shape};
use fuel_ir::shape::{Dim, Dims};
use std::sync::Arc;
#[derive(Clone, Debug)]
pub struct Tensor {
pub(crate) inner: fuel_graph::NodeHandle,
}
impl Tensor {
pub fn from_f32(
data: impl Into<Arc<[f32]>>,
shape: impl Into<Shape>,
device: &crate::Device,
) -> fuel_ir::Result<Self> {
Ok(Self {
inner: fuel_graph::NodeHandle::from_f32(data, shape, device.as_dyn())?,
})
}
pub fn graph_id(&self) -> fuel_graph::GraphId {
self.inner.graph_id()
}
pub fn graph(&self) -> &fuel_graph::SharedGraph {
self.inner.graph()
}
pub fn from_f32_on(
graph: &fuel_graph::SharedGraph,
data: impl Into<Arc<[f32]>>,
shape: impl Into<Shape>,
device: &crate::Device,
) -> fuel_ir::Result<Self> {
Ok(Self {
inner: fuel_graph::NodeHandle::from_f32_on(graph, data, shape, device.as_dyn())?,
})
}
pub fn from_f64_on(
graph: &fuel_graph::SharedGraph,
data: impl Into<Arc<[f64]>>,
shape: impl Into<Shape>,
device: &crate::Device,
) -> fuel_ir::Result<Self> {
Ok(Self {
inner: fuel_graph::NodeHandle::from_f64_on(graph, data, shape, device.as_dyn())?,
})
}
pub fn from_bf16_on(
graph: &fuel_graph::SharedGraph,
data: impl Into<Arc<[half::bf16]>>,
shape: impl Into<Shape>,
device: &crate::Device,
) -> fuel_ir::Result<Self> {
Ok(Self {
inner: fuel_graph::NodeHandle::from_bf16_on(graph, data, shape, device.as_dyn())?,
})
}
pub fn from_f16_on(
graph: &fuel_graph::SharedGraph,
data: impl Into<Arc<[half::f16]>>,
shape: impl Into<Shape>,
device: &crate::Device,
) -> fuel_ir::Result<Self> {
Ok(Self {
inner: fuel_graph::NodeHandle::from_f16_on(graph, data, shape, device.as_dyn())?,
})
}
pub fn from_u32_on(
graph: &fuel_graph::SharedGraph,
data: impl Into<Arc<[u32]>>,
shape: impl Into<Shape>,
device: &crate::Device,
) -> fuel_ir::Result<Self> {
Ok(Self {
inner: fuel_graph::NodeHandle::from_u32_on(graph, data, shape, device.as_dyn())?,
})
}
pub fn from_f64(
data: impl Into<Arc<[f64]>>,
shape: impl Into<Shape>,
device: &crate::Device,
) -> fuel_ir::Result<Self> {
Ok(Self {
inner: fuel_graph::NodeHandle::from_f64(data, shape, device.as_dyn())?,
})
}
pub fn from_bf16(
data: impl Into<Arc<[half::bf16]>>,
shape: impl Into<Shape>,
device: &crate::Device,
) -> fuel_ir::Result<Self> {
Ok(Self {
inner: fuel_graph::NodeHandle::from_bf16(data, shape, device.as_dyn())?,
})
}
pub fn from_f16(
data: impl Into<Arc<[half::f16]>>,
shape: impl Into<Shape>,
device: &crate::Device,
) -> fuel_ir::Result<Self> {
Ok(Self {
inner: fuel_graph::NodeHandle::from_f16(data, shape, device.as_dyn())?,
})
}
pub fn from_u32(
data: impl Into<Arc<[u32]>>,
shape: impl Into<Shape>,
device: &crate::Device,
) -> fuel_ir::Result<Self> {
Ok(Self {
inner: fuel_graph::NodeHandle::from_u32(data, shape, device.as_dyn())?,
})
}
pub fn const_f32_like(
&self,
data: impl Into<Arc<[f32]>>,
shape: impl Into<Shape>,
) -> std::result::Result<Self, fuel_ir::Error> {
Ok(Self {
inner: self.inner.const_f32_like(data, shape)?,
})
}
pub fn const_f16_like(
&self,
data: impl Into<Arc<[half::f16]>>,
shape: impl Into<Shape>,
) -> std::result::Result<Self, fuel_ir::Error> {
Ok(Self {
inner: self.inner.const_f16_like(data, shape)?,
})
}
pub fn const_bf16_like(
&self,
data: impl Into<Arc<[half::bf16]>>,
shape: impl Into<Shape>,
) -> std::result::Result<Self, fuel_ir::Error> {
Ok(Self {
inner: self.inner.const_bf16_like(data, shape)?,
})
}
pub fn const_like_dtype(
&self,
data: &[f32],
shape: impl Into<Shape>,
dtype: fuel_ir::DType,
) -> std::result::Result<Self, fuel_ir::Error> {
match dtype {
fuel_ir::DType::F32 => self.const_f32_like(data.to_vec(), shape),
fuel_ir::DType::BF16 => {
let converted: Vec<half::bf16> =
data.iter().map(|&v| half::bf16::from_f32(v)).collect();
self.const_bf16_like(converted, shape)
}
other => Err(fuel_ir::Error::Msg(format!(
"const_like_dtype: unsupported activation dtype {other:?} \
(expected F32 or BF16)",
))
.bt()),
}
}
pub fn into_graph_tensor(self) -> fuel_graph::NodeHandle {
self.inner
}
pub fn graph_tensor(&self) -> &fuel_graph::NodeHandle {
&self.inner
}
pub fn from_graph_tensor(t: fuel_graph::NodeHandle) -> Self {
Self { inner: t }
}
pub fn shape(&self) -> Shape {
self.inner.shape()
}
pub fn dtype(&self) -> DType {
self.inner.dtype()
}
pub fn graph_handle(&self) -> &fuel_graph::SharedGraph {
self.inner.graph()
}
pub fn node_id(&self) -> fuel_graph::NodeId {
self.inner.id()
}
pub fn rank(&self) -> usize {
self.inner.shape().dims().len()
}
pub fn elem_count(&self) -> usize {
self.inner.shape().elem_count()
}
pub fn numel(&self) -> usize {
self.elem_count()
}
pub fn dim<D: Dim>(&self, dim: D) -> std::result::Result<usize, fuel_ir::Error> {
let shape = self.inner.shape();
let dim = dim.to_index(&shape, "dim")?;
Ok(shape.dims()[dim])
}
pub fn add(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.check_strict_binary("add", other)?;
Ok(Self {
inner: self.inner.add(&other.inner),
})
}
pub fn sub(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.check_strict_binary("sub", other)?;
Ok(Self {
inner: self.inner.sub(&other.inner),
})
}
pub fn mul(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.check_strict_binary("mul", other)?;
Ok(Self {
inner: self.inner.mul(&other.inner),
})
}
pub fn div(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.check_strict_binary("div", other)?;
Ok(Self {
inner: self.inner.div(&other.inner),
})
}
pub fn maximum(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.check_strict_binary("maximum", other)?;
Ok(Self {
inner: self.inner.maximum(&other.inner),
})
}
pub fn minimum(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.check_strict_binary("minimum", other)?;
Ok(Self {
inner: self.inner.minimum(&other.inner),
})
}
pub fn fmax_ieee(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.check_strict_binary("fmax_ieee", other)?;
Ok(Self {
inner: self.inner.fmax_ieee(&other.inner),
})
}
pub fn fmin_ieee(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.check_strict_binary("fmin_ieee", other)?;
Ok(Self {
inner: self.inner.fmin_ieee(&other.inner),
})
}
pub fn rem_trunc(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.check_strict_binary("rem_trunc", other)?;
Ok(Self {
inner: self.inner.rem_trunc(&other.inner),
})
}
pub fn trunc(&self) -> Self {
Self {
inner: self.inner.trunc(),
}
}
pub fn eq(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.check_strict_binary("eq", other)?;
Ok(Self {
inner: self.inner.eq(&other.inner),
})
}
pub fn ne(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.check_strict_binary("ne", other)?;
Ok(Self {
inner: self.inner.ne(&other.inner),
})
}
pub fn lt(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.check_strict_binary("lt", other)?;
Ok(Self {
inner: self.inner.lt(&other.inner),
})
}
pub fn le(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.check_strict_binary("le", other)?;
Ok(Self {
inner: self.inner.le(&other.inner),
})
}
pub fn gt(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.check_strict_binary("gt", other)?;
Ok(Self {
inner: self.inner.gt(&other.inner),
})
}
pub fn ge(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.check_strict_binary("ge", other)?;
Ok(Self {
inner: self.inner.ge(&other.inner),
})
}
pub fn where_cond(&self, a: &Self, b: &Self) -> std::result::Result<Self, fuel_ir::Error> {
if self.inner.dtype() != fuel_ir::DType::Bool {
return Err(fuel_ir::Error::Msg(format!(
"where_cond: cond mask must be Bool, got {:?} — cast a numeric \
mask to Bool explicitly (GAP-168(c))",
self.inner.dtype(),
))
.bt());
}
if a.inner.dtype() != b.inner.dtype() {
return Err(fuel_ir::Error::Msg(format!(
"where_cond: branches must share dtype, got a={:?} b={:?}",
a.inner.dtype(),
b.inner.dtype(),
))
.bt());
}
let cond_dims = self.inner.shape();
let a_dims = a.inner.shape();
let b_dims = b.inner.shape();
if a_dims.dims() != cond_dims.dims() || b_dims.dims() != cond_dims.dims() {
return Err(fuel_ir::Error::Msg(format!(
"where_cond: shapes must match cond, got cond={:?} a={:?} b={:?}",
cond_dims.dims(),
a_dims.dims(),
b_dims.dims(),
))
.bt());
}
Ok(Self {
inner: self.inner.where_cond(&a.inner, &b.inner),
})
}
pub fn broadcast_add(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.check_broadcast_binary("broadcast_add", other)?;
Ok(Self {
inner: self.inner.broadcast_add(&other.inner),
})
}
pub fn broadcast_sub(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.check_broadcast_binary("broadcast_sub", other)?;
Ok(Self {
inner: self.inner.broadcast_sub(&other.inner),
})
}
pub fn broadcast_mul(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.check_broadcast_binary("broadcast_mul", other)?;
Ok(Self {
inner: self.inner.broadcast_mul(&other.inner),
})
}
pub fn broadcast_div(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.check_broadcast_binary("broadcast_div", other)?;
Ok(Self {
inner: self.inner.broadcast_div(&other.inner),
})
}
fn check_strict_binary(
&self,
name: &'static str,
other: &Self,
) -> std::result::Result<(), fuel_ir::Error> {
if self.inner.dtype() != other.inner.dtype() {
return Err(fuel_ir::Error::Msg(format!(
"{name}: dtype mismatch lhs={:?} rhs={:?}",
self.inner.dtype(),
other.inner.dtype(),
))
.bt());
}
let a_shape = self.inner.shape();
let b_shape = other.inner.shape();
if a_shape.dims() != b_shape.dims() {
return Err(fuel_ir::Error::Msg(format!(
"{name}: shape mismatch lhs={:?} rhs={:?}",
a_shape.dims(),
b_shape.dims(),
))
.bt());
}
Ok(())
}
fn check_broadcast_binary(
&self,
name: &'static str,
other: &Self,
) -> std::result::Result<(), fuel_ir::Error> {
if self.inner.dtype() != other.inner.dtype() {
return Err(fuel_ir::Error::Msg(format!(
"{name}: dtype mismatch lhs={:?} rhs={:?}",
self.inner.dtype(),
other.inner.dtype(),
))
.bt());
}
let a_shape = self.inner.shape();
let b_shape = other.inner.shape();
let a_dims = a_shape.dims();
let b_dims = b_shape.dims();
let rank = a_dims.len().max(b_dims.len());
for i in 0..rank {
let ad = a_dims
.get(a_dims.len().wrapping_sub(1 + i))
.copied()
.unwrap_or(1);
let bd = b_dims
.get(b_dims.len().wrapping_sub(1 + i))
.copied()
.unwrap_or(1);
if ad != bd && ad != 1 && bd != 1 {
return Err(fuel_ir::Error::Msg(format!(
"{name}: shapes {:?} and {:?} are not broadcast-compatible",
a_dims, b_dims,
))
.bt());
}
}
Ok(())
}
pub fn neg(&self) -> Self {
Self {
inner: self.inner.neg(),
}
}
pub fn sqr(&self) -> Self {
Self {
inner: self.inner.sqr(),
}
}
pub fn sqrt(&self) -> Self {
Self {
inner: self.inner.sqrt(),
}
}
pub fn exp(&self) -> Self {
Self {
inner: self.inner.exp(),
}
}
pub fn log(&self) -> Self {
Self {
inner: self.inner.log(),
}
}
pub fn relu(&self) -> Self {
Self {
inner: self.inner.relu(),
}
}
pub fn silu(&self) -> Self {
Self {
inner: self.inner.silu(),
}
}
pub fn gelu_tanh(&self) -> Self {
Self {
inner: self.inner.gelu_tanh(),
}
}
pub fn sigmoid(&self) -> Self {
Self {
inner: self.inner.sigmoid(),
}
}
pub fn tanh(&self) -> Self {
Self {
inner: self.inner.tanh(),
}
}
pub fn sin(&self) -> Self {
Self {
inner: self.inner.sin(),
}
}
pub fn cos(&self) -> Self {
Self {
inner: self.inner.cos(),
}
}
pub fn step(&self) -> Self {
Self {
inner: self.inner.step(),
}
}
pub fn recip(&self) -> Self {
Self {
inner: self.inner.recip(),
}
}
pub fn abs(&self) -> Self {
Self {
inner: self.inner.abs(),
}
}
pub fn floor(&self) -> Self {
Self {
inner: self.inner.floor(),
}
}
pub fn ceil(&self) -> Self {
Self {
inner: self.inner.ceil(),
}
}
pub fn round(&self) -> Self {
Self {
inner: self.inner.round(),
}
}
pub fn sign(&self) -> Self {
Self {
inner: self.inner.sign(),
}
}
pub fn erf(&self) -> Self {
Self {
inner: self.inner.erf(),
}
}
pub fn gelu(&self) -> Self {
Self {
inner: self.inner.gelu(),
}
}
pub fn pow(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
Ok(Self {
inner: self.inner.pow(&other.inner)?,
})
}
pub fn rsqrt(&self) -> Self {
Self {
inner: self.inner.rsqrt(),
}
}
pub fn rem(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
Ok(Self {
inner: self.inner.rem(&other.inner)?,
})
}
pub fn flip<D: Dim>(&self, dim: D) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dim = dim.to_index(&shape, "flip")?;
Ok(Self {
inner: self.inner.flip(dim)?,
})
}
pub fn roll<D: Dim>(&self, dim: D, shift: i64) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dim = dim.to_index(&shape, "roll")?;
Ok(Self {
inner: self.inner.roll(dim, shift)?,
})
}
pub fn cumsum<D: Dim>(&self, dim: D) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dim = dim.to_index(&shape, "cumsum")?;
Ok(Self {
inner: self.inner.cumsum(dim)?,
})
}
pub fn pad(
&self,
padding: Vec<(usize, usize)>,
mode: fuel_graph::PadMode,
value: f64,
) -> std::result::Result<Self, fuel_ir::Error> {
Ok(Self {
inner: self.inner.pad(padding, mode, value)?,
})
}
pub fn powi(&self, n: i32) -> Self {
Self {
inner: self.inner.powi(n),
}
}
pub fn matmul(&self, other: &Self) -> std::result::Result<Self, fuel_ir::Error> {
let a_dims = self.inner.shape().dims().to_vec();
let b_dims = other.inner.shape().dims().to_vec();
if a_dims.len() < 2 || b_dims.len() < 2 {
return Err(fuel_ir::Error::Msg(format!(
"matmul: both operands must be rank >= 2, got lhs={a_dims:?} rhs={b_dims:?}",
))
.bt());
}
let a_k = a_dims[a_dims.len() - 1];
let b_k = b_dims[b_dims.len() - 2];
if a_k != b_k {
return Err(fuel_ir::Error::Msg(format!(
"matmul: contracting dim mismatch lhs[..., M, {a_k}] vs rhs[..., {b_k}, N]",
))
.bt());
}
Ok(Self {
inner: self.inner.matmul(&other.inner),
})
}
pub fn matmul_dyn_m(
&self,
other: &Self,
row_count: fuel_ir::DynScalar,
) -> std::result::Result<Self, fuel_ir::Error> {
let a_dims = self.inner.shape().dims().to_vec();
let b_dims = other.inner.shape().dims().to_vec();
if a_dims.len() < 2 || b_dims.len() < 2 {
return Err(fuel_ir::Error::Msg(format!(
"matmul_dyn_m: both operands must be rank >= 2, got lhs={a_dims:?} rhs={b_dims:?}",
))
.bt());
}
let a_k = a_dims[a_dims.len() - 1];
let b_k = b_dims[b_dims.len() - 2];
if a_k != b_k {
return Err(fuel_ir::Error::Msg(format!(
"matmul_dyn_m: contracting dim mismatch lhs[..., M, {a_k}] vs rhs[..., {b_k}, N]",
))
.bt());
}
Ok(Self {
inner: self.inner.matmul_dyn_m(&other.inner, row_count),
})
}
pub fn qmatmul(
&self,
weight_bytes: &Self,
quant_type: fuel_graph::QuantType,
k: usize,
n: usize,
) -> std::result::Result<Self, fuel_ir::Error> {
if self.inner.dtype() != fuel_ir::DType::F32 {
return Err(fuel_ir::Error::Msg(format!(
"qmatmul: activations must be F32, got {:?}",
self.inner.dtype(),
))
.bt());
}
if weight_bytes.inner.dtype() != fuel_ir::DType::U32 {
return Err(fuel_ir::Error::Msg(format!(
"qmatmul: weight_bytes must be U32 (raw block bytes reinterpreted), got {:?}",
weight_bytes.inner.dtype(),
))
.bt());
}
let a_shape = self.inner.shape();
let a_dims = a_shape.dims();
if a_dims.len() < 2 {
return Err(fuel_ir::Error::Msg(format!(
"qmatmul: activations must be rank >= 2, got {a_dims:?}",
))
.bt());
}
if a_dims[a_dims.len() - 1] != k {
return Err(fuel_ir::Error::Msg(format!(
"qmatmul: last dim of activations ({}) must equal k ({k})",
a_dims[a_dims.len() - 1],
))
.bt());
}
let block_size = quant_type.elements_per_block();
if !k.is_multiple_of(block_size) {
return Err(fuel_ir::Error::Msg(format!(
"qmatmul: k={k} must be a multiple of {quant_type:?}'s block size ({block_size})",
))
.bt());
}
let expected_bytes = n * (k / block_size) * quant_type.bytes_per_block();
let expected_u32_elems = expected_bytes / 4;
let actual_elems = weight_bytes.inner.shape().elem_count();
if actual_elems != expected_u32_elems {
return Err(fuel_ir::Error::Msg(format!(
"qmatmul: weight_bytes has {actual_elems} u32 elements, expected {expected_u32_elems} for N={n}, K={k}, {quant_type:?}",
)).bt());
}
Ok(Self {
inner: self.inner.qmatmul(&weight_bytes.inner, quant_type, k, n),
})
}
pub fn transpose(&self) -> std::result::Result<Self, fuel_ir::Error> {
Ok(Self {
inner: self.inner.try_transpose()?,
})
}
pub fn permute<D: Dims>(&self, axes: D) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let axes = axes.to_indexes(&shape, "permute")?;
Ok(Self {
inner: self.inner.try_permute(&axes)?,
})
}
pub fn reshape(&self, shape: impl Into<Shape>) -> std::result::Result<Self, fuel_ir::Error> {
Ok(Self {
inner: self.inner.try_reshape(shape)?,
})
}
pub fn squeeze<D: Dim>(&self, dim: D) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dim = dim.to_index(&shape, "squeeze")?;
Ok(Self {
inner: self.inner.squeeze(dim)?,
})
}
pub fn broadcast_to(
&self,
shape: impl Into<Shape>,
) -> std::result::Result<Self, fuel_ir::Error> {
Ok(Self {
inner: self.inner.try_broadcast_to(shape)?,
})
}
pub fn layer_norm_affine(
&self,
gain: std::sync::Arc<[f32]>,
bias: std::sync::Arc<[f32]>,
eps: f64,
) -> std::result::Result<Self, fuel_ir::Error> {
let hidden = gain.len();
debug_assert_eq!(
bias.len(),
hidden,
"layer_norm_affine: gain ({}) and bias ({}) must have the same length",
gain.len(),
bias.len()
);
let normed = self.layer_norm_last_dim(eps)?;
let dims_v: Vec<usize> = self.inner.shape().dims().to_vec();
let mut affine_shape = vec![1_usize; dims_v.len()];
affine_shape[dims_v.len() - 1] = hidden;
let bc_shape = Shape::from_dims(&dims_v);
let g = normed
.const_f32_like(gain, Shape::from_dims(&[hidden]))?
.reshape(Shape::from_dims(&affine_shape))?
.broadcast_to(bc_shape.clone())?;
let b = normed
.const_f32_like(bias, Shape::from_dims(&[hidden]))?
.reshape(Shape::from_dims(&affine_shape))?
.broadcast_to(bc_shape)?;
normed.mul(&g)?.add(&b)
}
pub fn l2_normalize<D: Dim>(
&self,
dim: D,
eps: f64,
) -> std::result::Result<Self, fuel_ir::Error> {
let sq = self.sqr();
let summed = sq.sum_keepdim(dim)?;
let with_eps = if eps == 0.0 {
summed
} else {
summed.add_scalar(eps)
};
let l2 = with_eps.sqrt();
let dims_v: Vec<usize> = self.inner.shape().dims().to_vec();
let l2_bc = l2.broadcast_to(Shape::from_dims(&dims_v))?;
self.div(&l2_bc)
}
pub fn repeat_interleave<D: Dim>(
&self,
dim: D,
repeats: usize,
) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dim = dim.to_index(&shape, "repeat_interleave")?;
if repeats == 0 {
return Err(fuel_ir::Error::Msg("repeat_interleave: repeats must be ≥ 1".into()).bt());
}
if repeats == 1 {
return Ok(self.clone());
}
let dims_v: Vec<usize> = shape.dims().to_vec();
let mut unsq_shape = dims_v.clone();
unsq_shape.insert(dim + 1, 1);
let mut bc_shape = unsq_shape.clone();
bc_shape[dim + 1] = repeats;
let unsq = self.reshape(Shape::from_dims(&unsq_shape))?;
let bc = unsq.broadcast_to(Shape::from_dims(&bc_shape))?;
let mut out_shape = dims_v.clone();
out_shape[dim] *= repeats;
bc.reshape(Shape::from_dims(&out_shape))
}
pub fn slice<D: Dim>(
&self,
dim: D,
start: usize,
len: usize,
) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dim = dim.to_index(&shape, "slice")?;
let dim_size = shape.dims()[dim];
if start.saturating_add(len) > dim_size {
return Err(fuel_ir::Error::Msg(format!(
"slice: start={start} + len={len} exceeds dim {dim} size {dim_size}",
))
.bt());
}
Ok(Self {
inner: self.inner.slice(dim, start, len),
})
}
pub fn concat<D: Dim>(
&self,
other: &Self,
dim: D,
) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dim = dim.to_index(&shape, "concat")?;
let self_dims = shape.dims().to_vec();
let other_dims = other.inner.shape().dims().to_vec();
if self_dims.len() != other_dims.len() {
return Err(fuel_ir::Error::Msg(format!(
"concat: rank mismatch lhs={self_dims:?} rhs={other_dims:?}",
))
.bt());
}
for (i, (&a, &b)) in self_dims.iter().zip(other_dims.iter()).enumerate() {
if i != dim && a != b {
return Err(fuel_ir::Error::Msg(format!(
"concat: dim {i} mismatch lhs={a} rhs={b} (concat dim is {dim})",
))
.bt());
}
}
Ok(Self {
inner: self.inner.concat(&other.inner, dim),
})
}
pub fn add_scalar(&self, c: f64) -> Self {
Self {
inner: self.inner.add_scalar(c),
}
}
pub fn mul_scalar(&self, c: f64) -> Self {
Self {
inner: self.inner.mul_scalar(c),
}
}
pub fn argmax_dim<D: Dim>(&self, dim: D) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dim = dim.to_index(&shape, "argmax_dim")?;
Ok(Self {
inner: self.inner.argmax_dim(dim),
})
}
pub fn realize_u32(&self) -> Vec<u32> {
let graph = self.inner.graph().clone();
let target = self.inner.id();
let device = crate::Device::cpu();
crate::pipelined_bridge::realize_one_as::<u32>(&graph, target, &device)
.expect("realize_u32 via PipelinedExecutor")
}
pub fn realize_u8(&self) -> Vec<u8> {
let graph = self.inner.graph().clone();
let target = self.inner.id();
let device = crate::Device::cpu();
crate::pipelined_bridge::realize_one_bytes(&graph, target, &device)
.expect("realize_u8 via PipelinedExecutor")
}
pub fn sum_all(&self) -> Self {
Self {
inner: self.inner.sum_all(),
}
}
pub fn mean_all(&self) -> Self {
Self {
inner: self.inner.mean_all(),
}
}
pub fn max_all(&self) -> Self {
Self {
inner: self.inner.max_all(),
}
}
pub fn min_all(&self) -> Self {
Self {
inner: self.inner.min_all(),
}
}
pub fn sum_dim<D: Dim>(&self, dim: D) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dim = dim.to_index(&shape, "sum_dim")?;
Ok(Self {
inner: self.inner.sum_dim(dim),
})
}
pub fn max_dim<D: Dim>(&self, dim: D) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dim = dim.to_index(&shape, "max_dim")?;
Ok(Self {
inner: self.inner.max_dim(dim),
})
}
pub fn min_dim<D: Dim>(&self, dim: D) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dim = dim.to_index(&shape, "min_dim")?;
Ok(Self {
inner: self.inner.min_dim(dim),
})
}
pub fn clamp(&self, min: f64, max: f64) -> Self {
Self {
inner: self.inner.clamp(min, max),
}
}
pub fn mean_dim<D: Dim>(&self, dim: D) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dim = dim.to_index(&shape, "mean_dim")?;
Ok(Self {
inner: self.inner.mean_dim(dim),
})
}
pub fn reduce_sum_to(&self, target: impl Into<Shape>) -> Self {
Self {
inner: self.inner.reduce_sum_to(target),
}
}
pub fn reduce_max_to(&self, target: impl Into<Shape>) -> Self {
Self {
inner: self.inner.reduce_max_to(target),
}
}
pub fn softmax_last_dim(&self) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dims = shape.dims();
if dims.is_empty() {
return Err(fuel_ir::Error::Msg(
"softmax_last_dim: input must be rank >= 1, got scalar".into(),
)
.bt());
}
Ok(Self {
inner: self.inner.softmax_last_dim(),
})
}
pub fn nf4_matmul(&self, w_packed: &Self, absmax: &Self, block_size: usize) -> Self {
Self {
inner: self
.inner
.nf4_matmul(&w_packed.inner, &absmax.inner, block_size),
}
}
pub fn ssd_chunk_scan(
&self,
dt: &Self,
a: &Self,
b: &Self,
c: &Self,
chunk_size: usize,
) -> Self {
Self {
inner: self
.inner
.ssd_chunk_scan(&dt.inner, &a.inner, &b.inner, &c.inner, chunk_size),
}
}
pub fn selective_scan(
&self,
delta: &Self,
a: &Self,
b: &Self,
c: &Self,
delta_softplus: bool,
) -> Self {
Self {
inner: self.inner.selective_scan(
&delta.inner,
&a.inner,
&b.inner,
&c.inner,
delta_softplus,
),
}
}
pub fn selective_scan_bundled(
&self,
delta: &Self,
a: &Self,
b: &Self,
c: &Self,
delta_softplus: bool,
) -> std::result::Result<(Self, Self), fuel_ir::Error> {
let (y, last_state) = self.inner.selective_scan_bundled(
&delta.inner,
&a.inner,
&b.inner,
&c.inner,
delta_softplus,
)?;
Ok((Self { inner: y }, Self { inner: last_state }))
}
pub fn ssd_chunk_scan_bundled(
&self,
dt: &Self,
a: &Self,
b: &Self,
c: &Self,
chunk_size: usize,
) -> std::result::Result<(Self, Self), fuel_ir::Error> {
let (y, last_state) = self
.inner
.ssd_chunk_scan_bundled(&dt.inner, &a.inner, &b.inner, &c.inner, chunk_size)?;
Ok((Self { inner: y }, Self { inner: last_state }))
}
pub fn nonzero_indices_bundled(
&self,
count_sym: fuel_ir::SymId,
) -> std::result::Result<(Self, Self), fuel_ir::Error> {
let (indices, count) = self.inner.nonzero_indices_bundled(count_sym)?;
Ok((Self { inner: indices }, Self { inner: count }))
}
pub fn fresh_data_determined_sym(&self) -> fuel_ir::SymId {
let graph = self.inner.graph();
let g = graph.read().unwrap();
g.next_data_determined_sym()
}
pub fn causal_conv1d(&self, weight: &Self, bias: &Self, use_silu: bool) -> Self {
Self {
inner: self
.inner
.causal_conv1d(&weight.inner, &bias.inner, use_silu),
}
}
pub fn fused_softmax_cross_entropy(
&self,
targets: &Self,
reduction: fuel_graph::registry::Reduction,
ignore_index: i64,
) -> Self {
Self {
inner: self
.inner
.fused_softmax_cross_entropy(&targets.inner, reduction, ignore_index),
}
}
pub fn layer_norm_last_dim(&self, eps: f64) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dims = shape.dims();
if dims.last().copied().unwrap_or(0) == 0 {
return Err(fuel_ir::Error::Msg(format!(
"layer_norm_last_dim: input must have non-zero last dim, got {dims:?}",
))
.bt());
}
Ok(Self {
inner: self.inner.layer_norm_last_dim(eps),
})
}
pub fn rms_norm_last_dim(&self, eps: f64) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dims = shape.dims();
if dims.last().copied().unwrap_or(0) == 0 {
return Err(fuel_ir::Error::Msg(format!(
"rms_norm_last_dim: input must have non-zero last dim, got {dims:?}",
))
.bt());
}
Ok(Self {
inner: self.inner.rms_norm_last_dim(eps),
})
}
pub fn rope(&self, base: f64, start_pos: usize) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dims = shape.dims();
if dims.len() < 2 {
return Err(fuel_ir::Error::Msg(format!(
"rope: input must have rank >= 2, got {dims:?}",
))
.bt());
}
Ok(Self {
inner: self.inner.rope(base, start_pos),
})
}
pub fn rope_with_tables(
&self,
cos: &Self,
sin: &Self,
) -> std::result::Result<Self, fuel_ir::Error> {
if self.inner.dtype() != fuel_ir::DType::F32 {
return Err(fuel_ir::Error::Msg(format!(
"rope: only f32 is supported today, got {:?} (cast explicitly for other dtypes)",
self.inner.dtype(),
))
.bt());
}
let in_shape = self.inner.shape();
let dims = in_shape.dims();
let rank = dims.len();
if rank < 2 {
return Err(fuel_ir::Error::Msg(format!(
"rope: input must have rank >= 2, got {dims:?}",
))
.bt());
}
let seq = dims[rank - 2];
let d = dims[rank - 1];
if !d.is_multiple_of(2) {
return Err(fuel_ir::Error::Msg(format!("rope: feature dim {d} must be even",)).bt());
}
let cos_shape = cos.inner.shape();
let cos_dims = cos_shape.dims();
if cos_dims != [seq, d] {
return Err(fuel_ir::Error::Msg(format!(
"rope_with_tables: cos shape {cos_dims:?} does not match [seq, d] = [{seq}, {d}]",
))
.bt());
}
let sin_shape = sin.inner.shape();
let sin_dims = sin_shape.dims();
if sin_dims != [seq, d] {
return Err(fuel_ir::Error::Msg(format!(
"rope_with_tables: sin shape {sin_dims:?} does not match [seq, d] = [{seq}, {d}]",
))
.bt());
}
Ok(Self {
inner: self.inner.rope_with_tables(&cos.inner, &sin.inner),
})
}
pub fn rope_with_tables_decomposed(
&self,
cos: &Self,
sin: &Self,
) -> std::result::Result<Self, fuel_ir::Error> {
if self.inner.dtype() != fuel_ir::DType::F32 {
return Err(fuel_ir::Error::Msg(format!(
"rope: only f32 is supported today, got {:?} (cast explicitly for other dtypes)",
self.inner.dtype(),
))
.bt());
}
let in_shape = self.inner.shape();
let dims = in_shape.dims();
let rank = dims.len();
if rank < 2 {
return Err(fuel_ir::Error::Msg(format!(
"rope: input must have rank >= 2, got {dims:?}",
))
.bt());
}
let seq = dims[rank - 2];
let d = dims[rank - 1];
if !d.is_multiple_of(2) {
return Err(fuel_ir::Error::Msg(format!("rope: feature dim {d} must be even",)).bt());
}
let cos_dims = cos.inner.shape();
if cos_dims.dims() != [seq, d] {
return Err(fuel_ir::Error::Msg(format!(
"rope_with_tables_decomposed: cos shape {:?} does not match [seq, d] = [{seq}, {d}]",
cos_dims.dims(),
)).bt());
}
let sin_dims = sin.inner.shape();
if sin_dims.dims() != [seq, d] {
return Err(fuel_ir::Error::Msg(format!(
"rope_with_tables_decomposed: sin shape {:?} does not match [seq, d] = [{seq}, {d}]",
sin_dims.dims(),
)).bt());
}
Ok(Self {
inner: self
.inner
.rope_with_tables_decomposed(&cos.inner, &sin.inner),
})
}
pub fn index_select<D: Dim>(
&self,
dim: D,
indices: &Self,
) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dim = dim.to_index(&shape, "index_select")?;
if indices.inner.dtype() != fuel_ir::DType::U32 {
return Err(fuel_ir::Error::Msg(format!(
"index_select: index tensor must be U32, got {:?}",
indices.inner.dtype(),
))
.bt());
}
let idx_shape = indices.inner.shape();
let idx_dims = idx_shape.dims();
if idx_dims.len() != 1 {
return Err(fuel_ir::Error::Msg(format!(
"index_select: index tensor must be rank 1, got {idx_dims:?}",
))
.bt());
}
Ok(Self {
inner: self.inner.index_select(dim, &indices.inner),
})
}
pub fn gather<D: Dim>(
&self,
dim: D,
indices: &Self,
) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dim = dim.to_index(&shape, "gather")?;
if indices.inner.dtype() != fuel_ir::DType::U32 {
return Err(fuel_ir::Error::Msg(format!(
"gather: index tensor must be U32, got {:?}",
indices.inner.dtype(),
))
.bt());
}
let data_rank = shape.dims().len();
let idx_shape = indices.inner.shape();
let idx_rank = idx_shape.dims().len();
if data_rank != idx_rank {
return Err(fuel_ir::Error::Msg(format!(
"gather: data and index must have the same rank, got {data_rank} vs {idx_rank}",
))
.bt());
}
Ok(Self {
inner: self.inner.gather(dim, &indices.inner),
})
}
pub fn to_dtype(&self, dtype: DType) -> std::result::Result<Self, fuel_ir::Error> {
if self.inner.dtype() == dtype {
return Ok(self.clone());
}
Ok(Self {
inner: self.inner.cast(dtype),
})
}
pub fn detach(&self) -> Self {
self.clone()
}
pub fn track_op(&self) -> bool {
true
}
pub fn realize_f32(&self) -> Vec<f32> {
let graph = self.inner.graph().clone();
let target = self.inner.id();
let device = crate::Device::cpu();
crate::pipelined_bridge::realize_one_as::<f32>(&graph, target, &device)
.expect("realize_f32 via PipelinedExecutor")
}
pub fn realize_f32_reference(&self) -> Vec<f32> {
let graph = self.inner.graph().clone();
let target = self.inner.id();
let device = crate::Device::cpu();
crate::pipelined_bridge::realize_one_reference_as::<f32>(&graph, target, &device)
.expect("realize_f32_reference via PipelinedExecutor")
}
pub fn realize_f64(&self) -> Vec<f64> {
let graph = self.inner.graph().clone();
let target = self.inner.id();
let device = crate::Device::cpu();
crate::pipelined_bridge::realize_one_as::<f64>(&graph, target, &device)
.expect("realize_f64 via PipelinedExecutor")
}
pub fn realize_bf16(&self) -> Vec<half::bf16> {
let graph = self.inner.graph().clone();
let target = self.inner.id();
let device = crate::Device::cpu();
crate::pipelined_bridge::realize_one_as::<half::bf16>(&graph, target, &device)
.expect("realize_bf16 via PipelinedExecutor")
}
pub fn realize_f16(&self) -> Vec<half::f16> {
let graph = self.inner.graph().clone();
let target = self.inner.id();
let device = crate::Device::cpu();
crate::pipelined_bridge::realize_one_as::<half::f16>(&graph, target, &device)
.expect("realize_f16 via PipelinedExecutor")
}
#[cfg(feature = "cuda")]
pub fn realize_f32_cuda(&self, device: &fuel_cuda_backend::CudaDevice) -> Vec<f32> {
let graph = self.inner.graph().clone();
let target = self.inner.id();
let fc_device: crate::Device = device.clone().into();
crate::pipelined_bridge::realize_one_as::<f32>(&graph, target, &fc_device)
.expect("realize_f32_cuda via PipelinedExecutor")
}
#[cfg(feature = "vulkan")]
pub fn realize_f32_vulkan(
&self,
backend: &std::sync::Arc<fuel_vulkan_backend::VulkanBackend>,
) -> Vec<f32> {
let graph = self.inner.graph().clone();
let target = self.inner.id();
let fc_device: crate::Device = backend.clone().into();
crate::pipelined_bridge::realize_one_as::<f32>(&graph, target, &fc_device)
.expect("realize_f32_vulkan via PipelinedExecutor")
}
}
pub fn realize_many_f32(tensors: &[&Tensor]) -> Vec<Vec<f32>> {
if tensors.is_empty() {
return Vec::new();
}
let graph = tensors[0].inner.graph().clone();
let targets: Vec<fuel_graph::NodeId> = tensors.iter().map(|t| t.inner.id()).collect();
let device = crate::Device::cpu();
crate::pipelined_bridge::realize_many_as::<f32>(&graph, &targets, &device)
.expect("realize_many_f32 via PipelinedExecutor")
}
#[cfg(feature = "cuda")]
pub fn realize_many_f32_cuda(
tensors: &[&Tensor],
device: &fuel_cuda_backend::CudaDevice,
) -> Vec<Vec<f32>> {
if tensors.is_empty() {
return Vec::new();
}
let graph = tensors[0].inner.graph().clone();
let targets: Vec<fuel_graph::NodeId> = tensors.iter().map(|t| t.inner.id()).collect();
let fc_device: crate::Device = device.clone().into();
crate::pipelined_bridge::realize_many_as::<f32>(&graph, &targets, &fc_device)
.expect("realize_many_f32_cuda via PipelinedExecutor")
}
#[cfg(test)]
mod tests {
use super::*;
use crate::Shape;
#[test]
fn constructors_wrap_graph_tensor_correctly() {
let t =
Tensor::from_f32(vec![1.0, 2.0, 3.0], Shape::from_dims(&[3]), &Device::cpu()).unwrap();
assert_eq!(t.dtype(), DType::F32);
assert_eq!(t.shape().dims(), &[3]);
assert_eq!(t.rank(), 1);
assert_eq!(t.elem_count(), 3);
}
#[test]
fn add_builds_add_node_in_underlying_graph() {
let a =
Tensor::from_f32(vec![1.0, 2.0, 3.0], Shape::from_dims(&[3]), &Device::cpu()).unwrap();
let b = a
.const_f32_like(vec![4.0, 5.0, 6.0], Shape::from_dims(&[3]))
.unwrap();
let c = a.add(&b).unwrap();
assert_eq!(c.shape().dims(), &[3]);
assert!(std::sync::Arc::ptr_eq(
c.graph_tensor().graph(),
a.graph_tensor().graph(),
));
}
#[test]
fn chained_lazy_method_call_builds_sensible_graph() {
let x = Tensor::from_f32(
vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
Shape::from_dims(&[2, 3]),
&Device::cpu(),
)
.unwrap();
let w = x
.const_f32_like(
vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 1.0],
Shape::from_dims(&[3, 3]),
)
.unwrap();
let y = x
.rms_norm_last_dim(1e-6)
.unwrap()
.matmul(&w)
.unwrap()
.relu();
assert_eq!(y.shape().dims(), &[2, 3]);
assert_eq!(y.dtype(), DType::F32);
}
#[test]
fn rope_through_lazy_wrapper() {
let x = Tensor::from_f32(
vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0],
Shape::from_dims(&[2, 4]),
&Device::cpu(),
)
.unwrap();
let y = x.rope(10000.0, 0).unwrap();
assert_eq!(y.shape().dims(), &[2, 4]);
assert_eq!(y.dtype(), DType::F32);
}
#[test]
fn rope_delta_rotate_equals_direct_shift() {
let dev = Device::cpu();
let (n_kv_heads, head_dim, bs) = (2usize, 4usize, 4usize);
let theta = 10000.0;
let (p0, m) = (0usize, 8usize); let raw: Vec<f32> = (0..bs * n_kv_heads * head_dim)
.map(|i| (i as f32 * 0.017).sin())
.collect();
let rope_pool_at = |data: &[f32], start_pos: usize| -> Vec<f32> {
let k = Tensor::from_f32(
data.to_vec(),
Shape::from_dims(&[bs, n_kv_heads, head_dim]),
&dev,
)
.unwrap();
let k4 = k
.permute([1, 0, 2])
.unwrap()
.reshape(Shape::from_dims(&[1, n_kv_heads, bs, head_dim]))
.unwrap();
let (c, s) = k4.rope_tables_const(theta, start_pos, bs, head_dim);
k4.rope_with_tables_decomposed(&c, &s)
.unwrap()
.reshape(Shape::from_dims(&[n_kv_heads, bs, head_dim]))
.unwrap()
.permute([1, 0, 2])
.unwrap()
.realize_f32()
};
let direct = rope_pool_at(&raw, p0 + m); let at_p0 = rope_pool_at(&raw, p0); let shifted =
Tensor::rope_delta_rotate_block_f32(&dev, &at_p0, theta, m, bs, n_kv_heads, head_dim);
let maxdiff = direct
.iter()
.zip(&shifted)
.map(|(a, b)| (a - b).abs())
.fold(0.0f32, f32::max);
assert!(
maxdiff < 1e-5,
"delta-rotate == direct shift (maxdiff {maxdiff})"
);
}
#[test]
fn cast_switches_dtype_through_wrapper() {
let x =
Tensor::from_f32(vec![1.0, 2.0, 3.0], Shape::from_dims(&[3]), &Device::cpu()).unwrap();
let y = x.to_dtype(DType::F64).unwrap();
assert_eq!(y.dtype(), DType::F64);
assert_eq!(y.shape().dims(), &[3]);
}
#[test]
fn indexing_builds_correct_output_shape() {
let data =
Tensor::from_f32(vec![1.0; 12], Shape::from_dims(&[3, 4]), &Device::cpu()).unwrap();
let idx = data
.const_u32_like(vec![0, 2, 1], Shape::from_dims(&[3]))
.unwrap();
let out = data.index_select(0, &idx).unwrap();
assert_eq!(out.shape().dims(), &[3, 4]);
}
#[test]
fn realize_f32_executes_the_underlying_graph() {
let a =
Tensor::from_f32(vec![1.0, 2.0, 3.0], Shape::from_dims(&[3]), &Device::cpu()).unwrap();
let b = a
.const_f32_like(vec![4.0, 5.0, 6.0], Shape::from_dims(&[3]))
.unwrap();
let c = a.add(&b).unwrap().mul(&a).unwrap();
let result = c.realize_f32();
assert_eq!(result, vec![5.0, 14.0, 27.0]);
}
#[test]
fn realize_f32_relu_and_maximum_propagate_nan_end_to_end() {
let a = Tensor::from_f32(
vec![f32::NAN, -2.0, 3.0],
Shape::from_dims(&[3]),
&Device::cpu(),
)
.unwrap();
let relu_result = a.relu().realize_f32();
assert!(
relu_result[0].is_nan(),
"relu(NaN) must propagate NaN through realize, got {}",
relu_result[0]
);
assert_eq!(relu_result[1], 0.0, "relu(-2) must clip to 0");
assert_eq!(relu_result[2], 3.0, "relu(3) must pass through");
let b = a
.const_f32_like(vec![1.0, f32::NAN, 2.0], Shape::from_dims(&[3]))
.unwrap();
let max_result = a.maximum(&b).unwrap().realize_f32();
assert!(
max_result[0].is_nan(),
"maximum(NaN, 1) must propagate NaN through realize, got {}",
max_result[0]
);
assert!(
max_result[1].is_nan(),
"maximum(-2, NaN) must propagate NaN through realize, got {}",
max_result[1]
);
assert_eq!(max_result[2], 3.0, "maximum(3, 2) non-NaN sanity");
}
#[test]
fn fmax_fmin_ieee_suppress_nan_where_prop_propagates() {
let a = Tensor::from_f32(
vec![f32::NAN, 2.0, 5.0],
Shape::from_dims(&[3]),
&Device::cpu(),
)
.unwrap();
let b = a
.const_f32_like(vec![3.0, f32::NAN, 4.0], Shape::from_dims(&[3]))
.unwrap();
let fmax = a.fmax_ieee(&b).unwrap().realize_f32();
assert_eq!(fmax[0], 3.0, "fmax_ieee(NaN,3) suppresses NaN → 3");
assert_eq!(fmax[1], 2.0, "fmax_ieee(2,NaN) suppresses NaN → 2");
assert_eq!(fmax[2], 5.0, "fmax_ieee(5,4) → 5");
let max = a.maximum(&b).unwrap().realize_f32();
assert!(
max[0].is_nan() && max[1].is_nan(),
"maximum PROPAGATES NaN — the divergence that makes fmax_ieee distinct"
);
let fmin = a.fmin_ieee(&b).unwrap().realize_f32();
assert_eq!(fmin[0], 3.0, "fmin_ieee(NaN,3) suppresses NaN → 3");
assert_eq!(fmin[1], 2.0, "fmin_ieee(2,NaN) suppresses NaN → 2");
assert_eq!(fmin[2], 4.0, "fmin_ieee(5,4) → 4");
let min = a.minimum(&b).unwrap().realize_f32();
assert!(
min[0].is_nan() && min[1].is_nan(),
"minimum PROPAGATES NaN — the divergence that makes fmin_ieee distinct"
);
}
#[test]
fn fmax_fmin_ieee_signed_zero_tie_operand_a_wins_bitwise() {
let a = Tensor::from_f32(vec![-0.0, 0.0], Shape::from_dims(&[2]), &Device::cpu()).unwrap();
let b = a
.const_f32_like(vec![0.0, -0.0], Shape::from_dims(&[2]))
.unwrap();
let fmax = a.fmax_ieee(&b).unwrap().realize_f32();
assert_eq!(
fmax[0].to_bits(),
(-0.0f32).to_bits(),
"fmax_ieee(-0.0,+0.0): operand a wins the tie → -0.0 (bitwise)"
);
assert_eq!(
fmax[1].to_bits(),
(0.0f32).to_bits(),
"fmax_ieee(+0.0,-0.0): operand a wins the tie → +0.0 (bitwise)"
);
let fmin = a.fmin_ieee(&b).unwrap().realize_f32();
assert_eq!(
fmin[0].to_bits(),
(-0.0f32).to_bits(),
"fmin_ieee(-0.0,+0.0): operand a wins the tie → -0.0 (bitwise)"
);
assert_eq!(
fmin[1].to_bits(),
(0.0f32).to_bits(),
"fmin_ieee(+0.0,-0.0): operand a wins the tie → +0.0 (bitwise)"
);
}
#[test]
fn rem_trunc_diverges_from_floored_rem_on_opposite_signs() {
let a = Tensor::from_f32(
vec![-7.0, 7.0, -7.0],
Shape::from_dims(&[3]),
&Device::cpu(),
)
.unwrap();
let b = a
.const_f32_like(vec![3.0, -3.0, -3.0], Shape::from_dims(&[3]))
.unwrap();
let rt = a.rem_trunc(&b).unwrap().realize_f32();
assert_eq!(rt[0], -1.0, "-7 rem_trunc 3 = -1 (sign of dividend)");
assert_eq!(rt[1], 1.0, "7 rem_trunc -3 = 1 (sign of dividend)");
assert_eq!(rt[2], -1.0, "-7 rem_trunc -3 = -1");
let rf = a.rem(&b).unwrap().realize_f32();
assert_eq!(
rf[0], 2.0,
"-7 rem 3 = 2 (floored, sign of divisor) — the divergence from rem_trunc"
);
}
#[test]
fn trunc_edges_signed_zero_nan_inf_bitwise() {
let q = Tensor::from_f32(
vec![
2.7,
-2.7,
-0.0,
0.0,
f32::NAN,
f32::INFINITY,
f32::NEG_INFINITY,
],
Shape::from_dims(&[7]),
&Device::cpu(),
)
.unwrap();
let t = q.trunc().realize_f32();
assert_eq!(t[0], 2.0, "trunc(2.7) = 2 (toward zero)");
assert_eq!(t[1], -2.0, "trunc(-2.7) = -2 (toward zero, NOT -3)");
assert_eq!(
t[2].to_bits(),
(-0.0f32).to_bits(),
"trunc(-0.0) must PRESERVE -0.0 (bitwise, not +0.0)"
);
assert_eq!(
t[3].to_bits(),
(0.0f32).to_bits(),
"trunc(+0.0) = +0.0 (bitwise)"
);
assert!(t[4].is_nan(), "trunc(NaN) = NaN");
assert_eq!(t[5], f32::INFINITY, "trunc(+inf) = +inf");
assert_eq!(t[6], f32::NEG_INFINITY, "trunc(-inf) = -inf");
}
#[test]
fn realize_f32_matmul_hand_computed() {
let a = Tensor::from_f32(
vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
Shape::from_dims(&[2, 3]),
&Device::cpu(),
)
.unwrap();
let b = a
.const_f32_like(
vec![7.0, 8.0, 9.0, 10.0, 11.0, 12.0],
Shape::from_dims(&[3, 2]),
)
.unwrap();
let c = a.matmul(&b).unwrap();
assert_eq!(c.shape().dims(), &[2, 2]);
assert_eq!(c.realize_f32(), vec![58.0, 64.0, 139.0, 154.0]);
}
#[test]
fn fast_and_reference_agree_on_medium_matmul() {
let m = 64;
let k = 96;
let n = 32;
let a_data: Vec<f32> = (0..m * k).map(|i| (i as f32 * 0.01).sin()).collect();
let b_data: Vec<f32> = (0..k * n).map(|i| (i as f32 * 0.013).cos()).collect();
let a = Tensor::from_f32(a_data, Shape::from_dims(&[m, k]), &Device::cpu()).unwrap();
let b = a.const_f32_like(b_data, Shape::from_dims(&[k, n])).unwrap();
let c = a.matmul(&b).unwrap();
let fast = c.realize_f32();
let reference = c.realize_f32();
assert_eq!(fast.len(), reference.len());
for (i, (&f, &r)) in fast.iter().zip(&reference).enumerate() {
let diff = (f - r).abs();
let rel = if r.abs() > 1e-6 { diff / r.abs() } else { 0.0 };
assert!(
diff < 1e-4 || rel < 1e-3,
"at index {i}: fast={f}, reference={r}, diff={diff}, rel={rel}",
);
}
}
#[allow(clippy::too_many_arguments)]
fn ref_attention(
q: &[f32],
k: &[f32],
v: &[f32],
hq: usize,
hkv: usize,
sq: usize,
sk: usize,
d: usize,
scale: f32,
causal: bool,
softcap: Option<f32>,
alibi: Option<&[f32]>,
) -> Vec<f32> {
let g = hq / hkv;
let mut out = vec![0.0f32; hq * sq * d];
for h in 0..hq {
let hk = h / g;
let qh = &q[h * sq * d..(h + 1) * sq * d];
let kh = &k[hk * sk * d..(hk + 1) * sk * d];
let vh = &v[hk * sk * d..(hk + 1) * sk * d];
for i in 0..sq {
let mut scores = vec![0.0f32; sk];
for j in 0..sk {
let mut s = 0.0f32;
for l in 0..d {
s += qh[i * d + l] * kh[j * d + l];
}
let mut sc = scale * s;
if let Some(cap) = softcap {
sc = cap * (sc / cap).tanh();
}
if let Some(slopes) = alibi {
sc += slopes[h] * (j as f32 - i as f32);
}
scores[j] = if causal && j > i {
f32::NEG_INFINITY
} else {
sc
};
}
let m = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let mut sum = 0.0f32;
let exps: Vec<f32> = scores
.iter()
.map(|&s| {
let e = (s - m).exp();
sum += e;
e
})
.collect();
for l in 0..d {
let mut o = 0.0f32;
for j in 0..sk {
o += (exps[j] / sum) * vh[j * d + l];
}
out[(h * sq + i) * d + l] = o;
}
}
}
out
}
fn flash_decompose_vs_reference(
hq: usize,
hkv: usize,
causal: bool,
softcap: Option<f32>,
alibi: bool,
) {
let dev = Device::cpu();
let (sq, sk, d) = (2usize, 2usize, 2usize);
let q_data: Vec<f32> = (0..hq * sq * d).map(|i| (i as f32 * 0.1).sin()).collect();
let k_data: Vec<f32> = (0..hkv * sk * d).map(|i| (i as f32 * 0.13).cos()).collect();
let v_data: Vec<f32> = (0..hkv * sk * d).map(|i| i as f32 * 0.07 + 1.0).collect();
#[allow(clippy::approx_constant)]
let scale = 0.7071f32;
let q = Tensor::from_f32(q_data.clone(), Shape::from_dims(&[1, hq, sq, d]), &dev).unwrap();
let k = q
.const_f32_like(k_data.clone(), Shape::from_dims(&[1, hkv, sk, d]))
.unwrap();
let v = q
.const_f32_like(v_data.clone(), Shape::from_dims(&[1, hkv, sk, d]))
.unwrap();
let alibi_slopes: Option<Vec<f32>> = if alibi {
Some((0..hq).map(|h| 0.5f32.powi(h as i32 + 1)).collect())
} else {
None
};
let alibi_t = alibi_slopes.as_ref().map(|s| {
q.const_f32_like(s.clone(), Shape::from_dims(&[hq]))
.unwrap()
});
let attn = q
.flash_attn(&k, &v, alibi_t.as_ref(), scale, causal, None, None, softcap)
.unwrap();
let graph = attn.inner.graph().clone();
let id = attn.inner.id();
let roots =
fuel_graph::opt::RuleRegistry::lowering_only().optimize_to_fixpoint(&graph, &[id]);
assert_eq!(roots.len(), 1, "lowering should keep a single root");
let got = crate::pipelined_bridge::realize_one_as::<f32>(&graph, roots[0], &dev)
.expect("realize decomposed FlashAttn on CPU");
let expected = ref_attention(
&q_data,
&k_data,
&v_data,
hq,
hkv,
sq,
sk,
d,
scale,
causal,
softcap,
alibi_slopes.as_deref(),
);
assert_eq!(got.len(), expected.len());
for (i, (&g, &e)) in got.iter().zip(&expected).enumerate() {
assert!(
(g - e).abs() < 1e-4,
"FlashAttn decompose mismatch (hq={hq} hkv={hkv} causal={causal} \
softcap={softcap:?} alibi={alibi}) at {i}: got {g}, expected {e}",
);
}
}
#[test]
fn flash_attn_decompose_vanilla() {
flash_decompose_vs_reference(1, 1, false, None, false);
}
#[test]
fn flash_attn_decompose_causal() {
flash_decompose_vs_reference(1, 1, true, None, false);
}
#[test]
fn flash_attn_decompose_gqa_causal() {
flash_decompose_vs_reference(2, 1, true, None, false); }
#[test]
fn flash_attn_decompose_softcap() {
flash_decompose_vs_reference(1, 1, false, Some(0.5), false); }
#[test]
fn flash_attn_decompose_alibi() {
flash_decompose_vs_reference(2, 2, false, None, true);
}
#[test]
fn flash_attn_decompose_alibi_causal() {
flash_decompose_vs_reference(2, 2, true, None, true);
}
#[allow(clippy::too_many_arguments)]
fn ref_decode_attn(
q: &[f32],
k: &[f32],
v: &[f32],
hq: usize,
hkv: usize,
sq: usize,
cap: usize,
kl: usize,
d: usize,
scale: f32,
causal: bool,
) -> Vec<f32> {
let g = hq / hkv;
let offset = kl - sq;
let mut out = vec![0f32; hq * sq * d];
for h in 0..hq {
let hk = h / g;
for i in 0..sq {
let qh = &q[((h * sq) + i) * d..][..d];
let mut scores = vec![f32::NEG_INFINITY; kl];
for j in 0..kl {
if causal && j > offset + i {
continue;
}
let kh = &k[((hk * cap) + j) * d..][..d];
let s: f32 = (0..d).map(|l| qh[l] * kh[l]).sum();
scores[j] = scale * s;
}
let m = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let mut sum = 0f32;
let exps: Vec<f32> = scores
.iter()
.map(|&s| {
let e = (s - m).exp();
sum += e;
e
})
.collect();
for l in 0..d {
let mut o = 0f32;
for j in 0..kl {
o += (exps[j] / sum) * v[((hk * cap) + j) * d + l];
}
out[(h * sq + i) * d + l] = o;
}
}
}
out
}
#[test]
fn flash_attn_decompose_concrete_klen() {
use fuel_graph::Op;
use fuel_graph::registry::FusedOps;
use fuel_ir::DynScalar;
let dev = Device::cpu();
let (hq, hkv, sq, cap, kl, d) = (2usize, 1usize, 2usize, 4usize, 3usize, 2usize);
#[allow(clippy::approx_constant)]
let scale = 0.7071f32;
let causal = true;
let q_data: Vec<f32> = (0..hq * sq * d).map(|i| (i as f32 * 0.1).sin()).collect();
let k_data: Vec<f32> = (0..hkv * cap * d)
.map(|i| (i as f32 * 0.13).cos())
.collect();
let v_data: Vec<f32> = (0..hkv * cap * d).map(|i| i as f32 * 0.07 + 1.0).collect();
let q = Tensor::from_f32(q_data.clone(), Shape::from_dims(&[1, hq, sq, d]), &dev).unwrap();
let k = q
.inner
.const_f32_like(k_data.clone(), Shape::from_dims(&[1, hkv, cap, d]))
.unwrap();
let v = q
.inner
.const_f32_like(v_data.clone(), Shape::from_dims(&[1, hkv, cap, d]))
.unwrap();
let attn = q.inner.flash_attn_dyn(
&k,
&v,
None,
scale,
causal,
None,
None,
None,
DynScalar::Concrete(kl),
);
let graph = attn.graph().clone();
let id = attn.id();
let roots =
fuel_graph::opt::RuleRegistry::lowering_only().optimize_to_fixpoint(&graph, &[id]);
assert_eq!(roots.len(), 1, "lowering should keep a single root");
{
let g = graph.read().unwrap();
let mut stack = vec![roots[0]];
let mut seen = std::collections::HashSet::new();
while let Some(nid) = stack.pop() {
if !seen.insert(nid) {
continue;
}
let node = g.node(nid);
assert!(
!matches!(node.op, Op::Fused(fid, _) if fid == FusedOps::FLASH_ATTN),
"concrete-k_len FlashAttn still fused after lowering (self-return)",
);
for &inp in &node.inputs {
stack.push(inp);
}
}
}
let got = crate::pipelined_bridge::realize_one_as::<f32>(&graph, roots[0], &dev)
.expect("realize decomposed concrete-k_len FlashAttn on CPU");
let expected = ref_decode_attn(
&q_data, &k_data, &v_data, hq, hkv, sq, cap, kl, d, scale, causal,
);
assert_eq!(got.len(), expected.len());
for (i, (&gv, &ev)) in got.iter().zip(&expected).enumerate() {
assert!(
(gv - ev).abs() < 1e-4,
"concrete-k_len decode mismatch at {i}: got {gv}, expected {ev}",
);
}
}
fn assert_f32_bits_eq(got: &[f32], want: &[f32], ctx: &str) {
assert_eq!(got.len(), want.len(), "{ctx}: length mismatch");
for (i, (&g, &w)) in got.iter().zip(want).enumerate() {
assert_eq!(
g.to_bits(),
w.to_bits(),
"{ctx}: bit mismatch at {i}: got {g} (0x{:08x}), want {w} (0x{:08x})",
g.to_bits(),
w.to_bits(),
);
}
}
fn assert_f32_within_ulp(got: &[f32], want: &[f32], max_ulp: i64, ctx: &str) {
assert_eq!(got.len(), want.len(), "{ctx}: length mismatch");
for (i, (&g, &w)) in got.iter().zip(want).enumerate() {
let d = (g.to_bits() as i64 - w.to_bits() as i64).abs();
assert!(
d <= max_ulp,
"{ctx}: {d} ULP at {i} (> {max_ulp}): got {g} (0x{:08x}), want {w} (0x{:08x})",
g.to_bits(),
w.to_bits(),
);
}
}
#[test]
fn reduce_swap_sumdim_vs_reducesumto_bit_exact_on_cpu_backend() {
let dev = Device::cpu();
let (rows, cols) = (2usize, 4usize);
let data: Vec<f32> = vec![1e8, 1.0, -1e8, 1.0, 3.0e7, 2.0, -3.0e7, -1.0];
let mut left_fold = vec![0.0f32; rows];
let mut right_fold = vec![0.0f32; rows];
for r in 0..rows {
let row = &data[r * cols..(r + 1) * cols];
let mut acc = 0.0f32;
for &v in row.iter() {
acc += v;
}
left_fold[r] = acc;
let mut acc_r = 0.0f32;
for &v in row.iter().rev() {
acc_r += v;
}
right_fold[r] = acc_r;
}
assert_ne!(
left_fold, right_fold,
"test data must be accumulation-order-sensitive so the parity has teeth",
);
let x = Tensor::from_f32(data.clone(), Shape::from_dims(&[rows, cols]), &dev).unwrap();
let sumdim = x.inner.sum_dim(1);
let got_sumdim = crate::pipelined_bridge::realize_one_as::<f32>(
&sumdim.graph().clone(),
sumdim.id(),
&dev,
)
.expect("realize SumDim(last) on CPU");
let reducesumto = x.inner.reduce_sum_to(Shape::from_dims(&[rows, 1]));
let got_reducesumto = crate::pipelined_bridge::realize_one_as::<f32>(
&reducesumto.graph().clone(),
reducesumto.id(),
&dev,
)
.expect("realize ReduceSumTo(keepdim) on CPU");
assert_f32_bits_eq(&got_sumdim, &got_reducesumto, "SumDim vs ReduceSumTo");
assert_f32_bits_eq(&got_sumdim, &left_fold, "SumDim vs left-fold");
let maxdim = x.inner.max_dim(1);
let got_maxdim = crate::pipelined_bridge::realize_one_as::<f32>(
&maxdim.graph().clone(),
maxdim.id(),
&dev,
)
.expect("realize MaxDim(last) on CPU");
let reducemaxto = x.inner.reduce_max_to(Shape::from_dims(&[rows, 1]));
let got_reducemaxto = crate::pipelined_bridge::realize_one_as::<f32>(
&reducemaxto.graph().clone(),
reducemaxto.id(),
&dev,
)
.expect("realize ReduceMaxTo(keepdim) on CPU");
let row_max: Vec<f32> = (0..rows)
.map(|r| {
data[r * cols..(r + 1) * cols]
.iter()
.cloned()
.fold(f32::NEG_INFINITY, f32::max)
})
.collect();
assert_f32_bits_eq(&got_maxdim, &got_reducemaxto, "MaxDim vs ReduceMaxTo");
assert_f32_bits_eq(&got_maxdim, &row_max, "MaxDim vs per-row max");
}
#[test]
fn softmax_last_dim_recipe_decompose_matches_legacy_on_cpu() {
let dev = Device::cpu();
let (rows, cols) = (3usize, 5usize);
let data: Vec<f32> = (0..rows * cols)
.map(|i| ((i as f32) * 0.37).sin() * 4.0 - (i as f32) * 0.11)
.collect();
let shape = || Shape::from_dims(&[rows, cols]);
let keepdim = || Shape::from_dims(&[rows, 1]);
let recipe = Tensor::from_f32(data.clone(), shape(), &dev)
.unwrap()
.softmax_last_dim()
.expect("softmax recipe build");
let graph = recipe.inner.graph().clone();
let id = recipe.inner.id();
let roots =
fuel_graph::opt::RuleRegistry::lowering_only().optimize_to_fixpoint(&graph, &[id]);
assert_eq!(roots.len(), 1, "lowering keeps a single root");
let recipe_out = crate::pipelined_bridge::realize_one_as::<f32>(&graph, roots[0], &dev)
.expect("realize recipe-decomposed softmax on CPU");
let x = Tensor::from_f32(data.clone(), shape(), &dev).unwrap().inner;
let m = x.reduce_max_to(keepdim());
let mb = m.broadcast_to(shape());
let s = x.sub(&mb);
let e = s.exp();
let d = e.reduce_sum_to(keepdim());
let db = d.broadcast_to(shape());
let legacy = e.div(&db);
let legacy_out = crate::pipelined_bridge::realize_one_as::<f32>(
&legacy.graph().clone(),
legacy.id(),
&dev,
)
.expect("realize legacy softmax on CPU");
assert_f32_bits_eq(&recipe_out, &legacy_out, "softmax recipe vs legacy");
for r in 0..rows {
let sum: f32 = recipe_out[r * cols..(r + 1) * cols].iter().sum();
assert!((sum - 1.0).abs() < 1e-5, "row {r} softmax sum = {sum}");
}
let fused_out = Tensor::from_f32(data.clone(), shape(), &dev)
.unwrap()
.softmax_last_dim()
.expect("softmax fused build")
.realize_f32();
assert_f32_within_ulp(&recipe_out, &fused_out, 1, "softmax recipe vs fused kernel");
}
#[test]
fn rms_norm_last_dim_recipe_decompose_matches_legacy_on_cpu() {
let dev = Device::cpu();
let (rows, cols) = (3usize, 6usize);
let eps = 1e-6f64;
let data: Vec<f32> = (0..rows * cols)
.map(|i| ((i as f32) * 0.29).cos() * 2.5 + 0.5)
.collect();
let shape = || Shape::from_dims(&[rows, cols]);
let keepdim = || Shape::from_dims(&[rows, 1]);
let recipe = Tensor::from_f32(data.clone(), shape(), &dev)
.unwrap()
.rms_norm_last_dim(eps)
.expect("rms_norm recipe build");
let graph = recipe.inner.graph().clone();
let id = recipe.inner.id();
let roots =
fuel_graph::opt::RuleRegistry::lowering_only().optimize_to_fixpoint(&graph, &[id]);
assert_eq!(roots.len(), 1, "lowering keeps a single root");
let recipe_out = crate::pipelined_bridge::realize_one_as::<f32>(&graph, roots[0], &dev)
.expect("realize recipe-decomposed rms_norm on CPU");
let x = Tensor::from_f32(data.clone(), shape(), &dev).unwrap().inner;
let sq = x.sqr();
let mean = sq.mean_dim(1);
let mean_kd = mean.reshape(keepdim());
let denom_sq = mean_kd.add_scalar(eps);
let denom = denom_sq.sqrt();
let denom_b = denom.broadcast_to(shape());
let legacy = x.div(&denom_b);
let legacy_out = crate::pipelined_bridge::realize_one_as::<f32>(
&legacy.graph().clone(),
legacy.id(),
&dev,
)
.expect("realize legacy rms_norm on CPU");
assert_f32_bits_eq(&recipe_out, &legacy_out, "rms_norm recipe vs legacy");
let fused_out = Tensor::from_f32(data.clone(), shape(), &dev)
.unwrap()
.rms_norm_last_dim(eps)
.expect("rms_norm fused build")
.realize_f32();
assert_f32_within_ulp(
&recipe_out,
&fused_out,
1,
"rms_norm recipe vs fused kernel",
);
}
#[test]
fn selective_scan_decompose_lowers_to_scan_and_matches() {
use fuel_graph::Op;
use fuel_graph::registry::FusedOps;
let dev = Device::cpu();
let u = Tensor::from_f32(vec![2.0f32], Shape::from_dims(&[1, 1, 1]), &dev).unwrap();
let delta = u
.const_f32_like(vec![0.5f32], Shape::from_dims(&[1, 1, 1]))
.unwrap();
let a = u
.const_f32_like(vec![-1.0f32], Shape::from_dims(&[1, 1]))
.unwrap();
let b = u
.const_f32_like(vec![3.0f32], Shape::from_dims(&[1, 1, 1]))
.unwrap();
let c = u
.const_f32_like(vec![4.0f32], Shape::from_dims(&[1, 1, 1]))
.unwrap();
let y = u.selective_scan(&delta, &a, &b, &c, false);
let got = y.realize_f32();
assert_eq!(got.len(), 1);
assert!((got[0] - 12.0).abs() < 1e-4, "fused kernel y: {}", got[0]);
let graph = y.inner.graph().clone();
let id = y.inner.id();
let roots =
fuel_graph::opt::RuleRegistry::lowering_only().optimize_to_fixpoint(&graph, &[id]);
assert_eq!(roots.len(), 1);
let scan_id = {
let g = graph.read().unwrap();
let mut stack = vec![roots[0]];
let mut seen = std::collections::HashSet::new();
let mut scan_id = None;
while let Some(nid) = stack.pop() {
if !seen.insert(nid) {
continue;
}
let node = g.node(nid);
assert!(
!matches!(node.op, Op::Fused(fid, _) if fid == FusedOps::SELECTIVE_SCAN),
"SelectiveScan must lower to Op::Scan, not remain fused"
);
if matches!(node.op, Op::Scan { .. }) {
scan_id = Some(nid);
}
for &inp in &node.inputs {
stack.push(inp);
}
}
scan_id.expect("an Op::Scan terminal must be present after lowering")
};
let ys = {
let mut g = graph.write().unwrap();
fuel_graph::scan::unroll_scan(&mut g, scan_id, 1)
.expect("unroll")
.0[0] };
let oracle = crate::pipelined_bridge::realize_one_as::<f32>(&graph, ys, &dev)
.expect("realize unrolled selective_scan oracle on CPU");
assert!(
oracle.iter().any(|&v| (v - 12.0).abs() < 1e-4),
"unroll oracle must contain y = 12, got {oracle:?}"
);
}
#[test]
fn ssd_chunk_scan_decompose_lowers_to_scan_and_matches() {
use fuel_graph::Op;
use fuel_graph::registry::FusedOps;
let dev = Device::cpu();
let x = Tensor::from_f32(vec![2.0f32], Shape::from_dims(&[1, 1, 1, 1]), &dev).unwrap();
let dt = x
.const_f32_like(vec![0.5f32], Shape::from_dims(&[1, 1, 1]))
.unwrap();
let a = x
.const_f32_like(vec![-1.0f32], Shape::from_dims(&[1]))
.unwrap();
let b = x
.const_f32_like(vec![3.0f32], Shape::from_dims(&[1, 1, 1, 1]))
.unwrap();
let c = x
.const_f32_like(vec![4.0f32], Shape::from_dims(&[1, 1, 1, 1]))
.unwrap();
let y = x.ssd_chunk_scan(&dt, &a, &b, &c, 1);
let got = y.realize_f32();
assert!(
(got[0] - 12.0).abs() < 1e-4,
"fused ssd kernel y: {}",
got[0]
);
let graph = y.inner.graph().clone();
let id = y.inner.id();
let roots =
fuel_graph::opt::RuleRegistry::lowering_only().optimize_to_fixpoint(&graph, &[id]);
let scan_id = {
let g = graph.read().unwrap();
let mut stack = vec![roots[0]];
let mut seen = std::collections::HashSet::new();
let mut scan_id = None;
while let Some(nid) = stack.pop() {
if !seen.insert(nid) {
continue;
}
let node = g.node(nid);
assert!(
!matches!(node.op, Op::Fused(fid, _) if fid == FusedOps::SSD_CHUNK_SCAN),
"SsdChunkScan must lower to Op::Scan, not remain fused"
);
if matches!(node.op, Op::Scan { .. }) {
scan_id = Some(nid);
}
for &inp in &node.inputs {
stack.push(inp);
}
}
scan_id.expect("an Op::Scan terminal must be present after lowering")
};
let ys = {
let mut g = graph.write().unwrap();
fuel_graph::scan::unroll_scan(&mut g, scan_id, 1)
.expect("unroll")
.0[0] };
let oracle = crate::pipelined_bridge::realize_one_as::<f32>(&graph, ys, &dev)
.expect("realize unrolled ssd_chunk_scan oracle on CPU");
assert!(
oracle.iter().any(|&v| (v - 12.0).abs() < 1e-4),
"oracle y=12, got {oracle:?}"
);
}
#[test]
fn affine_scan_bptt_matches_finite_difference() {
use fuel_graph::{Node, NodeHandle, Op, ScanEmit, ScanRole};
use fuel_ir::{DType, Shape};
let dev = Device::cpu();
let build = |init_v: f32,
a_v: f32,
b_v: f32|
-> (
std::sync::Arc<std::sync::RwLock<fuel_graph::Graph>>,
fuel_graph::NodeId,
fuel_graph::NodeHandle,
fuel_graph::NodeHandle,
fuel_graph::NodeHandle,
) {
let init =
NodeHandle::from_f32(vec![init_v], Shape::from_dims(&[1]), dev.as_dyn()).unwrap();
let g = init.graph().clone();
let a = NodeHandle::from_existing(g.clone(), init.id())
.const_f32_like(vec![a_v], Shape::from_dims(&[1]))
.unwrap();
let b = NodeHandle::from_existing(g.clone(), init.id())
.const_f32_like(vec![b_v], Shape::from_dims(&[1]))
.unwrap();
let nc = {
let mut gw = g.write().unwrap();
let s = Shape::from_dims(&[1]);
let hole = gw.push(Node {
op: Op::ScanPlaceholder {
role: ScanRole::Carry,
index: 0,
},
inputs: vec![],
shape: s.clone(),
dtype: DType::F32,
});
let ac = gw.push(Node {
op: Op::Mul,
inputs: vec![a.id(), hole],
shape: s.clone(),
dtype: DType::F32,
});
gw.push(Node {
op: Op::Add,
inputs: vec![ac, b.id()],
shape: s.clone(),
dtype: DType::F32,
})
};
let nc_t = NodeHandle::from_existing(g.clone(), nc);
let out = init
.scan(
&[],
&[a.clone(), b.clone()],
&nc_t,
&nc_t,
3,
ScanEmit::Final,
)
.expect("scan");
(g, out.id(), init, a, b)
};
let fwd = |init_v: f32, a_v: f32, b_v: f32| -> f32 {
let (g, out_id, _i, _a, _b) = build(init_v, a_v, b_v);
let scan_id = {
let gr = g.read().unwrap();
gr.node(out_id).inputs[0]
};
let carry = {
let mut gw = g.write().unwrap();
fuel_graph::scan::unroll_scan(&mut gw, scan_id, 3)
.expect("unroll")
.0[0] }; crate::pipelined_bridge::realize_one_as::<f32>(&g, carry, &dev).expect("realize")[0]
};
let (g, out_id, init, a, b) = build(1.0, 0.5, 0.1);
let out = fuel_graph::NodeHandle::from_existing(g.clone(), out_id);
let grads = out.backward();
let realize_grad = |t: &fuel_graph::NodeHandle| {
crate::pipelined_bridge::realize_one_as::<f32>(
&g,
grads.get(t).expect("grad").id(),
&dev,
)
.expect("realize grad")[0]
};
let g_init = realize_grad(&init);
let g_a = realize_grad(&a);
let g_b = realize_grad(&b);
let h = 1e-3f32;
let fd = |dinit: f32, da: f32, db: f32| {
(fwd(1.0 + dinit, 0.5 + da, 0.1 + db) - fwd(1.0 - dinit, 0.5 - da, 0.1 - db))
/ (2.0 * h)
};
assert!(
(g_init - fd(h, 0.0, 0.0)).abs() < 2e-2,
"d/dinit: autograd {g_init} vs FD {}",
fd(h, 0.0, 0.0)
);
assert!(
(g_a - fd(0.0, h, 0.0)).abs() < 2e-2,
"d/da: autograd {g_a} vs FD {}",
fd(0.0, h, 0.0)
);
assert!(
(g_b - fd(0.0, 0.0, h)).abs() < 2e-2,
"d/db: autograd {g_b} vs FD {}",
fd(0.0, 0.0, h)
);
}
#[test]
fn selective_scan_is_differentiable_backward_matches_fd() {
let dev = Device::cpu();
let fwd = |u_v: f32| -> f32 {
let u = Tensor::from_f32(vec![u_v], Shape::from_dims(&[1, 1, 1]), &dev).unwrap();
let delta = u
.const_f32_like(vec![0.5f32], Shape::from_dims(&[1, 1, 1]))
.unwrap();
let a = u
.const_f32_like(vec![-1.0f32], Shape::from_dims(&[1, 1]))
.unwrap();
let b = u
.const_f32_like(vec![3.0f32], Shape::from_dims(&[1, 1, 1]))
.unwrap();
let c = u
.const_f32_like(vec![4.0f32], Shape::from_dims(&[1, 1, 1]))
.unwrap();
u.selective_scan(&delta, &a, &b, &c, false).realize_f32()[0]
};
let u = Tensor::from_f32(vec![2.0f32], Shape::from_dims(&[1, 1, 1]), &dev).unwrap();
let delta = u
.const_f32_like(vec![0.5f32], Shape::from_dims(&[1, 1, 1]))
.unwrap();
let a = u
.const_f32_like(vec![-1.0f32], Shape::from_dims(&[1, 1]))
.unwrap();
let b = u
.const_f32_like(vec![3.0f32], Shape::from_dims(&[1, 1, 1]))
.unwrap();
let c = u
.const_f32_like(vec![4.0f32], Shape::from_dims(&[1, 1, 1]))
.unwrap();
let y = u.selective_scan(&delta, &a, &b, &c, false);
let grads = y.inner.backward();
let g_u_id = grads.get(u.graph_tensor()).expect("grad u").id();
let g_u =
crate::pipelined_bridge::realize_one_as::<f32>(&u.inner.graph().clone(), g_u_id, &dev)
.expect("realize")[0];
let h = 1e-3;
let fd = (fwd(2.0 + h) - fwd(2.0 - h)) / (2.0 * h);
assert!(
(g_u - fd).abs() < 5e-2,
"selective_scan d/du: autograd {g_u} vs FD {fd}"
);
}
#[test]
fn ssd_chunk_scan_is_differentiable_backward_matches_fd() {
let dev = Device::cpu();
let fwd = |x_v: f32| -> f32 {
let x = Tensor::from_f32(vec![x_v], Shape::from_dims(&[1, 1, 1, 1]), &dev).unwrap();
let dt = x
.const_f32_like(vec![0.5f32], Shape::from_dims(&[1, 1, 1]))
.unwrap();
let a = x
.const_f32_like(vec![-1.0f32], Shape::from_dims(&[1]))
.unwrap();
let b = x
.const_f32_like(vec![3.0f32], Shape::from_dims(&[1, 1, 1, 1]))
.unwrap();
let c = x
.const_f32_like(vec![4.0f32], Shape::from_dims(&[1, 1, 1, 1]))
.unwrap();
x.ssd_chunk_scan(&dt, &a, &b, &c, 1).realize_f32()[0]
};
let x = Tensor::from_f32(vec![2.0f32], Shape::from_dims(&[1, 1, 1, 1]), &dev).unwrap();
let dt = x
.const_f32_like(vec![0.5f32], Shape::from_dims(&[1, 1, 1]))
.unwrap();
let a = x
.const_f32_like(vec![-1.0f32], Shape::from_dims(&[1]))
.unwrap();
let b = x
.const_f32_like(vec![3.0f32], Shape::from_dims(&[1, 1, 1, 1]))
.unwrap();
let c = x
.const_f32_like(vec![4.0f32], Shape::from_dims(&[1, 1, 1, 1]))
.unwrap();
let y = x.ssd_chunk_scan(&dt, &a, &b, &c, 1);
let grads = y.inner.backward();
let g_x_id = grads.get(x.graph_tensor()).expect("grad x").id();
let g_x =
crate::pipelined_bridge::realize_one_as::<f32>(&x.inner.graph().clone(), g_x_id, &dev)
.expect("realize")[0];
let h = 1e-3;
let fd = (fwd(2.0 + h) - fwd(2.0 - h)) / (2.0 * h);
assert!(
(g_x - fd).abs() < 5e-2,
"ssd_chunk_scan d/dx: autograd {g_x} vs FD {fd}"
);
}
#[test]
fn selective_scan_seqlen2_bptt_matches_finite_difference() {
let dev = Device::cpu();
let fwd = |u_vals: &[f32]| -> f32 {
let u = Tensor::from_f32(u_vals.to_vec(), Shape::from_dims(&[1, 2, 1]), &dev).unwrap();
let delta = u
.const_f32_like(vec![0.7f32, 0.7], Shape::from_dims(&[1, 2, 1]))
.unwrap();
let a = u
.const_f32_like(vec![-0.5f32], Shape::from_dims(&[1, 1]))
.unwrap();
let b = u
.const_f32_like(vec![1.0f32, 1.0], Shape::from_dims(&[1, 2, 1]))
.unwrap();
let c = u
.const_f32_like(vec![1.0f32, 1.0], Shape::from_dims(&[1, 2, 1]))
.unwrap();
u.selective_scan(&delta, &a, &b, &c, false)
.realize_f32()
.iter()
.sum()
};
let u0 = vec![1.0f32, 2.0];
let u = Tensor::from_f32(u0.clone(), Shape::from_dims(&[1, 2, 1]), &dev).unwrap();
let delta = u
.const_f32_like(vec![0.7f32, 0.7], Shape::from_dims(&[1, 2, 1]))
.unwrap();
let a = u
.const_f32_like(vec![-0.5f32], Shape::from_dims(&[1, 1]))
.unwrap();
let b = u
.const_f32_like(vec![1.0f32, 1.0], Shape::from_dims(&[1, 2, 1]))
.unwrap();
let c = u
.const_f32_like(vec![1.0f32, 1.0], Shape::from_dims(&[1, 2, 1]))
.unwrap();
let y = u.selective_scan(&delta, &a, &b, &c, false);
let grads = y.inner.backward(); let g_u_id = grads.get(u.graph_tensor()).expect("grad u").id();
let g_u =
crate::pipelined_bridge::realize_one_as::<f32>(&u.inner.graph().clone(), g_u_id, &dev)
.expect("realize");
let h = 1e-3f32;
let mut up = u0.clone();
up[0] += h;
let mut um = u0.clone();
um[0] -= h;
let fd0 = (fwd(&up) - fwd(&um)) / (2.0 * h);
assert!(
(g_u[0] - fd0).abs() < 5e-2,
"selective_scan seqlen2 dL/du_1: autograd {} vs FD {fd0}",
g_u[0]
);
assert!(
g_u[0].abs() > 0.5,
"recurrent-gate grad must be non-trivial, got {g_u:?}"
);
}
#[test]
fn ssd_chunk_scan_seqlen4_chunk2_bptt_matches_finite_difference() {
let dev = Device::cpu();
let fwd = |x_vals: &[f32]| -> f32 {
let x =
Tensor::from_f32(x_vals.to_vec(), Shape::from_dims(&[1, 4, 1, 1]), &dev).unwrap();
let dt = x
.const_f32_like(vec![0.6f32, 0.6, 0.6, 0.6], Shape::from_dims(&[1, 4, 1]))
.unwrap();
let a = x
.const_f32_like(vec![-0.5f32], Shape::from_dims(&[1]))
.unwrap();
let b = x
.const_f32_like(vec![1.0f32, 1.0, 1.0, 1.0], Shape::from_dims(&[1, 4, 1, 1]))
.unwrap();
let c = x
.const_f32_like(vec![1.0f32, 1.0, 1.0, 1.0], Shape::from_dims(&[1, 4, 1, 1]))
.unwrap();
x.ssd_chunk_scan(&dt, &a, &b, &c, 2)
.realize_f32()
.iter()
.sum()
};
let x0 = vec![1.0f32, 2.0, 3.0, 4.0];
let x = Tensor::from_f32(x0.clone(), Shape::from_dims(&[1, 4, 1, 1]), &dev).unwrap();
let dt = x
.const_f32_like(vec![0.6f32, 0.6, 0.6, 0.6], Shape::from_dims(&[1, 4, 1]))
.unwrap();
let a = x
.const_f32_like(vec![-0.5f32], Shape::from_dims(&[1]))
.unwrap();
let b = x
.const_f32_like(vec![1.0f32, 1.0, 1.0, 1.0], Shape::from_dims(&[1, 4, 1, 1]))
.unwrap();
let c = x
.const_f32_like(vec![1.0f32, 1.0, 1.0, 1.0], Shape::from_dims(&[1, 4, 1, 1]))
.unwrap();
let y = x.ssd_chunk_scan(&dt, &a, &b, &c, 2);
let grads = y.inner.backward();
let g_x_id = grads.get(x.graph_tensor()).expect("grad x").id();
let g_x =
crate::pipelined_bridge::realize_one_as::<f32>(&x.inner.graph().clone(), g_x_id, &dev)
.expect("realize");
let h = 1e-3f32;
let mut xp = x0.clone();
xp[0] += h;
let mut xm = x0.clone();
xm[0] -= h;
let fd0 = (fwd(&xp) - fwd(&xm)) / (2.0 * h);
assert!(
(g_x[0] - fd0).abs() < 5e-2,
"ssd_chunk_scan seqlen4 dL/dx_0: autograd {} vs FD {fd0}",
g_x[0]
);
assert!(
g_x[0].abs() > 0.5,
"cross-chunk-gate grad must be non-trivial, got {g_x:?}"
);
}
fn find_scan_terminal(
graph: &std::sync::Arc<std::sync::RwLock<fuel_graph::Graph>>,
root: fuel_graph::NodeId,
fused_id: fuel_graph::registry::FusedOpId,
) -> fuel_graph::NodeId {
use fuel_graph::Op;
let g = graph.read().unwrap();
let mut stack = vec![root];
let mut seen = std::collections::HashSet::new();
let mut scan_id = None;
while let Some(nid) = stack.pop() {
if !seen.insert(nid) {
continue;
}
let node = g.node(nid);
assert!(
!matches!(node.op, Op::Fused(fid, _) if fid == fused_id),
"SSM op must lower to Op::Scan, not remain fused",
);
if matches!(node.op, Op::Scan { .. }) {
scan_id = Some(nid);
}
for &inp in &node.inputs {
stack.push(inp);
}
}
scan_id.expect("an Op::Scan terminal must be present after lowering")
}
fn scan_ys_to_fused_layout(ys: &[f32], seqlen: usize, batch: usize, tail: usize) -> Vec<f32> {
assert_eq!(ys.len(), seqlen * batch * tail, "ys layout mismatch");
let mut out = vec![0f32; ys.len()];
for s in 0..seqlen {
for bb in 0..batch {
for t in 0..tail {
out[(bb * seqlen + s) * tail + t] = ys[(s * batch + bb) * tail + t];
}
}
}
out
}
#[test]
fn selective_scan_multidim_unroll_matches_fused_kernel() {
use fuel_graph::registry::FusedOps;
let dev = Device::cpu();
let (batch, seqlen, dim, dstate) = (2usize, 3usize, 2usize, 2usize);
let u = Tensor::from_f32(
vec![
0.5, -0.3, 1.0, 0.2, -0.5, 0.7, 0.25, 0.9, -0.8, 0.4, 0.6, -0.1,
],
Shape::from_dims(&[batch, seqlen, dim]),
&dev,
)
.unwrap();
let delta = u
.const_f32_like(
vec![
0.1, 0.2, 0.3, 0.15, 0.25, 0.4, 0.35, 0.05, 0.2, 0.3, 0.12, 0.28,
],
Shape::from_dims(&[batch, seqlen, dim]),
)
.unwrap();
let a = u
.const_f32_like(
vec![-0.7, -0.3, -0.5, -0.9],
Shape::from_dims(&[dim, dstate]),
)
.unwrap();
let b = u
.const_f32_like(
vec![
1.0, 0.5, 0.25, 0.75, 0.6, 0.2, 0.8, 0.4, 0.3, 0.9, 0.55, 0.15,
],
Shape::from_dims(&[batch, seqlen, dstate]),
)
.unwrap();
let c = u
.const_f32_like(
vec![2.0, 1.0, 0.5, 1.5, 0.8, 0.3, 1.2, 0.6, 0.9, 0.4, 0.7, 1.1],
Shape::from_dims(&[batch, seqlen, dstate]),
)
.unwrap();
let y = u.selective_scan(&delta, &a, &b, &c, false);
let fused = y.realize_f32();
assert_eq!(fused.len(), batch * seqlen * dim);
let graph = y.inner.graph().clone();
let roots = fuel_graph::opt::RuleRegistry::lowering_only()
.optimize_to_fixpoint(&graph, &[y.inner.id()]);
assert_eq!(roots.len(), 1);
let scan_id = find_scan_terminal(&graph, roots[0], FusedOps::SELECTIVE_SCAN);
let ys = {
let mut g = graph.write().unwrap();
fuel_graph::scan::unroll_scan(&mut g, scan_id, seqlen)
.expect("unroll")
.0[0] };
let ys_flat = crate::pipelined_bridge::realize_one_as::<f32>(&graph, ys, &dev)
.expect("realize selective_scan unroll oracle on CPU");
let oracle = scan_ys_to_fused_layout(&ys_flat, seqlen, batch, dim);
assert_eq!(oracle.len(), fused.len());
let mut max_drift = 0f32;
for (i, (o, f)) in oracle.iter().zip(fused.iter()).enumerate() {
let d = (o - f).abs();
max_drift = max_drift.max(d);
assert!(
d < 1e-3,
"selective_scan multidim mismatch at {i}: oracle {o} vs fused {f} (drift {d})",
);
}
eprintln!("selective_scan multidim: max honest drift = {max_drift:e}");
}
#[test]
fn selective_scan_softplus_multitoken_unroll_matches_fused_kernel() {
use fuel_graph::registry::FusedOps;
let dev = Device::cpu();
let (batch, seqlen, dim, dstate) = (1usize, 3usize, 2usize, 2usize);
let u = Tensor::from_f32(
vec![0.5, -0.3, 1.0, 0.2, -0.5, 0.7],
Shape::from_dims(&[batch, seqlen, dim]),
&dev,
)
.unwrap();
let delta = u
.const_f32_like(
vec![-0.5, 0.2, 0.8, -1.0, 0.3, -0.2],
Shape::from_dims(&[batch, seqlen, dim]),
)
.unwrap();
let a = u
.const_f32_like(
vec![-0.7, -0.3, -0.5, -0.9],
Shape::from_dims(&[dim, dstate]),
)
.unwrap();
let b = u
.const_f32_like(
vec![1.0, 0.5, 0.25, 0.75, 0.6, 0.2],
Shape::from_dims(&[batch, seqlen, dstate]),
)
.unwrap();
let c = u
.const_f32_like(
vec![2.0, 1.0, 0.5, 1.5, 0.8, 0.3],
Shape::from_dims(&[batch, seqlen, dstate]),
)
.unwrap();
let y = u.selective_scan(&delta, &a, &b, &c, true);
let fused = y.realize_f32();
assert_eq!(fused.len(), batch * seqlen * dim);
let graph = y.inner.graph().clone();
let roots = fuel_graph::opt::RuleRegistry::lowering_only()
.optimize_to_fixpoint(&graph, &[y.inner.id()]);
let scan_id = find_scan_terminal(&graph, roots[0], FusedOps::SELECTIVE_SCAN);
let ys = {
let mut g = graph.write().unwrap();
fuel_graph::scan::unroll_scan(&mut g, scan_id, seqlen)
.expect("unroll")
.0[0] };
let ys_flat = crate::pipelined_bridge::realize_one_as::<f32>(&graph, ys, &dev)
.expect("realize selective_scan softplus oracle on CPU");
let oracle = scan_ys_to_fused_layout(&ys_flat, seqlen, batch, dim);
let mut max_drift = 0f32;
for (i, (o, f)) in oracle.iter().zip(fused.iter()).enumerate() {
let d = (o - f).abs();
max_drift = max_drift.max(d);
assert!(
d < 1e-3,
"selective_scan softplus mismatch at {i}: oracle {o} vs fused {f} (drift {d})",
);
}
eprintln!("selective_scan softplus: max honest drift = {max_drift:e}");
}
#[test]
fn ssd_chunk_scan_multidim_unroll_matches_fused_kernel() {
use fuel_graph::registry::FusedOps;
let dev = Device::cpu();
let (batch, seqlen, heads, head_dim, state_dim, chunk) =
(2usize, 4usize, 2usize, 2usize, 2usize, 2usize);
let x_data: Vec<f32> = (0..batch * seqlen * heads * head_dim)
.map(|i| 0.20 + 0.05 * i as f32)
.collect();
let dt_data: Vec<f32> = (0..batch * seqlen * heads)
.map(|i| 0.10 + 0.04 * i as f32)
.collect();
let b_data: Vec<f32> = (0..batch * seqlen * heads * state_dim)
.map(|i| 0.15 + 0.03 * i as f32)
.collect();
let c_data: Vec<f32> = (0..batch * seqlen * heads * state_dim)
.map(|i| 0.25 + 0.035 * i as f32)
.collect();
let x = Tensor::from_f32(
x_data,
Shape::from_dims(&[batch, seqlen, heads, head_dim]),
&dev,
)
.unwrap();
let dt = x
.const_f32_like(dt_data, Shape::from_dims(&[batch, seqlen, heads]))
.unwrap();
let a = x
.const_f32_like(vec![-0.6f32, -1.0], Shape::from_dims(&[heads]))
.unwrap();
let b = x
.const_f32_like(b_data, Shape::from_dims(&[batch, seqlen, heads, state_dim]))
.unwrap();
let c = x
.const_f32_like(c_data, Shape::from_dims(&[batch, seqlen, heads, state_dim]))
.unwrap();
let y = x.ssd_chunk_scan(&dt, &a, &b, &c, chunk);
let fused = y.realize_f32();
assert_eq!(fused.len(), batch * seqlen * heads * head_dim);
let graph = y.inner.graph().clone();
let roots = fuel_graph::opt::RuleRegistry::lowering_only()
.optimize_to_fixpoint(&graph, &[y.inner.id()]);
assert_eq!(roots.len(), 1);
let scan_id = find_scan_terminal(&graph, roots[0], FusedOps::SSD_CHUNK_SCAN);
let ys = {
let mut g = graph.write().unwrap();
fuel_graph::scan::unroll_scan(&mut g, scan_id, seqlen)
.expect("unroll")
.0[0] };
let ys_flat = crate::pipelined_bridge::realize_one_as::<f32>(&graph, ys, &dev)
.expect("realize ssd_chunk_scan unroll oracle on CPU");
let oracle = scan_ys_to_fused_layout(&ys_flat, seqlen, batch, heads * head_dim);
assert_eq!(oracle.len(), fused.len());
let mut max_drift = 0f32;
for (i, (o, f)) in oracle.iter().zip(fused.iter()).enumerate() {
let d = (o - f).abs();
max_drift = max_drift.max(d);
assert!(
d < 1e-3,
"ssd_chunk_scan multidim mismatch at {i}: oracle {o} vs fused {f} (drift {d})",
);
}
eprintln!("ssd_chunk_scan multidim: max honest drift = {max_drift:e}");
}
#[test]
fn nf4_matmul_decompose_matches_kernel() {
use fuel_graph::Op;
use fuel_graph::registry::FusedOps;
let dev = Device::cpu();
let weight = crate::nf4::nf4_from_bytes(
vec![247_u8, 247, 127, 127],
vec![1.0_f32, 2.0, 10.0, 20.0],
2,
4,
2,
&dev,
)
.expect("nf4_from_bytes");
let act = Tensor::from_graph_tensor(
weight
.w_packed
.graph_tensor()
.const_f32_like(vec![1.0_f32, 2.0, 2.0, 4.0], Shape::from_dims(&[1, 4]))
.unwrap(),
);
let y = weight.matmul(&act);
let graph = y.inner.graph().clone();
let id = y.inner.id();
let roots =
fuel_graph::opt::RuleRegistry::lowering_only().optimize_to_fixpoint(&graph, &[id]);
assert_eq!(roots.len(), 1, "lowering should keep a single root");
{
let g = graph.read().unwrap();
let mut stack = vec![roots[0]];
let mut seen = std::collections::HashSet::new();
while let Some(nid) = stack.pop() {
if !seen.insert(nid) {
continue;
}
let node = g.node(nid);
assert!(
!matches!(node.op, Op::Fused(fid, _) if fid == FusedOps::NF4_MATMUL),
"decomposed graph still contains an Op::Fused(NF4_MATMUL) island",
);
for &inp in &node.inputs {
stack.push(inp);
}
}
}
let got = crate::pipelined_bridge::realize_one_as::<f32>(&graph, roots[0], &dev)
.expect("realize decomposed Nf4Matmul on CPU");
assert_eq!(got.len(), 2);
assert!((got[0] - 10.0).abs() < 1e-4, "out 0: {}", got[0]);
assert!((got[1] - 50.0).abs() < 1e-4, "out 1: {}", got[1]);
}
fn lower_realize_fused(
anchor: &Tensor,
op: fuel_graph::Op,
inputs: Vec<fuel_graph::NodeId>,
shape: Shape,
) -> Vec<f32> {
let graph = anchor.inner.graph().clone();
let fused_id = {
let mut g = graph.write().unwrap();
g.push(fuel_graph::Node {
op,
inputs,
shape,
dtype: DType::F32,
})
};
let roots = fuel_graph::opt::RuleRegistry::lowering_only()
.optimize_to_fixpoint(&graph, &[fused_id]);
assert_eq!(roots.len(), 1, "lowering keeps a single root");
crate::pipelined_bridge::realize_one_as::<f32>(&graph, roots[0], &Device::cpu())
.expect("realize decomposed fused op on CPU")
}
#[test]
fn powi_backward_decompose_matches_reference() {
use fuel_graph::registry::{FusedOpParams, FusedOps};
let dev = Device::cpu();
let exp = 3i32;
let x_data = vec![1.5f32, -2.0, 0.5, 3.0];
let up_data = vec![1.0f32, 0.5, 2.0, -1.0];
let shape = Shape::from_dims(&[4]);
let x = Tensor::from_f32(x_data.clone(), shape.clone(), &dev).unwrap();
let up = x.const_f32_like(up_data.clone(), shape.clone()).unwrap();
let got = lower_realize_fused(
&x,
fuel_graph::Op::Fused(FusedOps::POWI_BACKWARD, FusedOpParams::PowIBackward { exp }),
vec![x.inner.id(), up.inner.id()],
shape,
);
let expected: Vec<f32> = x_data
.iter()
.zip(&up_data)
.map(|(&x, &u)| exp as f32 * x.powi(exp - 1) * u)
.collect();
assert_eq!(got.len(), expected.len());
for (i, (&g, &e)) in got.iter().zip(&expected).enumerate() {
assert!(
(g - e).abs() < 1e-4,
"powi_backward at {i}: got {g}, expected {e}"
);
}
}
#[test]
fn inplace_affine_decompose_matches_affine_reference() {
use fuel_graph::registry::{FusedOpParams, FusedOps};
let dev = Device::cpu();
let (mul, add) = (2.5f64, -1.0f64);
let x_data = vec![1.5f32, -2.0, 0.5, 3.0];
let shape = Shape::from_dims(&[4]);
let x = Tensor::from_f32(x_data.clone(), shape.clone(), &dev).unwrap();
let got = lower_realize_fused(
&x,
fuel_graph::Op::Fused(
FusedOps::INPLACE_AFFINE,
FusedOpParams::InplaceAffine { mul, add },
),
vec![x.inner.id()],
shape,
);
let expected: Vec<f32> = x_data
.iter()
.map(|&v| (mul as f32) * v + add as f32)
.collect();
assert_eq!(got.len(), expected.len());
for (i, (&g, &e)) in got.iter().zip(&expected).enumerate() {
assert_eq!(g, e, "inplace_affine at {i}: got {g}, expected {e}");
}
}
#[test]
fn softmax_last_dim_backward_decompose_matches_reference() {
use fuel_graph::registry::{FusedOpParams, FusedOps};
let dev = Device::cpu();
let (rows, cols) = (2usize, 3usize);
let logits = [[0.5f32, -1.0, 2.0], [1.0, 0.0, -0.5]];
let mut s_data = Vec::new();
for row in &logits {
let m = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let exps: Vec<f32> = row.iter().map(|&v| (v - m).exp()).collect();
let sum: f32 = exps.iter().sum();
s_data.extend(exps.iter().map(|&e| e / sum));
}
let g_data = vec![1.0f32, -0.5, 0.3, 0.2, 1.5, -1.0];
let shape = Shape::from_dims(&[rows, cols]);
let s = Tensor::from_f32(s_data.clone(), shape.clone(), &dev).unwrap();
let g = s.const_f32_like(g_data.clone(), shape.clone()).unwrap();
let got = lower_realize_fused(
&s,
fuel_graph::Op::Fused(
FusedOps::SOFTMAX_LAST_DIM_BACKWARD,
FusedOpParams::SoftmaxLastDimBackward,
),
vec![s.inner.id(), g.inner.id()],
shape,
);
let mut expected = vec![0.0f32; rows * cols];
for r in 0..rows {
let dot: f32 = (0..cols)
.map(|c| g_data[r * cols + c] * s_data[r * cols + c])
.sum();
for c in 0..cols {
expected[r * cols + c] = s_data[r * cols + c] * (g_data[r * cols + c] - dot);
}
}
for (i, (&gv, &ev)) in got.iter().zip(&expected).enumerate() {
assert!(
(gv - ev).abs() < 1e-5,
"softmax_bwd at {i}: got {gv}, expected {ev}"
);
}
}
#[test]
fn rms_norm_last_dim_backward_decompose_matches_reference() {
use fuel_graph::registry::{FusedOpParams, FusedOps};
let dev = Device::cpu();
let (rows, cols) = (2usize, 3usize);
let eps = 1e-5f64;
let x_data = vec![1.0f32, -2.0, 0.5, 3.0, 0.25, -1.5];
let g_data = vec![0.5f32, 1.0, -0.3, 0.2, -1.0, 0.7];
let shape = Shape::from_dims(&[rows, cols]);
let x = Tensor::from_f32(x_data.clone(), shape.clone(), &dev).unwrap();
let g = x.const_f32_like(g_data.clone(), shape.clone()).unwrap();
let got = lower_realize_fused(
&x,
fuel_graph::Op::Fused(
FusedOps::RMS_NORM_LAST_DIM_BACKWARD,
FusedOpParams::RmsNormLastDimBackward { eps },
),
vec![x.inner.id(), g.inner.id()],
shape,
);
let n = cols as f32;
let mut expected = vec![0.0f32; rows * cols];
for r in 0..rows {
let meansq: f32 = (0..cols).map(|c| x_data[r * cols + c].powi(2)).sum::<f32>() / n;
let denom = meansq + eps as f32;
let rrms = 1.0 / denom.sqrt();
let s: f32 = (0..cols)
.map(|c| g_data[r * cols + c] * x_data[r * cols + c])
.sum();
for c in 0..cols {
let term = x_data[r * cols + c] * s / (n * denom);
expected[r * cols + c] = rrms * (g_data[r * cols + c] - term);
}
}
for (i, (&gv, &ev)) in got.iter().zip(&expected).enumerate() {
assert!(
(gv - ev).abs() < 1e-4,
"rms_norm_bwd at {i}: got {gv}, expected {ev}"
);
}
}
#[test]
fn layer_norm_last_dim_backward_decompose_matches_reference() {
use fuel_graph::registry::{FusedOpParams, FusedOps};
let dev = Device::cpu();
let (rows, cols) = (2usize, 3usize);
let eps = 1e-5f64;
let x_data = vec![1.0f32, -2.0, 0.5, 3.0, 0.25, -1.5];
let g_data = vec![0.5f32, 1.0, -0.3, 0.2, -1.0, 0.7];
let shape = Shape::from_dims(&[rows, cols]);
let x = Tensor::from_f32(x_data.clone(), shape.clone(), &dev).unwrap();
let g = x.const_f32_like(g_data.clone(), shape.clone()).unwrap();
let got = lower_realize_fused(
&x,
fuel_graph::Op::Fused(
FusedOps::LAYER_NORM_LAST_DIM_BACKWARD,
FusedOpParams::LayerNormLastDimBackward { eps },
),
vec![x.inner.id(), g.inner.id()],
shape,
);
let n = cols as f32;
let mut expected = vec![0.0f32; rows * cols];
for r in 0..rows {
let mean_x: f32 = (0..cols).map(|c| x_data[r * cols + c]).sum::<f32>() / n;
let var: f32 = (0..cols)
.map(|c| (x_data[r * cols + c] - mean_x).powi(2))
.sum::<f32>()
/ n;
let istd = 1.0 / (var + eps as f32).sqrt();
let xhat: Vec<f32> = (0..cols)
.map(|c| (x_data[r * cols + c] - mean_x) * istd)
.collect();
let mean_g: f32 = (0..cols).map(|c| g_data[r * cols + c]).sum::<f32>() / n;
let mean_gxh: f32 = (0..cols)
.map(|c| g_data[r * cols + c] * xhat[c])
.sum::<f32>()
/ n;
for c in 0..cols {
expected[r * cols + c] =
istd * (g_data[r * cols + c] - mean_g - xhat[c] * mean_gxh);
}
}
for (i, (&gv, &ev)) in got.iter().zip(&expected).enumerate() {
assert!(
(gv - ev).abs() < 1e-4,
"layer_norm_bwd at {i}: got {gv}, expected {ev}"
);
}
}
#[test]
fn reduce_max_to_backward_decompose_matches_reference() {
use fuel_graph::registry::{FusedOpParams, FusedOps};
let dev = Device::cpu();
let x_data = vec![1.0f32, 3.0, 3.0, 2.0, 0.0, 1.0]; let up_data = vec![10.0f32, 5.0]; let x_shape = Shape::from_dims(&[2, 3]);
let up_shape = Shape::from_dims(&[2, 1]);
let x = Tensor::from_f32(x_data.clone(), x_shape.clone(), &dev).unwrap();
let up = x.const_f32_like(up_data.clone(), up_shape).unwrap();
let got = lower_realize_fused(
&x,
fuel_graph::Op::Fused(
FusedOps::REDUCE_MAX_TO_BACKWARD,
FusedOpParams::ReduceMaxToBackward,
),
vec![x.inner.id(), up.inner.id()],
x_shape,
);
let expected = vec![0.0f32, 5.0, 5.0, 5.0, 0.0, 0.0];
for (i, (&gv, &ev)) in got.iter().zip(&expected).enumerate() {
assert!(
(gv - ev).abs() < 1e-5,
"reduce_max_bwd at {i}: got {gv}, expected {ev}"
);
}
}
fn causal_conv1d_check(use_silu: bool) {
use fuel_graph::registry::{FusedOpParams, FusedOps};
let dev = Device::cpu();
let (b, c, k, seq) = (1usize, 2usize, 3usize, 2usize);
let x_seq = seq + (k - 1); let x_data = vec![0.0f32, 1.0, 2.0, 3.0, 1.0, 0.0, -1.0, -2.0];
let w_data = vec![0.5f32, 1.0, -0.5, 1.0, 0.0, 2.0];
let bias_data = vec![0.1f32, -0.2];
let x = Tensor::from_f32(x_data.clone(), Shape::from_dims(&[b, c, x_seq]), &dev).unwrap();
let w = x
.const_f32_like(w_data.clone(), Shape::from_dims(&[c, 1, k]))
.unwrap();
let bias = x
.const_f32_like(bias_data.clone(), Shape::from_dims(&[c]))
.unwrap();
let got = lower_realize_fused(
&x,
fuel_graph::Op::Fused(
FusedOps::CAUSAL_CONV1D,
FusedOpParams::CausalConv1d { use_silu },
),
vec![x.inner.id(), w.inner.id(), bias.inner.id()],
Shape::from_dims(&[b, c, seq]),
);
let mut expected = vec![0.0f32; b * c * seq];
for ch in 0..c {
for t in 0..seq {
let mut acc = 0.0f32;
for kk in 0..k {
acc += w_data[ch * k + kk] * x_data[ch * x_seq + t + kk];
}
acc += bias_data[ch];
if use_silu {
acc *= 1.0 / (1.0 + (-acc).exp());
}
expected[ch * seq + t] = acc;
}
}
for (i, (&gv, &ev)) in got.iter().zip(&expected).enumerate() {
assert!(
(gv - ev).abs() < 1e-4,
"causal_conv1d (silu={use_silu}) at {i}: got {gv}, expected {ev}",
);
}
}
#[test]
fn causal_conv1d_decompose_matches_reference() {
causal_conv1d_check(false);
causal_conv1d_check(true);
}
#[test]
fn flash_attn_backward_decompose_matches_reference() {
use fuel_graph::registry::{FusedOpParams, FusedOps};
let dev = Device::cpu();
let (sq, sk, d) = (2usize, 2usize, 2usize);
#[allow(clippy::approx_constant)]
let scale = 0.7071f32;
let q_data = vec![0.1f32, -0.2, 0.3, 0.5]; let k_data = vec![0.4f32, 0.1, -0.3, 0.2]; let v_data = vec![1.0f32, 2.0, -1.0, 0.5]; let do_data = vec![0.5f32, -1.0, 0.2, 0.3]; let qshape = Shape::from_dims(&[1, 1, sq, d]);
let kshape = Shape::from_dims(&[1, 1, sk, d]);
let q = Tensor::from_f32(q_data.clone(), qshape.clone(), &dev).unwrap();
let k = q.const_f32_like(k_data.clone(), kshape.clone()).unwrap();
let v = q.const_f32_like(v_data.clone(), kshape.clone()).unwrap();
let dout = q.const_f32_like(do_data.clone(), qshape.clone()).unwrap();
let params = FusedOpParams::FlashAttnBackward {
softmax_scale: scale,
causal: false,
window_size_left: None,
window_size_right: None,
softcap: None,
};
let inputs = vec![q.inner.id(), k.inner.id(), v.inner.id(), dout.inner.id()];
let dq = lower_realize_fused(
&q,
fuel_graph::Op::Fused(FusedOps::FLASH_ATTN_BACKWARD_Q, params.clone()),
inputs.clone(),
qshape.clone(),
);
let dk = lower_realize_fused(
&q,
fuel_graph::Op::Fused(FusedOps::FLASH_ATTN_BACKWARD_K, params.clone()),
inputs.clone(),
kshape.clone(),
);
let dv = lower_realize_fused(
&q,
fuel_graph::Op::Fused(FusedOps::FLASH_ATTN_BACKWARD_V, params),
inputs,
kshape,
);
let mut p = vec![0.0f32; sq * sk];
for i in 0..sq {
let mut scores = vec![0.0f32; sk];
for j in 0..sk {
let mut s = 0.0f32;
for l in 0..d {
s += q_data[i * d + l] * k_data[j * d + l];
}
scores[j] = scale * s;
}
let m = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let mut sum = 0.0f32;
let e: Vec<f32> = scores
.iter()
.map(|&s| {
let x = (s - m).exp();
sum += x;
x
})
.collect();
for j in 0..sk {
p[i * sk + j] = e[j] / sum;
}
}
let mut ref_dv = vec![0.0f32; sk * d];
for j in 0..sk {
for l in 0..d {
ref_dv[j * d + l] = (0..sq).map(|i| p[i * sk + j] * do_data[i * d + l]).sum();
}
}
let mut dp = vec![0.0f32; sq * sk];
for i in 0..sq {
for j in 0..sk {
dp[i * sk + j] = (0..d).map(|l| do_data[i * d + l] * v_data[j * d + l]).sum();
}
}
let mut dsr = vec![0.0f32; sq * sk];
for i in 0..sq {
let rowdot: f32 = (0..sk).map(|j| dp[i * sk + j] * p[i * sk + j]).sum();
for j in 0..sk {
dsr[i * sk + j] = scale * p[i * sk + j] * (dp[i * sk + j] - rowdot);
}
}
let mut ref_dq = vec![0.0f32; sq * d];
for i in 0..sq {
for l in 0..d {
ref_dq[i * d + l] = (0..sk).map(|j| dsr[i * sk + j] * k_data[j * d + l]).sum();
}
}
let mut ref_dk = vec![0.0f32; sk * d];
for j in 0..sk {
for l in 0..d {
ref_dk[j * d + l] = (0..sq).map(|i| dsr[i * sk + j] * q_data[i * d + l]).sum();
}
}
let check = |name: &str, got: &[f32], exp: &[f32]| {
assert_eq!(got.len(), exp.len(), "{name} length");
for (i, (&gv, &ev)) in got.iter().zip(exp).enumerate() {
assert!(
(gv - ev).abs() < 1e-4,
"{name} at {i}: got {gv}, expected {ev}"
);
}
};
check("dQ", &dq, &ref_dq);
check("dK", &dk, &ref_dk);
check("dV", &dv, &ref_dv);
}
#[test]
fn paged_attn_decompose_matches_reference() {
use fuel_graph::registry::{FusedOpParams, FusedOps};
let dev = Device::cpu();
let block_size = 2usize;
let scale = 1.0f32;
let kc = vec![
1.0f32, 0.0, 0.0, 1.0, 9.0, 9.0, 9.0, 9.0, 1.0, 1.0, 2.0, 0.0,
];
let vc = vec![
1.0f32, 2.0, 3.0, 4.0, 0.0, 0.0, 0.0, 0.0, 5.0, 6.0, 7.0, 8.0,
];
let q_data = vec![0.5f32, 0.5]; let bt = vec![0u32, 2u32]; let cl = vec![3u32];
let q = Tensor::from_f32(q_data.clone(), Shape::from_dims(&[1, 1, 1, 2]), &dev).unwrap();
let kcache = q
.const_f32_like(kc.clone(), Shape::from_dims(&[3, 2, 1, 2]))
.unwrap();
let vcache = q
.const_f32_like(vc.clone(), Shape::from_dims(&[3, 2, 1, 2]))
.unwrap();
let block_table = q.const_u32_like(bt, Shape::from_dims(&[1, 2])).unwrap();
let context_lens = q.const_u32_like(cl, Shape::from_dims(&[1])).unwrap();
let params = FusedOpParams::PagedAttn {
softmax_scale: scale,
block_size,
softcap: None,
};
let inputs = vec![
q.inner.id(),
kcache.inner.id(),
vcache.inner.id(),
block_table.inner.id(),
context_lens.inner.id(),
];
let got = lower_realize_fused(
&q,
fuel_graph::Op::Fused(FusedOps::PAGED_ATTN, params),
inputs,
Shape::from_dims(&[1, 1, 1, 2]),
);
let (kv_len, ctx) = (4usize, 3usize);
let mut k_seq = vec![];
let mut v_seq = vec![];
for &blk in &[0usize, 2usize] {
for p in 0..block_size {
let base = (blk * block_size + p) * 2; k_seq.push([kc[base], kc[base + 1]]);
v_seq.push([vc[base], vc[base + 1]]);
}
}
let mut scores = vec![0.0f32; kv_len];
for j in 0..kv_len {
let s = q_data[0] * k_seq[j][0] + q_data[1] * k_seq[j][1];
scores[j] = if j >= ctx {
f32::NEG_INFINITY
} else {
scale * s
};
}
let m = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let mut sum = 0.0f32;
let e: Vec<f32> = scores
.iter()
.map(|&s| {
let x = (s - m).exp();
sum += x;
x
})
.collect();
let mut expected = vec![0.0f32; 2];
for j in 0..kv_len {
for l in 0..2 {
expected[l] += (e[j] / sum) * v_seq[j][l];
}
}
for (i, (&gv, &ev)) in got.iter().zip(&expected).enumerate() {
assert!(
(gv - ev).abs() < 1e-4,
"paged_attn at {i}: got {gv}, expected {ev}"
);
}
}
#[test]
#[cfg(feature = "cuda")]
fn cuda_executor_matches_cpu_on_add_mul() {
let a =
Tensor::from_f32(vec![1.0, 2.0, 3.0], Shape::from_dims(&[3]), &Device::cpu()).unwrap();
let b = a
.const_f32_like(vec![4.0, 5.0, 6.0], Shape::from_dims(&[3]))
.unwrap();
let c = a.add(&b).unwrap().mul(&a).unwrap();
let cpu_result = c.realize_f32();
let executor = fuel_cuda_backend::CudaDevice::new(0).unwrap();
let cuda_result = c.realize_f32_cuda(&executor);
assert_eq!(cpu_result, cuda_result);
}
#[test]
#[cfg(feature = "cuda")]
#[ignore = "requires a live CUDA device"]
fn maximum_minimum_nan_convention_lazy_realize_smoke() {
let cuda = match fuel_cuda_backend::CudaDevice::new(0) {
Ok(d) => d,
Err(e) => {
return fuel_test_support::hardware::skip(
fuel_test_support::hardware::Hardware::Cuda,
fuel_test_support::hardware::Missing::device(format!(
"CudaDevice::new(0): {e:?}"
)),
);
}
};
let a = Tensor::from_f32(
vec![f32::NAN, -2.0, 3.0, f32::NAN],
Shape::from_dims(&[4]),
&Device::cpu(),
)
.unwrap();
let b = a
.const_f32_like(vec![1.0, f32::NAN, 2.0, f32::NAN], Shape::from_dims(&[4]))
.unwrap();
let max_lazy = a.maximum(&b).unwrap();
let cpu_max = max_lazy.realize_f32();
let cuda_max = max_lazy.realize_f32_cuda(&cuda);
for i in 0..4 {
assert_eq!(
cpu_max[i].is_nan(),
cuda_max[i].is_nan(),
"maximum[{i}] NaN-ness must match: cpu={}, cuda={}",
cpu_max[i],
cuda_max[i],
);
if !cpu_max[i].is_nan() {
assert_eq!(cpu_max[i], cuda_max[i], "maximum[{i}]");
}
}
let min_lazy = a.minimum(&b).unwrap();
let cpu_min = min_lazy.realize_f32();
let cuda_min = min_lazy.realize_f32_cuda(&cuda);
for i in 0..4 {
assert_eq!(
cpu_min[i].is_nan(),
cuda_min[i].is_nan(),
"minimum[{i}] NaN-ness must match: cpu={}, cuda={}",
cpu_min[i],
cuda_min[i],
);
if !cpu_min[i].is_nan() {
assert_eq!(cpu_min[i], cuda_min[i], "minimum[{i}]");
}
}
}
#[test]
#[cfg(feature = "cuda")]
#[ignore = "requires a live CUDA device"]
fn relu_nan_convention_lazy_realize_smoke() {
let cuda = match fuel_cuda_backend::CudaDevice::new(0) {
Ok(d) => d,
Err(e) => {
return fuel_test_support::hardware::skip(
fuel_test_support::hardware::Hardware::Cuda,
fuel_test_support::hardware::Missing::device(format!(
"CudaDevice::new(0): {e:?}"
)),
);
}
};
let a = Tensor::from_f32(
vec![f32::NAN, -2.0, 3.0],
Shape::from_dims(&[3]),
&Device::cpu(),
)
.unwrap();
let relu_lazy = a.relu();
let cpu_result = relu_lazy.realize_f32();
let cuda_result = relu_lazy.realize_f32_cuda(&cuda);
for i in 0..3 {
assert_eq!(
cpu_result[i].is_nan(),
cuda_result[i].is_nan(),
"relu[{i}] NaN-ness must match: cpu={}, cuda={}",
cpu_result[i],
cuda_result[i],
);
if !cpu_result[i].is_nan() {
assert_eq!(cpu_result[i], cuda_result[i], "relu[{i}]");
}
}
}
#[test]
#[cfg(feature = "cuda")]
fn cuda_executor_matches_cpu_on_matmul() {
let a = Tensor::from_f32(
vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
Shape::from_dims(&[2, 3]),
&Device::cpu(),
)
.unwrap();
let b = a
.const_f32_like(
vec![7.0, 8.0, 9.0, 10.0, 11.0, 12.0],
Shape::from_dims(&[3, 2]),
)
.unwrap();
let c = a.matmul(&b).unwrap();
let cpu = c.realize_f32();
let exe = fuel_cuda_backend::CudaDevice::new(0).unwrap();
let cuda = c.realize_f32_cuda(&exe);
assert_eq!(cpu.len(), cuda.len());
for (i, (a, b)) in cpu.iter().zip(cuda.iter()).enumerate() {
assert!((a - b).abs() < 1e-3, "matmul[{i}]: cpu={a}, cuda={b}",);
}
}
#[test]
#[cfg(feature = "cuda")]
fn cuda_executor_matches_cpu_on_broadcast_matmul() {
let x = Tensor::from_f32(
(0..12).map(|i| i as f32 * 0.1).collect::<Vec<_>>(),
Shape::from_dims(&[1, 3, 4]),
&Device::cpu(),
)
.unwrap();
let w = x
.const_f32_like(
(0..8).map(|i| i as f32 * 0.2).collect::<Vec<_>>(),
Shape::from_dims(&[4, 2]),
)
.unwrap();
let y = x.matmul(&w).unwrap();
let cpu = y.realize_f32();
let exe = fuel_cuda_backend::CudaDevice::new(0).unwrap();
let cuda = y.realize_f32_cuda(&exe);
assert_eq!(cpu.len(), cuda.len());
for (i, (&a, &b)) in cpu.iter().zip(cuda.iter()).enumerate() {
assert!((a - b).abs() < 1e-3, "bcast_mm[{i}]: cpu={a}, cuda={b}");
}
}
#[test]
#[cfg(feature = "cuda")]
fn cuda_executor_matches_cpu_on_permute() {
let x = Tensor::from_f32(
(0..24).map(|i| i as f32).collect::<Vec<_>>(),
Shape::from_dims(&[1, 2, 3, 4]),
&Device::cpu(),
)
.unwrap();
let y = x.permute([0, 2, 1, 3_usize]).unwrap();
let cpu = y.realize_f32();
let exe = fuel_cuda_backend::CudaDevice::new(0).unwrap();
let cuda = y.realize_f32_cuda(&exe);
assert_eq!(cpu, cuda, "permute mismatch");
}
#[test]
#[cfg(feature = "cuda")]
fn cuda_executor_matches_cpu_on_softmax() {
let x = Tensor::from_f32(
vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
Shape::from_dims(&[2, 3]),
&Device::cpu(),
)
.unwrap();
let y = x.softmax_last_dim().unwrap();
let cpu = y.realize_f32();
let exe = fuel_cuda_backend::CudaDevice::new(0).unwrap();
let cuda = y.realize_f32_cuda(&exe);
assert_eq!(cpu.len(), cuda.len());
for (i, (&a, &b)) in cpu.iter().zip(cuda.iter()).enumerate() {
assert!((a - b).abs() < 1e-4, "softmax[{i}]: cpu={a}, cuda={b}");
}
}
#[test]
#[cfg(feature = "cuda")]
fn cuda_executor_matches_cpu_on_softmax_via_lowering() {
let n = 24;
let last = 5;
let data: Vec<f32> = (0..n * last)
.map(|i| ((i as f32) * 0.13).sin() * 2.0 - 0.7)
.collect();
let x = Tensor::from_f32(data, Shape::from_dims(&[n, last]), &Device::cpu()).unwrap();
let y = x.softmax_last_dim().unwrap();
let cpu = y.realize_f32();
let graph = y.inner.graph().clone();
let registry = fuel_graph::opt::RuleRegistry::lowering_only();
let remapped = registry.optimize_to_fixpoint(&graph, &[y.inner.id()]);
let dev = fuel_cuda_backend::CudaDevice::new(0).unwrap();
let fc_device: crate::Device = dev.clone().into();
let cuda = crate::pipelined_bridge::realize_one_as::<f32>(&graph, remapped[0], &fc_device)
.expect("realize lowered softmax on CUDA");
assert_eq!(cpu.len(), cuda.len());
let mut max_abs_err = 0.0_f32;
for (i, (&a, &b)) in cpu.iter().zip(cuda.iter()).enumerate() {
let err = (a - b).abs();
if err > max_abs_err {
max_abs_err = err;
}
assert!(
err < 1e-5,
"lowered softmax[{i}]: cpu={a} (fused), cuda={b} (composed), err={err}",
);
}
eprintln!("max_abs_err over lowered-vs-fused softmax: {max_abs_err:.3e}");
}
#[test]
#[cfg(feature = "cuda")]
fn cuda_executor_matches_cpu_on_concat_slice() {
let a = Tensor::from_f32(
vec![1.0, 2.0, 3.0, 4.0],
Shape::from_dims(&[2, 2]),
&Device::cpu(),
)
.unwrap();
let b = a
.const_f32_like(vec![5.0, 6.0, 7.0, 8.0], Shape::from_dims(&[2, 2]))
.unwrap();
let cat = a.concat(&b, 1).unwrap(); let sliced = cat.slice(1, 1, 2).unwrap(); let cpu = sliced.realize_f32();
let exe = fuel_cuda_backend::CudaDevice::new(0).unwrap();
let cuda = sliced.realize_f32_cuda(&exe);
assert_eq!(cpu, cuda, "concat+slice mismatch");
}
#[test]
#[cfg(feature = "cuda")]
fn cuda_executor_matches_cpu_on_rms_norm() {
let x = Tensor::from_f32(
(0..8).map(|i| i as f32 * 0.5 - 1.5).collect::<Vec<_>>(),
Shape::from_dims(&[2, 4]),
&Device::cpu(),
)
.unwrap();
let y = x.rms_norm_last_dim(1e-5).unwrap();
let cpu = y.realize_f32();
let exe = fuel_cuda_backend::CudaDevice::new(0).unwrap();
let cuda = y.realize_f32_cuda(&exe);
assert_eq!(cpu.len(), cuda.len());
for (i, (&a, &b)) in cpu.iter().zip(cuda.iter()).enumerate() {
assert!((a - b).abs() < 1e-3, "rms_norm[{i}]: cpu={a}, cuda={b}");
}
}
#[test]
fn realize_f64_through_bridge() {
let a =
Tensor::from_f64(vec![1.5, 2.5, 3.5], Shape::from_dims(&[3]), &Device::cpu()).unwrap();
let b = a.mul(&a).unwrap();
assert_eq!(b.realize_f64(), vec![2.25, 6.25, 12.25]);
}
#[test]
fn gap327_realize_f32_on_f64_root_must_reject_not_reinterpret() {
let t =
Tensor::from_f64(vec![1.0_f64, 2.0], Shape::from_dims(&[2]), &Device::cpu()).unwrap();
let graph = t.inner.graph().clone();
let target = t.inner.id();
let device = Device::cpu();
let prev_hook = std::panic::take_hook();
std::panic::set_hook(Box::new(|_| {}));
let ctl_f64 = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| t.realize_f64()));
let ctl_f16 = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| t.realize_f16()));
let subj_method =
std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| t.realize_f32()));
std::panic::set_hook(prev_hook);
let subj_raw = crate::pipelined_bridge::realize_one_as::<f32>(&graph, target, &device);
let bytes_view = crate::pipelined_bridge::realize_one_bytes(&graph, target, &device);
assert!(
ctl_f64.is_ok(),
"control: realize_f64 on an F64 root must succeed; got {ctl_f64:?}"
);
assert!(
ctl_f16.is_err(),
"positive control: realize_f16() on an F64 root must panic (guard fires); if it \
does not, the harness is not observing panics and the subject asserts are vacuous"
);
assert!(
subj_method.is_err(),
"realize_f32() on an F64 root must REJECT (panic), not silently reinterpret bytes"
);
assert!(
matches!(subj_raw, Err(fuel_ir::Error::UnexpectedDType { .. })),
"realize_one_as::<f32> on an F64 root must return Err(UnexpectedDType), not \
Ok(reinterpreted bytes); got {subj_raw:?}"
);
assert!(
matches!(&bytes_view, Ok(b) if b.len() == 16),
"byte view realize_one_bytes on an F64 [1.0, 2.0] root must return 16 raw bytes \
(unguarded), not an error; got {bytes_view:?}"
);
}
#[test]
fn lazy_tensor_mini_llama_block_forward() {
let seq = 3;
let d_head = 4; let num_heads = 2;
let d_model = num_heads * d_head;
let x_data: Vec<f32> = (0..seq * d_model).map(|i| i as f32 * 0.01).collect();
let x =
Tensor::from_f32(x_data, Shape::from_dims(&[1, seq, d_model]), &Device::cpu()).unwrap();
let w_q = x
.const_f32_like(
identity_matrix(d_model),
Shape::from_dims(&[d_model, d_model]),
)
.unwrap();
let w_k = x
.const_f32_like(
identity_matrix(d_model),
Shape::from_dims(&[d_model, d_model]),
)
.unwrap();
let w_v = x
.const_f32_like(
identity_matrix(d_model),
Shape::from_dims(&[d_model, d_model]),
)
.unwrap();
let w_o = x
.const_f32_like(
identity_matrix(d_model),
Shape::from_dims(&[d_model, d_model]),
)
.unwrap();
let x_norm = x.rms_norm_last_dim(1e-6).unwrap();
let q = x_norm.matmul(&w_q).unwrap();
let k = x_norm.matmul(&w_k).unwrap();
let v = x_norm.matmul(&w_v).unwrap();
let q_h = q
.reshape(Shape::from_dims(&[1, seq, num_heads, d_head]))
.unwrap()
.permute([0, 2, 1, 3_usize])
.unwrap();
let k_h = k
.reshape(Shape::from_dims(&[1, seq, num_heads, d_head]))
.unwrap()
.permute([0, 2, 1, 3_usize])
.unwrap();
let v_h = v
.reshape(Shape::from_dims(&[1, seq, num_heads, d_head]))
.unwrap()
.permute([0, 2, 1, 3_usize])
.unwrap();
let q_r = q_h.rope(10000.0, 0).unwrap();
let k_r = k_h.rope(10000.0, 0).unwrap();
let k_t = k_r.transpose().unwrap();
let scores = q_r.matmul(&k_t).unwrap();
let attn = scores.softmax_last_dim().unwrap();
let attn_v = attn.matmul(&v_h).unwrap();
let merged = attn_v
.permute([0, 2, 1, 3_usize])
.unwrap()
.reshape(Shape::from_dims(&[1, seq, d_model]))
.unwrap();
let attn_out = merged.matmul(&w_o).unwrap();
let h = x.add(&attn_out).unwrap();
let result = h.realize_f32();
assert_eq!(result.len(), seq * d_model);
for &v in &result {
assert!(
v.is_finite(),
"bridge-based LLaMA block output non-finite: {v}"
);
}
}
fn identity_matrix(n: usize) -> Vec<f32> {
let mut out = vec![0.0_f32; n * n];
for i in 0..n {
out[i * n + i] = 1.0;
}
out
}
}
impl Tensor {
pub fn const_u32_like(
&self,
data: impl Into<Arc<[u32]>>,
shape: impl Into<Shape>,
) -> std::result::Result<Self, fuel_ir::Error> {
Ok(Self {
inner: self.inner.const_u32_like(data, shape)?,
})
}
pub fn const_placeholder_like(&self, shape: impl Into<Shape>, dtype: fuel_ir::DType) -> Self {
Self {
inner: self.inner.const_placeholder_like(shape, dtype),
}
}
pub fn write_slice(&self, source: &Self, ranges: Vec<(usize, usize)>) -> crate::Result<Self> {
let inner = self.inner.write_slice(&source.inner, ranges)?;
Ok(Self { inner })
}
pub fn write_slice_dyn(
&self,
source: &Self,
ranges: Vec<(usize, usize)>,
dyn_axis: usize,
offset: fuel_ir::DynScalar,
) -> crate::Result<Self> {
let inner = self
.inner
.write_slice_dyn(&source.inner, ranges, dyn_axis, offset)?;
Ok(Self { inner })
}
pub fn write_slice_rotating(
&self,
source: &Self,
position: &Self,
axis: usize,
modulus: usize,
ranges: Vec<(usize, usize)>,
) -> crate::Result<Self> {
let inner = self.inner.write_slice_rotating(
&source.inner,
&position.inner,
axis,
modulus,
ranges,
)?;
Ok(Self { inner })
}
pub fn write_slice_doff(
&self,
source: &Self,
offset: &Self,
axis: usize,
ranges: Vec<(usize, usize)>,
) -> crate::Result<Self> {
let inner = self
.inner
.write_slice_doff(&source.inner, &offset.inner, axis, ranges)?;
Ok(Self { inner })
}
pub fn conv2d(
&self,
weight: &Self,
bias: Option<&Self>,
stride: (usize, usize),
padding: (usize, usize),
groups: usize,
) -> std::result::Result<Self, fuel_ir::Error> {
if groups < 1 {
return Err(
fuel_ir::Error::Msg(format!("conv2d: groups must be >= 1, got {groups}",)).bt(),
);
}
let x_shape = self.inner.shape();
let x_dims = x_shape.dims();
let w_shape = weight.inner.shape();
let w_dims = w_shape.dims();
if x_dims.len() != 4 {
return Err(fuel_ir::Error::Msg(format!(
"conv2d: x must be rank 4 [N, Cin, H, W], got {x_dims:?}",
))
.bt());
}
if w_dims.len() != 4 {
return Err(fuel_ir::Error::Msg(format!(
"conv2d: weight must be rank 4 [Cout, Cin/groups, Kh, Kw], got {w_dims:?}",
))
.bt());
}
let (cin, h_in, w_in) = (x_dims[1], x_dims[2], x_dims[3]);
let (cout, cin_per_g, kh, kw) = (w_dims[0], w_dims[1], w_dims[2], w_dims[3]);
if cin != cin_per_g * groups {
return Err(fuel_ir::Error::Msg(format!(
"conv2d: x has {cin} in-channels but weight expects {} ({cin_per_g}*{groups})",
cin_per_g * groups,
))
.bt());
}
if cout % groups != 0 {
return Err(fuel_ir::Error::Msg(format!(
"conv2d: Cout={cout} must be divisible by groups={groups}",
))
.bt());
}
if let Some(b) = bias {
let b_shape = b.inner.shape();
let b_dims = b_shape.dims();
if b_dims != [cout] {
return Err(fuel_ir::Error::Msg(format!(
"conv2d: bias shape {b_dims:?} must match [Cout={cout}]",
))
.bt());
}
}
let (stride_h, stride_w) = stride;
let (pad_h, pad_w) = padding;
if stride_h < 1 || stride_w < 1 {
return Err(fuel_ir::Error::Msg(format!(
"conv2d: stride must be >= 1, got ({stride_h}, {stride_w})",
))
.bt());
}
let h_padded = h_in + 2 * pad_h;
let w_padded = w_in + 2 * pad_w;
if h_padded < kh || w_padded < kw {
return Err(fuel_ir::Error::Msg(format!(
"conv2d: padded input ({h_padded}x{w_padded}) smaller than kernel ({kh}x{kw})",
))
.bt());
}
Ok(Self {
inner: self.inner.conv2d(
&weight.inner,
bias.map(|b| &b.inner),
stride,
padding,
groups,
),
})
}
#[allow(clippy::too_many_arguments)]
pub fn flash_attn(
&self,
k: &Self,
v: &Self,
alibi_slopes: Option<&Self>,
softmax_scale: f32,
causal: bool,
window_size_left: Option<usize>,
window_size_right: Option<usize>,
softcap: Option<f32>,
) -> std::result::Result<Self, fuel_ir::Error> {
let q_shape = self.inner.shape();
let q_dims = q_shape.dims();
let k_shape = k.inner.shape();
let k_dims = k_shape.dims();
let v_shape = v.inner.shape();
let v_dims = v_shape.dims();
if q_dims.len() != 4 {
return Err(fuel_ir::Error::Msg(format!(
"flash_attn: q must be rank 4 [B, Hq, Sq, D], got {q_dims:?}",
))
.bt());
}
if k_dims.len() != 4 {
return Err(fuel_ir::Error::Msg(format!(
"flash_attn: k must be rank 4 [B, Hkv, Sk, D], got {k_dims:?}",
))
.bt());
}
if v_dims.len() != 4 {
return Err(fuel_ir::Error::Msg(format!(
"flash_attn: v must be rank 4 [B, Hkv, Sk, D], got {v_dims:?}",
))
.bt());
}
let (b, hq, _sq, d) = (q_dims[0], q_dims[1], q_dims[2], q_dims[3]);
let (bk, hkv, sk, dk) = (k_dims[0], k_dims[1], k_dims[2], k_dims[3]);
let (bv, hkv_v, sk_v, dv) = (v_dims[0], v_dims[1], v_dims[2], v_dims[3]);
if b != bk || b != bv {
return Err(fuel_ir::Error::Msg(
format!("flash_attn: B mismatch q={b} k={bk} v={bv}",),
)
.bt());
}
if hkv != hkv_v {
return Err(fuel_ir::Error::Msg(format!(
"flash_attn: Hkv mismatch k={hkv} vs v={hkv_v}",
))
.bt());
}
if sk != sk_v {
return Err(fuel_ir::Error::Msg(
format!("flash_attn: Sk mismatch k={sk} vs v={sk_v}",),
)
.bt());
}
if d != dk || d != dv {
return Err(fuel_ir::Error::Msg(format!(
"flash_attn: head_dim mismatch q={d} k={dk} v={dv}",
))
.bt());
}
if hkv == 0 || hq % hkv != 0 {
return Err(fuel_ir::Error::Msg(format!(
"flash_attn: Hq={hq} must be a positive multiple of Hkv={hkv}",
))
.bt());
}
if let Some(a) = alibi_slopes {
let a_shape = a.inner.shape();
let a_dims = a_shape.dims();
if a_dims != [hq] {
return Err(fuel_ir::Error::Msg(format!(
"flash_attn: alibi_slopes must be [Hq={hq}], got {a_dims:?}",
))
.bt());
}
}
Ok(Self {
inner: self.inner.flash_attn(
&k.inner,
&v.inner,
alibi_slopes.map(|t| &t.inner),
softmax_scale,
causal,
window_size_left,
window_size_right,
softcap,
),
})
}
#[allow(clippy::too_many_arguments)]
pub fn paged_attn(
&self,
k_cache: &Self,
v_cache: &Self,
block_table: &Self,
context_lens: &Self,
alibi_slopes: Option<&Self>,
softmax_scale: f32,
block_size: usize,
softcap: Option<f32>,
) -> std::result::Result<Self, fuel_ir::Error> {
if block_size < 1 {
return Err(fuel_ir::Error::Msg("paged_attn: block_size must be >= 1".into()).bt());
}
let q_shape = self.inner.shape();
let q_dims = q_shape.dims();
let kc_shape = k_cache.inner.shape();
let kc_dims = kc_shape.dims();
let vc_shape = v_cache.inner.shape();
let vc_dims = vc_shape.dims();
let bt_shape = block_table.inner.shape();
let bt_dims = bt_shape.dims();
let cl_shape = context_lens.inner.shape();
let cl_dims = cl_shape.dims();
if q_dims.len() != 4 {
return Err(fuel_ir::Error::Msg(format!(
"paged_attn: q must be rank 4 [B, Hq, Sq, D], got {q_dims:?}",
))
.bt());
}
if kc_dims.len() != 4 {
return Err(fuel_ir::Error::Msg(format!(
"paged_attn: k_cache must be rank 4 [num_blocks, block_size, Hkv, D], got {kc_dims:?}",
)).bt());
}
if vc_dims.len() != 4 {
return Err(fuel_ir::Error::Msg(format!(
"paged_attn: v_cache must be rank 4 [num_blocks, block_size, Hkv, D], got {vc_dims:?}",
)).bt());
}
if bt_dims.len() != 2 {
return Err(fuel_ir::Error::Msg(format!(
"paged_attn: block_table must be rank 2 [B, max_blocks], got {bt_dims:?}",
))
.bt());
}
if cl_dims.len() != 1 {
return Err(fuel_ir::Error::Msg(format!(
"paged_attn: context_lens must be rank 1 [B], got {cl_dims:?}",
))
.bt());
}
let (b, hq, _sq, d) = (q_dims[0], q_dims[1], q_dims[2], q_dims[3]);
if kc_dims[1] != block_size {
return Err(fuel_ir::Error::Msg(format!(
"paged_attn: k_cache block dim {} != block_size {block_size}",
kc_dims[1],
))
.bt());
}
if vc_dims[1] != block_size {
return Err(fuel_ir::Error::Msg(format!(
"paged_attn: v_cache block dim {} != block_size {block_size}",
vc_dims[1],
))
.bt());
}
let hkv = kc_dims[2];
if vc_dims[2] != hkv {
return Err(fuel_ir::Error::Msg(format!(
"paged_attn: Hkv mismatch k_cache={hkv} vs v_cache={}",
vc_dims[2],
))
.bt());
}
if kc_dims[3] != d || vc_dims[3] != d {
return Err(fuel_ir::Error::Msg(format!(
"paged_attn: D mismatch q={d} k={} v={}",
kc_dims[3], vc_dims[3],
))
.bt());
}
if hkv == 0 || hq % hkv != 0 {
return Err(fuel_ir::Error::Msg(format!(
"paged_attn: Hq={hq} must be a positive multiple of Hkv={hkv}",
))
.bt());
}
if bt_dims[0] != b {
return Err(fuel_ir::Error::Msg(format!(
"paged_attn: block_table batch dim {} != B={b}",
bt_dims[0],
))
.bt());
}
if cl_dims[0] != b {
return Err(fuel_ir::Error::Msg(format!(
"paged_attn: context_lens len {} != B={b}",
cl_dims[0],
))
.bt());
}
if block_table.inner.dtype() != fuel_ir::DType::U32 {
return Err(fuel_ir::Error::Msg(format!(
"paged_attn: block_table must be U32, got {:?}",
block_table.inner.dtype(),
))
.bt());
}
if context_lens.inner.dtype() != fuel_ir::DType::U32 {
return Err(fuel_ir::Error::Msg(format!(
"paged_attn: context_lens must be U32, got {:?}",
context_lens.inner.dtype(),
))
.bt());
}
if let Some(a) = alibi_slopes {
let a_shape = a.inner.shape();
let a_dims = a_shape.dims();
if a_dims != [hq] {
return Err(fuel_ir::Error::Msg(format!(
"paged_attn: alibi_slopes must be [Hq={hq}], got {a_dims:?}",
))
.bt());
}
}
Ok(Self {
inner: self.inner.paged_attn(
&k_cache.inner,
&v_cache.inner,
&block_table.inner,
&context_lens.inner,
alibi_slopes.map(|t| &t.inner),
softmax_scale,
block_size,
softcap,
),
})
}
pub fn conv_transpose2d(
&self,
weight: &Self,
stride: (usize, usize),
padding: (usize, usize),
output_padding: (usize, usize),
dilation: (usize, usize),
groups: usize,
) -> std::result::Result<Self, fuel_ir::Error> {
if groups < 1 {
return Err(fuel_ir::Error::Msg(format!(
"conv_transpose2d: groups must be >= 1, got {groups}",
))
.bt());
}
let x_shape = self.inner.shape();
let x_dims = x_shape.dims();
let w_shape = weight.inner.shape();
let w_dims = w_shape.dims();
if x_dims.len() != 4 {
return Err(fuel_ir::Error::Msg(format!(
"conv_transpose2d: x must be rank 4 [N, Cin, H, W], got {x_dims:?}",
))
.bt());
}
if w_dims.len() != 4 {
return Err(fuel_ir::Error::Msg(format!(
"conv_transpose2d: weight must be rank 4 [Cin, Cout/groups, Kh, Kw], got {w_dims:?}",
)).bt());
}
let (cin, h_in, w_in) = (x_dims[1], x_dims[2], x_dims[3]);
let (cin_w, cout_per_g, kh, kw) = (w_dims[0], w_dims[1], w_dims[2], w_dims[3]);
if cin != cin_w {
return Err(fuel_ir::Error::Msg(format!(
"conv_transpose2d: x has {cin} in-channels but weight has {cin_w}",
))
.bt());
}
if cin % groups != 0 {
return Err(fuel_ir::Error::Msg(format!(
"conv_transpose2d: Cin={cin} must be divisible by groups={groups}",
))
.bt());
}
let (stride_h, stride_w) = stride;
let (pad_h, pad_w) = padding;
let (out_pad_h, out_pad_w) = output_padding;
let (dil_h, dil_w) = dilation;
if stride_h < 1 || stride_w < 1 {
return Err(fuel_ir::Error::Msg(format!(
"conv_transpose2d: stride must be >= 1, got ({stride_h}, {stride_w})",
))
.bt());
}
if dil_h < 1 || dil_w < 1 {
return Err(fuel_ir::Error::Msg(format!(
"conv_transpose2d: dilation must be >= 1, got ({dil_h}, {dil_w})",
))
.bt());
}
let h_out = h_in.saturating_sub(1) * stride_h + dil_h * (kh - 1) + out_pad_h + 1;
let w_out = w_in.saturating_sub(1) * stride_w + dil_w * (kw - 1) + out_pad_w + 1;
if h_out <= 2 * pad_h || w_out <= 2 * pad_w {
return Err(fuel_ir::Error::Msg(format!(
"conv_transpose2d: padding ({pad_h}x{pad_w}) is larger than the produced output dims ({h_out}x{w_out})",
)).bt());
}
let _ = cout_per_g;
Ok(Self {
inner: self.inner.conv_transpose2d(
&weight.inner,
stride,
padding,
output_padding,
dilation,
groups,
),
})
}
pub fn conv_transpose1d(
&self,
weight: &Self,
stride: usize,
padding: usize,
output_padding: usize,
dilation: usize,
groups: usize,
) -> std::result::Result<Self, fuel_ir::Error> {
if groups < 1 {
return Err(fuel_ir::Error::Msg(format!(
"conv_transpose1d: groups must be >= 1, got {groups}",
))
.bt());
}
let x_shape = self.inner.shape();
let x_dims = x_shape.dims();
let w_shape = weight.inner.shape();
let w_dims = w_shape.dims();
if x_dims.len() != 3 {
return Err(fuel_ir::Error::Msg(format!(
"conv_transpose1d: x must be rank 3 [N, Cin, Lin], got {x_dims:?}",
))
.bt());
}
if w_dims.len() != 3 {
return Err(fuel_ir::Error::Msg(format!(
"conv_transpose1d: weight must be rank 3 [Cin, Cout/groups, K], got {w_dims:?}",
))
.bt());
}
if stride < 1 {
return Err(fuel_ir::Error::Msg(format!(
"conv_transpose1d: stride must be >= 1, got {stride}",
))
.bt());
}
if dilation < 1 {
return Err(fuel_ir::Error::Msg(format!(
"conv_transpose1d: dilation must be >= 1, got {dilation}",
))
.bt());
}
let cin = x_dims[1];
let cin_w = w_dims[0];
if cin != cin_w {
return Err(fuel_ir::Error::Msg(format!(
"conv_transpose1d: x has {cin} in-channels but weight has {cin_w}",
))
.bt());
}
if !cin.is_multiple_of(groups) {
return Err(fuel_ir::Error::Msg(format!(
"conv_transpose1d: Cin={cin} must be divisible by groups={groups}",
))
.bt());
}
Ok(Self {
inner: self.inner.conv_transpose1d(
&weight.inner,
stride,
padding,
output_padding,
dilation,
groups,
),
})
}
}
impl Tensor {
pub fn unsqueeze<D: Dim>(&self, dim: D) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dim = dim.to_index_plus_one(&shape, "unsqueeze")?;
Ok(Self {
inner: self.inner.try_unsqueeze(dim)?,
})
}
pub fn triu(&self, diagonal: i64) -> std::result::Result<Self, fuel_ir::Error> {
Ok(Self {
inner: self.inner.triu(diagonal)?,
})
}
pub fn tril(&self, diagonal: i64) -> std::result::Result<Self, fuel_ir::Error> {
Ok(Self {
inner: self.inner.tril(diagonal)?,
})
}
pub fn log_softmax_last_dim(&self) -> std::result::Result<Self, fuel_ir::Error> {
Ok(Self {
inner: self.inner.log_softmax_last_dim()?,
})
}
pub fn softmax<D: Dim>(&self, dim: D) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let axis: usize = dim.to_index(&shape, "softmax")?;
let m = self.max_keepdim(axis)?;
let shifted = self.broadcast_sub(&m)?;
let e = shifted.exp();
let s = e.sum_keepdim(axis)?;
e.broadcast_div(&s)
}
pub fn log_softmax<D: Dim>(&self, dim: D) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let axis: usize = dim.to_index(&shape, "log_softmax")?;
let m = self.max_keepdim(axis)?;
let shifted = self.broadcast_sub(&m)?;
let lse = shifted.exp().sum_keepdim(axis)?.log();
shifted.broadcast_sub(&lse)
}
pub fn argmin_dim<D: Dim>(&self, dim: D) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dim = dim.to_index(&shape, "argmin_dim")?;
Ok(Self {
inner: self.inner.argmin_dim(dim),
})
}
pub fn masked_fill(
&self,
mask: &Self,
value: fuel_ir::Scalar,
) -> std::result::Result<Self, fuel_ir::Error> {
Ok(Self {
inner: self.inner.masked_fill(&mask.inner, value)?,
})
}
pub fn index_add<D: Dim>(
&self,
dim: D,
indices: &Self,
src: &Self,
) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dim = dim.to_index(&shape, "index_add")?;
if indices.inner.dtype() != fuel_ir::DType::U32 {
return Err(fuel_ir::Error::Msg(format!(
"index_add: index must be U32, got {:?}",
indices.inner.dtype(),
))
.bt());
}
if self.inner.dtype() != src.inner.dtype() {
return Err(fuel_ir::Error::Msg(format!(
"index_add: base and src dtypes must match, got {:?} vs {:?}",
self.inner.dtype(),
src.inner.dtype(),
))
.bt());
}
let base_dims = shape.dims();
let src_shape = src.inner.shape();
let src_dims = src_shape.dims();
if base_dims.len() != src_dims.len() {
return Err(fuel_ir::Error::Msg(format!(
"index_add: base and src must have the same rank, got {} vs {}",
base_dims.len(),
src_dims.len(),
))
.bt());
}
let idx_shape = indices.inner.shape();
let idx_dims = idx_shape.dims();
if idx_dims.len() != 1 {
return Err(fuel_ir::Error::Msg(format!(
"index_add: index must be rank 1, got {idx_dims:?}",
))
.bt());
}
if src_dims[dim] != idx_dims[0] {
return Err(fuel_ir::Error::Msg(format!(
"index_add: src dim {dim} ({}) must match index length ({})",
src_dims[dim], idx_dims[0],
))
.bt());
}
Ok(Self {
inner: self.inner.index_add(dim, &indices.inner, &src.inner),
})
}
pub fn scatter_add<D: Dim>(
&self,
dim: D,
indices: &Self,
src: &Self,
) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.inner.shape();
let dim = dim.to_index(&shape, "scatter_add")?;
if indices.inner.dtype() != fuel_ir::DType::U32 {
return Err(fuel_ir::Error::Msg(format!(
"scatter_add: index must be U32, got {:?}",
indices.inner.dtype(),
))
.bt());
}
if self.inner.dtype() != src.inner.dtype() {
return Err(fuel_ir::Error::Msg(format!(
"scatter_add: base and src dtypes must match, got {:?} vs {:?}",
self.inner.dtype(),
src.inner.dtype(),
))
.bt());
}
let idx_shape = indices.inner.shape();
let src_shape = src.inner.shape();
if idx_shape.dims() != src_shape.dims() {
return Err(fuel_ir::Error::Msg(format!(
"scatter_add: index and src must have the same shape, got {:?} vs {:?}",
idx_shape.dims(),
src_shape.dims(),
))
.bt());
}
Ok(Self {
inner: self.inner.scatter_add(dim, &indices.inner, &src.inner),
})
}
pub fn relu_inplace(&self) -> Self {
Self {
inner: self.inner.relu_inplace(),
}
}
pub fn silu_inplace(&self) -> Self {
Self {
inner: self.inner.silu_inplace(),
}
}
pub fn gelu_tanh_inplace(&self) -> Self {
Self {
inner: self.inner.gelu_tanh_inplace(),
}
}
pub fn tanh_inplace(&self) -> Self {
Self {
inner: self.inner.tanh_inplace(),
}
}
pub fn sigmoid_inplace(&self) -> Self {
Self {
inner: self.inner.sigmoid_inplace(),
}
}
pub fn affine_inplace(&self, mul: f64, add: f64) -> Self {
Self {
inner: self.inner.affine_inplace(mul, add),
}
}
pub fn const_f64_like(
&self,
data: impl Into<Arc<[f64]>>,
shape: impl Into<Shape>,
) -> std::result::Result<Self, fuel_ir::Error> {
Ok(Self {
inner: self.inner.const_f64_like(data, shape)?,
})
}
pub fn const_i64_like(
&self,
data: impl Into<Arc<[i64]>>,
shape: impl Into<Shape>,
) -> std::result::Result<Self, fuel_ir::Error> {
Ok(Self {
inner: self.inner.const_i64_like(data, shape)?,
})
}
pub fn on_device(self, device: &Device) -> Self {
Self {
inner: self.inner.on_device(device.location()),
}
}
pub fn release(&self) -> Self {
Self {
inner: self.inner.release(),
}
}
pub fn move_to_device(&self, device: &Device) -> Self {
Self {
inner: self.inner.move_to_device(device.location()),
}
}
pub fn copy_to_device(&self, device: &Device) -> Self {
Self {
inner: self.inner.copy_to_device(device.location()),
}
}
pub fn backward(&self) -> fuel_graph::GradMap {
self.inner.backward()
}
}
impl Tensor {
pub fn transpose_last_two(&self) -> std::result::Result<Self, fuel_ir::Error> {
self.transpose()
}
pub fn t(&self) -> std::result::Result<Self, fuel_ir::Error> {
self.transpose()
}
pub fn transpose_dims<D1: Dim, D2: Dim>(
&self,
dim1: D1,
dim2: D2,
) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.shape();
let dim1 = dim1.to_index(&shape, "transpose_dims")?;
let dim2 = dim2.to_index(&shape, "transpose_dims")?;
if dim1 == dim2 {
return Ok(self.clone());
}
let rank = shape.dims().len();
let mut axes: Vec<usize> = (0..rank).collect();
axes.swap(dim1, dim2);
self.permute(axes.as_slice())
}
pub fn flatten<D1: Dim, D2: Dim>(
&self,
start_dim: D1,
end_dim: D2,
) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.shape();
let rank = shape.dims().len();
if rank == 0 {
return Ok(self.clone());
}
let start_dim = start_dim.to_index(&shape, "flatten")?;
let end_dim = end_dim.to_index(&shape, "flatten")?;
if start_dim > end_dim {
return Err(fuel_ir::Error::Msg(format!(
"flatten: start_dim={start_dim} > end_dim={end_dim}",
))
.bt());
}
let dims = shape.dims();
let merged: usize = dims[start_dim..=end_dim].iter().product();
let mut new_dims: Vec<usize> = Vec::with_capacity(rank - (end_dim - start_dim));
new_dims.extend_from_slice(&dims[..start_dim]);
new_dims.push(merged);
new_dims.extend_from_slice(&dims[end_dim + 1..]);
self.reshape(new_dims)
}
pub fn flatten_to<D: Dim>(&self, end_dim: D) -> std::result::Result<Self, fuel_ir::Error> {
self.flatten(0_usize, end_dim)
}
pub fn flatten_from<D: Dim>(&self, start_dim: D) -> std::result::Result<Self, fuel_ir::Error> {
let rank = self.shape().dims().len();
if rank == 0 {
return Ok(self.clone());
}
self.flatten(start_dim, rank - 1)
}
pub fn flatten_all(&self) -> std::result::Result<Self, fuel_ir::Error> {
let rank = self.shape().dims().len();
if rank == 0 {
return Ok(self.clone());
}
self.flatten(0, rank - 1)
}
pub fn stack<D: Dim>(args: &[&Self], dim: D) -> std::result::Result<Self, fuel_ir::Error> {
if args.is_empty() {
return Err(fuel_ir::Error::Msg("stack: requires at least one tensor".into()).bt());
}
let reference_shape = args[0].shape();
let reference_dims = reference_shape.dims().to_vec();
let dim = dim.to_index_plus_one(&reference_shape, "stack")?;
for (idx, t) in args.iter().enumerate().skip(1) {
if t.shape().dims() != reference_dims.as_slice() {
return Err(fuel_ir::Error::Msg(format!(
"stack: tensor {idx} shape {:?} != reference shape {:?}",
t.shape().dims(),
reference_dims,
))
.bt());
}
}
let mut iter = args.iter();
let first = iter.next().unwrap().unsqueeze(dim)?;
let mut acc = first;
for t in iter {
let u = t.unsqueeze(dim)?;
acc = acc.concat(&u, dim)?;
}
Ok(acc)
}
pub fn sum_keepdim<D: Dim>(&self, dim: D) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.shape();
let dim = dim.to_index(&shape, "sum_keepdim")?;
self.sum_dim(dim).unwrap().unsqueeze(dim)
}
pub fn mean_keepdim<D: Dim>(&self, dim: D) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.shape();
let dim = dim.to_index(&shape, "mean_keepdim")?;
self.mean_dim(dim).unwrap().unsqueeze(dim)
}
pub fn max_keepdim<D: Dim>(&self, dim: D) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.shape();
let dim = dim.to_index(&shape, "max_keepdim")?;
self.max_dim(dim).unwrap().unsqueeze(dim)
}
pub fn min_keepdim<D: Dim>(&self, dim: D) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.shape();
let dim = dim.to_index(&shape, "min_keepdim")?;
self.min_dim(dim).unwrap().unsqueeze(dim)
}
pub fn var_keepdim<D: Dim>(
&self,
dim: D,
correction: f64,
) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.shape();
let dim = dim.to_index(&shape, "var_keepdim")?;
let dims = shape.dims();
let n = dims[dim];
let mean = self.mean_keepdim(dim)?;
let deviation = self.broadcast_sub(&mean).unwrap();
let squares = deviation.sqr();
let summed = squares.sum_keepdim(dim)?;
let divisor = n as f64 - correction;
Ok(summed.mul_scalar(1.0 / divisor))
}
pub fn var<D: Dim>(
&self,
dim: D,
correction: f64,
) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.shape();
let dim = dim.to_index(&shape, "var")?;
self.var_keepdim(dim, correction)?.squeeze(dim)
}
pub fn affine(&self, mul: f64, add: f64) -> Self {
self.mul_scalar(mul).add_scalar(add)
}
pub fn scale_and_shift(&self, scale: f64, shift: f64) -> Self {
self.affine(scale, shift)
}
pub fn elu(&self, alpha: f64) -> Self {
let neg_branch = self.exp().affine(alpha, -alpha);
let zero = self
.const_f32_like(vec![0.0; self.elem_count()], self.shape())
.expect("elu: vec![_; self.elem_count()] against self.shape() -- same source");
let mask = self.gt(&zero).unwrap();
mask.where_cond(self, &neg_branch).unwrap()
}
pub fn dot(&self, rhs: &Self) -> std::result::Result<Self, fuel_ir::Error> {
let a = self.shape().dims().to_vec();
let b = rhs.shape().dims().to_vec();
if a.len() != 1 || b.len() != 1 {
return Err(fuel_ir::Error::Msg(format!(
"dot: requires rank-1 inputs, got lhs={a:?} rhs={b:?}",
))
.bt());
}
if a[0] != b[0] {
return Err(fuel_ir::Error::Msg(format!(
"dot: length mismatch lhs={} rhs={}",
a[0], b[0],
))
.bt());
}
Ok(self.mul(rhs).unwrap().sum_all())
}
pub fn mv(&self, rhs: &Self) -> std::result::Result<Self, fuel_ir::Error> {
let a = self.shape().dims().to_vec();
let b = rhs.shape().dims().to_vec();
if a.len() != 2 || b.len() != 1 || a[1] != b[0] {
return Err(fuel_ir::Error::Msg(format!(
"mv: shape mismatch lhs={a:?} rhs={b:?} (need [m,n] · [n])",
))
.bt());
}
let rhs_col = rhs.unsqueeze(1_usize)?;
let prod = self.matmul(&rhs_col).unwrap();
prod.squeeze(1_usize)
}
pub fn matvec(&self, rhs: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.mv(rhs)
}
pub fn broadcast_matmul(&self, rhs: &Self) -> std::result::Result<Self, fuel_ir::Error> {
self.matmul(rhs)
}
pub fn ones_like(&self) -> std::result::Result<Self, fuel_ir::Error> {
let n = self.elem_count();
let shape = self.shape();
match self.dtype() {
DType::F32 => self.const_f32_like(vec![1.0_f32; n], shape),
DType::F64 => self.const_f64_like(vec![1.0_f64; n], shape),
DType::BF16 => self.const_bf16_like(vec![half::bf16::ONE; n], shape),
DType::F16 => self.const_f16_like(vec![half::f16::ONE; n], shape),
DType::U32 => self.const_u32_like(vec![1_u32; n], shape),
DType::I64 => self.const_i64_like(vec![1_i64; n], shape),
other => {
Err(fuel_ir::Error::Msg(format!("ones_like: unsupported dtype {other:?}",)).bt())
}
}
}
pub fn zeros_like(&self) -> std::result::Result<Self, fuel_ir::Error> {
let n = self.elem_count();
let shape = self.shape();
match self.dtype() {
DType::F32 => self.const_f32_like(vec![0.0_f32; n], shape),
DType::F64 => self.const_f64_like(vec![0.0_f64; n], shape),
DType::BF16 => self.const_bf16_like(vec![half::bf16::ZERO; n], shape),
DType::F16 => self.const_f16_like(vec![half::f16::ZERO; n], shape),
DType::U32 => self.const_u32_like(vec![0_u32; n], shape),
DType::I64 => self.const_i64_like(vec![0_i64; n], shape),
other => {
Err(fuel_ir::Error::Msg(format!("zeros_like: unsupported dtype {other:?}",)).bt())
}
}
}
pub fn ones(
shape: impl Into<Shape>,
dtype: DType,
device: &Device,
) -> std::result::Result<Self, fuel_ir::Error> {
let shape = shape.into();
let n = shape.elem_count();
match dtype {
DType::F32 => Self::from_f32(vec![1.0_f32; n], shape, device),
DType::F64 => Self::from_f64(vec![1.0_f64; n], shape, device),
DType::BF16 => Self::from_bf16(vec![half::bf16::ONE; n], shape, device),
DType::F16 => Self::from_f16(vec![half::f16::ONE; n], shape, device),
DType::U32 => Self::from_u32(vec![1_u32; n], shape, device),
other => Err(fuel_ir::Error::Msg(format!("ones: unsupported dtype {other:?}",)).bt()),
}
}
pub fn zeros(
shape: impl Into<Shape>,
dtype: DType,
device: &Device,
) -> std::result::Result<Self, fuel_ir::Error> {
let shape = shape.into();
let n = shape.elem_count();
match dtype {
DType::F32 => Self::from_f32(vec![0.0_f32; n], shape, device),
DType::F64 => Self::from_f64(vec![0.0_f64; n], shape, device),
DType::BF16 => Self::from_bf16(vec![half::bf16::ZERO; n], shape, device),
DType::F16 => Self::from_f16(vec![half::f16::ZERO; n], shape, device),
DType::U32 => Self::from_u32(vec![0_u32; n], shape, device),
other => Err(fuel_ir::Error::Msg(format!("zeros: unsupported dtype {other:?}",)).bt()),
}
}
pub fn full(
shape: impl Into<Shape>,
value: fuel_ir::Scalar,
device: &Device,
) -> std::result::Result<Self, fuel_ir::Error> {
let shape = shape.into();
let n = shape.elem_count();
match value {
fuel_ir::Scalar::F32(v) => Self::from_f32(vec![v; n], shape, device),
fuel_ir::Scalar::F64(v) => Self::from_f64(vec![v; n], shape, device),
fuel_ir::Scalar::BF16(v) => Self::from_bf16(vec![v; n], shape, device),
fuel_ir::Scalar::F16(v) => Self::from_f16(vec![v; n], shape, device),
fuel_ir::Scalar::U32(v) => Self::from_u32(vec![v; n], shape, device),
other => Err(fuel_ir::Error::Msg(format!(
"full: unsupported scalar dtype {:?}",
other.dtype(),
))
.bt()),
}
}
pub fn eye(n: usize, dtype: DType, device: &Device) -> Self {
let mut data = vec![0.0_f32; n * n];
for i in 0..n {
data[i * n + i] = 1.0;
}
let base = Self::from_f32(data, vec![n, n], device).expect(
"eye: data is vec![_; n*n] and shape is [n, n] (elem_count n*n), \n both computed from `n` in this function -- lengths cannot disagree",
);
if dtype == DType::F32 {
base
} else {
base.to_dtype(dtype).unwrap()
}
}
pub fn split_heads(
&self,
num_heads: usize,
head_dim: usize,
) -> std::result::Result<Self, fuel_ir::Error> {
let dims = self.inner.shape().dims().to_vec();
debug_assert_eq!(
dims.len(),
3,
"split_heads: input must be rank 3 (B, N, embed), got {dims:?}"
);
debug_assert_eq!(
dims[2],
num_heads * head_dim,
"split_heads: trailing dim ({}) != num_heads * head_dim ({} * {} = {})",
dims[2],
num_heads,
head_dim,
num_heads * head_dim
);
let b = dims[0];
let n = dims[1];
self.reshape(Shape::from_dims(&[b, n, num_heads, head_dim]))?
.permute([0, 2, 1, 3_usize])
}
pub fn merge_heads(&self) -> std::result::Result<Self, fuel_ir::Error> {
let dims = self.inner.shape().dims().to_vec();
debug_assert_eq!(
dims.len(),
4,
"merge_heads: input must be rank 4 (B, heads, N, head_dim), got {dims:?}"
);
let b = dims[0];
let num_heads = dims[1];
let n = dims[2];
let head_dim = dims[3];
self.permute([0, 2, 1, 3_usize])?
.reshape(Shape::from_dims(&[b, n, num_heads * head_dim]))
}
pub fn add_trailing_bias(
&self,
bias: std::sync::Arc<[f32]>,
) -> std::result::Result<Self, fuel_ir::Error> {
let n = bias.len();
let bias_t = self.const_like_dtype(&bias, Shape::from_dims(&[n]), self.dtype())?;
self.broadcast_add(&bias_t)
}
pub fn embed_tokens(
embed_table: std::sync::Arc<[f32]>,
vocab_size: usize,
hidden: usize,
tokens: &[u32],
device: &crate::Device,
) -> std::result::Result<Self, fuel_ir::Error> {
let seq = tokens.len();
if seq == 0 {
return Err(fuel_ir::Error::Msg("embed_tokens: tokens must be non-empty".into()).bt());
}
let embed = Self::from_f32(embed_table, Shape::from_dims(&[vocab_size, hidden]), device)?;
let token_ids = embed.const_u32_like(tokens.to_vec(), Shape::from_dims(&[seq]))?;
embed
.index_select(0_usize, &token_ids)?
.reshape(Shape::from_dims(&[1, seq, hidden]))
}
pub fn embed_tokens_anchored(
&self,
embed_table: std::sync::Arc<[f32]>,
vocab_size: usize,
hidden: usize,
tokens: &[u32],
) -> std::result::Result<Self, fuel_ir::Error> {
let seq = tokens.len();
if seq == 0 {
return Err(fuel_ir::Error::Msg(
"embed_tokens_anchored: tokens must be non-empty".into(),
)
.bt());
}
let embed = self.const_f32_like(embed_table, Shape::from_dims(&[vocab_size, hidden]))?;
let token_ids = self.const_u32_like(tokens.to_vec(), Shape::from_dims(&[seq]))?;
embed
.index_select(0_usize, &token_ids)?
.reshape(Shape::from_dims(&[1, seq, hidden]))
}
pub fn rope_tables_const(
&self,
theta: f64,
start_pos: usize,
seq: usize,
head_dim: usize,
) -> (Self, Self) {
let (cos_data, sin_data) = fuel_graph::build_rope_tables(theta, start_pos, seq, head_dim);
let rope_shape = Shape::from_dims(&[seq, head_dim]);
let rope_cos = self
.const_f32_like(cos_data, rope_shape.clone())
.expect("rope tables: build_rope_tables must return seq*head_dim elements");
let rope_sin = self
.const_f32_like(sin_data, rope_shape)
.expect("rope tables: build_rope_tables must return seq*head_dim elements");
(rope_cos, rope_sin)
}
pub fn rope_delta_tables_const(
&self,
theta: f64,
delta: usize,
rows: usize,
head_dim: usize,
) -> (Self, Self) {
let (c1, s1) = fuel_graph::build_rope_tables(theta, delta, 1, head_dim);
let mut cos = Vec::with_capacity(rows * head_dim);
let mut sin = Vec::with_capacity(rows * head_dim);
for _ in 0..rows {
cos.extend_from_slice(&c1);
sin.extend_from_slice(&s1);
}
let shape = Shape::from_dims(&[rows, head_dim]);
(
self.const_f32_like(cos, shape.clone()).expect(
"rope tables: rows appends of c1 (local); c1 has head_dim elements (not local)",
),
self.const_f32_like(sin, shape).expect(
"rope tables: rows appends of s1 (local); s1 has head_dim elements (not local)",
),
)
}
pub fn rope_delta_rotate_block_f32(
dev: &crate::Device,
k_block: &[f32],
theta: f64,
delta: usize,
block_size: usize,
n_kv_heads: usize,
head_dim: usize,
) -> Vec<f32> {
let pool_shape = Shape::from_dims(&[block_size, n_kv_heads, head_dim]);
let k = Tensor::from_f32(k_block.to_vec(), pool_shape, dev).expect(
"rope_delta_rotate_block_f32: k_block length must equal block_size*n_kv_heads*head_dim -- both caller-supplied, so this function cannot prove it; the -> Vec<f32> signature has no error channel (see the sibling permute/reshape expects)",
);
let k4 = k
.permute([1, 0, 2])
.expect("rope_delta_rotate_block_f32: permute [bs,Hkv,D]->[Hkv,bs,D]")
.reshape(Shape::from_dims(&[1, n_kv_heads, block_size, head_dim]))
.expect("rope_delta_rotate_block_f32: reshape to [1,Hkv,bs,D]");
let (cos, sin) = k4.rope_delta_tables_const(theta, delta, block_size, head_dim);
k4.rope_with_tables_decomposed(&cos, &sin)
.expect("rope_delta_rotate_block_f32: rope_with_tables_decomposed")
.reshape(Shape::from_dims(&[n_kv_heads, block_size, head_dim]))
.expect("rope_delta_rotate_block_f32: reshape back to [Hkv,bs,D]")
.permute([1, 0, 2])
.expect("rope_delta_rotate_block_f32: permute [Hkv,bs,D]->[bs,Hkv,D]")
.realize_f32()
}
pub fn rope_tables_const_batched(
&self,
theta: f64,
positions: &[usize],
head_dim: usize,
) -> (Self, Self) {
let b = positions.len();
let mut cos_data = Vec::with_capacity(b * head_dim);
let mut sin_data = Vec::with_capacity(b * head_dim);
for &pos in positions {
let (c, s) = fuel_graph::build_rope_tables(theta, pos, 1, head_dim);
cos_data.extend_from_slice(&c);
sin_data.extend_from_slice(&s);
}
let shape = Shape::from_dims(&[b, 1, 1, head_dim]);
let rope_cos = self
.const_f32_like(cos_data, shape.clone())
.expect("batched rope tables: b IS positions.len() (local); build_rope_tables must yield head_dim per position (not local)");
let rope_sin = self
.const_f32_like(sin_data, shape)
.expect("batched rope tables: b IS positions.len() (local); build_rope_tables must yield head_dim per position (not local)");
(rope_cos, rope_sin)
}
pub fn rope_batched(
&self,
cos: &Self,
sin: &Self,
) -> std::result::Result<Self, fuel_ir::Error> {
let dims: Vec<usize> = self.inner.shape().dims().to_vec();
let rank = dims.len();
if rank < 2 {
return Err(fuel_ir::Error::Msg(format!(
"rope_batched: expected rank >= 2, got {dims:?}",
))
.bt());
}
let d = dims[rank - 1];
if !d.is_multiple_of(2) {
return Err(
fuel_ir::Error::Msg(format!("rope_batched: head_dim {d} must be even",)).bt(),
);
}
let half = d / 2;
let target = Shape::from_dims(&dims);
let cos_b = cos.broadcast_to(target.clone())?;
let sin_b = sin.broadcast_to(target)?;
let first = self.slice(rank - 1, 0, half)?;
let second = self.slice(rank - 1, half, half)?;
let rotated = second.neg().concat(&first, rank - 1)?;
let left = self.mul(&cos_b)?;
let right = rotated.mul(&sin_b)?;
left.add(&right)
}
pub fn rope_partial(
&self,
rope_cos: &Self,
rope_sin: &Self,
rope_dim: usize,
) -> std::result::Result<Self, fuel_ir::Error> {
let dims = self.inner.shape();
let dims = dims.dims();
let head_dim = *dims.last().ok_or_else(|| {
fuel_ir::Error::Msg("rope_partial: receiver must have at least one dimension".into())
.bt()
})?;
if rope_dim == head_dim {
return self.rope_with_tables(rope_cos, rope_sin);
}
if rope_dim > head_dim {
return Err(fuel_ir::Error::Msg(format!(
"rope_partial: rope_dim={rope_dim} exceeds head_dim={head_dim}",
))
.bt());
}
let last = dims.len() - 1;
let pass_dim = head_dim - rope_dim;
let rot = self.slice(last, 0, rope_dim)?;
let pass = self.slice(last, rope_dim, pass_dim)?;
let rot_rotated = rot.rope_with_tables(rope_cos, rope_sin)?;
rot_rotated.concat(&pass, last)
}
pub fn add_optional_trailing_bias(
&self,
bias: Option<&std::sync::Arc<[f32]>>,
) -> std::result::Result<Self, fuel_ir::Error> {
match bias {
None => Ok(self.clone()),
Some(b) => self.add_trailing_bias(std::sync::Arc::clone(b)),
}
}
pub fn softcap(&self, cap: f64) -> Self {
self.mul_scalar(1.0 / cap).tanh().mul_scalar(cap)
}
pub fn softcap_optional(&self, cap: Option<f64>) -> Self {
match cap {
Some(c) if c > 0.0 => self.softcap(c),
_ => self.clone(),
}
}
pub fn rms_norm_affine_with_offset(
&self,
gain: &[f32],
offset: f32,
eps: f64,
) -> std::result::Result<Self, fuel_ir::Error> {
let shifted: std::sync::Arc<[f32]> =
std::sync::Arc::from(gain.iter().map(|g| *g + offset).collect::<Vec<_>>());
self.rms_norm_affine(shifted, eps)
}
pub fn rms_norm_affine(
&self,
gain: std::sync::Arc<[f32]>,
eps: f64,
) -> std::result::Result<Self, fuel_ir::Error> {
let hidden = gain.len();
let normed = self.rms_norm_last_dim(eps)?;
let gain_t = self.const_f32_like(gain, Shape::from_dims(&[hidden]))?;
normed.broadcast_mul(&gain_t)
}
pub fn global_avg_pool_2d(&self) -> std::result::Result<Self, fuel_ir::Error> {
let dims = self.inner.shape().dims().to_vec();
debug_assert_eq!(
dims.len(),
4,
"global_avg_pool_2d: input must be rank 4 (B, C, H, W), got {dims:?}"
);
self.mean_dim(3_usize)?.mean_dim(2_usize)
}
pub fn channel_affine_4d(
&self,
gain: std::sync::Arc<[f32]>,
bias: std::sync::Arc<[f32]>,
) -> std::result::Result<Self, fuel_ir::Error> {
let dims = self.inner.shape().dims().to_vec();
debug_assert_eq!(
dims.len(),
4,
"channel_affine_4d: input must be rank 4 (B, C, H, W), got {dims:?}"
);
let channels = dims[1];
debug_assert_eq!(
gain.len(),
channels,
"channel_affine_4d: gain len ({}) != C ({})",
gain.len(),
channels
);
debug_assert_eq!(
bias.len(),
channels,
"channel_affine_4d: bias len ({}) != C ({})",
bias.len(),
channels
);
let w_t = self
.const_f32_like(gain, Shape::from_dims(&[channels]))?
.reshape(Shape::from_dims(&[1, channels, 1, 1]))?;
let b_t = self
.const_f32_like(bias, Shape::from_dims(&[channels]))?
.reshape(Shape::from_dims(&[1, channels, 1, 1]))?;
self.broadcast_mul(&w_t)?.broadcast_add(&b_t)
}
pub fn additive_causal_mask_like(anchor: &Tensor, seq_len: usize) -> Self {
let mut data = vec![0.0_f32; seq_len * seq_len];
for i in 0..seq_len {
for j in (i + 1)..seq_len {
data[i * seq_len + j] = f32::NEG_INFINITY;
}
}
anchor
.const_f32_like(
std::sync::Arc::from(data),
Shape::from_dims(&[seq_len, seq_len]),
)
.expect(
"additive_causal_mask_like: vec![_; seq_len*seq_len] against [seq_len, seq_len]",
)
}
pub fn tril2(n: usize, dtype: DType, device: &Device) -> Self {
let mut data = vec![0.0_f32; n * n];
for i in 0..n {
for j in 0..=i {
data[i * n + j] = 1.0;
}
}
let base = Self::from_f32(data, vec![n, n], device).expect(
"tril2: data is vec![_; n*n] and shape is [n, n] (elem_count n*n), \n both computed from `n` in this function -- lengths cannot disagree",
);
if dtype == DType::F32 {
base
} else {
base.to_dtype(dtype).unwrap()
}
}
pub fn triu2(n: usize, dtype: DType, device: &Device) -> Self {
let mut data = vec![0.0_f32; n * n];
for i in 0..n {
for j in i..n {
data[i * n + j] = 1.0;
}
}
let base = Self::from_f32(data, vec![n, n], device).expect(
"triu2: data is vec![_; n*n] and shape is [n, n] (elem_count n*n), \n both computed from `n` in this function -- lengths cannot disagree",
);
if dtype == DType::F32 {
base
} else {
base.to_dtype(dtype).unwrap()
}
}
pub fn narrow<D: Dim>(
&self,
dim: D,
start: usize,
len: usize,
) -> std::result::Result<Self, fuel_ir::Error> {
self.slice(dim, start, len)
}
pub fn chunk<D: Dim>(
&self,
chunks: usize,
dim: D,
) -> std::result::Result<Vec<Self>, fuel_ir::Error> {
let shape = self.shape();
let dim = dim.to_index(&shape, "chunk")?;
if chunks == 0 {
return Err(fuel_ir::Error::Msg("chunk: chunks must be > 0".into()).bt());
}
let dims = shape.dims();
let size = dims[dim];
if size < chunks {
return Ok((0..size).map(|i| self.slice(dim, i, 1).unwrap()).collect());
}
let base = size / chunks;
let extra = size % chunks;
let mut out = Vec::with_capacity(chunks);
let mut start = 0;
for i in 0..chunks {
let len = if i < extra { base + 1 } else { base };
out.push(self.slice(dim, start, len).unwrap());
start += len;
}
Ok(out)
}
pub fn get(&self, i: usize) -> std::result::Result<Self, fuel_ir::Error> {
let dims = self.shape().dims().to_vec();
if dims.is_empty() {
return Ok(self.clone());
}
self.slice(0, i, 1).unwrap().squeeze(0)
}
pub fn get_on_dim<D: Dim>(
&self,
dim: D,
index: usize,
) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.shape();
let dim = dim.to_index(&shape, "get_on_dim")?;
self.slice(dim, index, 1).unwrap().squeeze(dim)
}
pub fn sum_dims<D: Dims>(&self, dims: D) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.shape();
let mut sorted = dims.to_indexes(&shape, "sum_dims")?;
sorted.sort_by(|a, b| b.cmp(a));
sorted.dedup();
let mut acc = self.clone();
for d in sorted {
acc = acc.sum_dim(d)?;
}
Ok(acc)
}
pub fn mean_dims<D: Dims>(&self, dims: D) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.shape();
let mut sorted = dims.to_indexes(&shape, "mean_dims")?;
sorted.sort_by(|a, b| b.cmp(a));
sorted.dedup();
let mut acc = self.clone();
for d in sorted {
acc = acc.mean_dim(d)?;
}
Ok(acc)
}
pub fn sum_dims_keepdim(&self, dims: &[usize]) -> std::result::Result<Self, fuel_ir::Error> {
let mut sorted: Vec<usize> = dims.to_vec();
sorted.sort();
sorted.dedup();
let mut acc = self.clone();
for d in sorted {
acc = acc.sum_keepdim(d)?;
}
Ok(acc)
}
pub fn mean_dims_keepdim(&self, dims: &[usize]) -> std::result::Result<Self, fuel_ir::Error> {
let mut sorted: Vec<usize> = dims.to_vec();
sorted.sort();
sorted.dedup();
let mut acc = self.clone();
for d in sorted {
acc = acc.mean_keepdim(d)?;
}
Ok(acc)
}
pub fn rand_like(&self, lo: f64, up: f64) -> std::result::Result<Self, fuel_ir::Error> {
Self::rand(self.shape(), lo, up, self.dtype(), &Device::cpu())
}
pub fn randn_like(&self, mean: f64, stdev: f64) -> std::result::Result<Self, fuel_ir::Error> {
Self::randn(self.shape(), mean, stdev, self.dtype(), &Device::cpu())
}
pub fn rand(
shape: impl Into<Shape>,
lo: f64,
up: f64,
dtype: DType,
device: &Device,
) -> std::result::Result<Self, fuel_ir::Error> {
use rand::Rng;
let shape = shape.into();
let n = shape.elem_count();
let mut rng = rand::rng();
match dtype {
DType::F32 => {
let data: Vec<f32> = (0..n).map(|_| rng.random_range(lo..up) as f32).collect();
Self::from_f32(data, shape, device)
}
DType::F64 => {
let data: Vec<f64> = (0..n).map(|_| rng.random_range(lo..up)).collect();
Self::from_f64(data, shape, device)
}
DType::BF16 => {
let data: Vec<half::bf16> = (0..n)
.map(|_| half::bf16::from_f64(rng.random_range(lo..up)))
.collect();
Self::from_bf16(data, shape, device)
}
DType::F16 => {
let data: Vec<half::f16> = (0..n)
.map(|_| half::f16::from_f64(rng.random_range(lo..up)))
.collect();
Self::from_f16(data, shape, device)
}
other => {
Err(fuel_ir::Error::Msg(format!("Tensor::rand: unsupported dtype {other:?}",)).bt())
}
}
}
pub fn randn(
shape: impl Into<Shape>,
mean: f64,
stdev: f64,
dtype: DType,
device: &Device,
) -> std::result::Result<Self, fuel_ir::Error> {
use rand_distr::{Distribution, Normal};
let shape = shape.into();
let n = shape.elem_count();
let normal = Normal::new(mean, stdev).map_err(|e| {
fuel_ir::Error::Msg(format!("Tensor::randn: invalid stdev={stdev}: {e}",)).bt()
})?;
let mut rng = rand::rng();
match dtype {
DType::F32 => {
let data: Vec<f32> = (0..n).map(|_| normal.sample(&mut rng) as f32).collect();
Self::from_f32(data, shape, device)
}
DType::F64 => {
let data: Vec<f64> = (0..n).map(|_| normal.sample(&mut rng)).collect();
Self::from_f64(data, shape, device)
}
DType::BF16 => {
let data: Vec<half::bf16> = (0..n)
.map(|_| half::bf16::from_f64(normal.sample(&mut rng)))
.collect();
Self::from_bf16(data, shape, device)
}
DType::F16 => {
let data: Vec<half::f16> = (0..n)
.map(|_| half::f16::from_f64(normal.sample(&mut rng)))
.collect();
Self::from_f16(data, shape, device)
}
other => Err(fuel_ir::Error::Msg(format!(
"Tensor::randn: unsupported dtype {other:?}",
))
.bt()),
}
}
pub fn arange(start: f32, end: f32, device: &Device) -> Self {
Self::arange_step(start, end, 1.0, device)
}
pub fn arange_step(start: f32, end: f32, step: f32, device: &Device) -> Self {
assert!(step != 0.0, "arange_step: step must be non-zero");
let mut data = Vec::new();
let mut current = start;
if step > 0.0 {
while current < end {
data.push(current);
current += step;
}
} else {
while current > end {
data.push(current);
current += step;
}
}
let n = data.len();
Self::from_f32(data, vec![n], device).expect(
"arange_step: shape is [n] where n IS data.len(), computed one line above from that same buffer -- the lengths cannot disagree",
)
}
pub fn linspace(start: f32, end: f32, n: usize, device: &Device) -> Self {
assert!(n >= 1, "linspace: n must be >= 1");
if n == 1 {
return Self::from_f32(vec![start], vec![1], device)
.expect("linspace: vec![start] has exactly 1 element and the shape is [1]");
}
let step = (end - start) / ((n - 1) as f32);
let data: Vec<f32> = (0..n).map(|i| start + step * (i as f32)).collect();
Self::from_f32(data, vec![n], device).expect(
"linspace: data is (0..n).map(..).collect() so it has exactly n elements, and the shape is [n] -- both from the same `n`",
)
}
pub fn norm(&self) -> Self {
self.sqr().sum_all().sqrt()
}
pub fn conv1d(
&self,
weight: &Self,
bias: Option<&Self>,
stride: usize,
padding: usize,
groups: usize,
) -> std::result::Result<Self, fuel_ir::Error> {
let x_dims = self.shape().dims().to_vec();
let w_dims = weight.shape().dims().to_vec();
if x_dims.len() != 3 {
return Err(fuel_ir::Error::Msg(format!(
"conv1d: x must be rank 3 [N, Cin, T], got {x_dims:?}",
))
.bt());
}
if w_dims.len() != 3 {
return Err(fuel_ir::Error::Msg(format!(
"conv1d: weight must be rank 3 [Cout, Cin/groups, K], got {w_dims:?}",
))
.bt());
}
if groups < 1 {
return Err(fuel_ir::Error::Msg("conv1d: groups must be ≥ 1".into()).bt());
}
if stride < 1 {
return Err(fuel_ir::Error::Msg("conv1d: stride must be ≥ 1".into()).bt());
}
let x_4d = self.unsqueeze(2_usize)?;
let w_4d = weight.unsqueeze(2_usize)?;
let out_4d = x_4d.conv2d(&w_4d, bias, (1, stride), (0, padding), groups)?;
out_4d.squeeze(2)
}
pub fn conv1d_with_algo<A>(
&self,
weight: &Self,
bias: Option<&Self>,
stride: usize,
padding: usize,
groups: usize,
_algo: A,
) -> std::result::Result<Self, fuel_ir::Error> {
self.conv1d(weight, bias, stride, padding, groups)
}
pub fn avg_pool2d(
&self,
kernel: (usize, usize),
stride: (usize, usize),
padding: (usize, usize),
count_include_pad: bool,
) -> std::result::Result<Self, fuel_ir::Error> {
let dims = self.shape().dims().to_vec();
if dims.len() != 4 {
return Err(fuel_ir::Error::Msg(format!(
"avg_pool2d: input must be rank 4 [N, C, H, W], got {dims:?}",
))
.bt());
}
let (c, h, w) = (dims[1], dims[2], dims[3]);
let (kh, kw) = kernel;
if kh == 0 || kw == 0 {
return Err(
fuel_ir::Error::Msg("avg_pool2d: kernel sizes must be positive".into()).bt(),
);
}
if count_include_pad || padding == (0, 0) {
let inv = 1.0_f32 / ((kh * kw) as f32);
let weight =
self.const_f32_like(vec![inv; c * kh * kw], Shape::from_dims(&[c, 1, kh, kw]))?;
return self.conv2d(&weight, None, stride, padding, c);
}
if let Some((axis, out_idx)) = first_zero_tap_output_index(h, kh, stride.0, padding.0)
.map(|i| ("H", i))
.or_else(|| first_zero_tap_output_index(w, kw, stride.1, padding.1).map(|i| ("W", i)))
{
return Err(fuel_ir::Error::Msg(format!(
"avg_pool2d: count_include_pad=false has no valid (non-padded) \
taps for output position {out_idx} on the {axis} axis \
(kernel={kernel:?}, stride={stride:?}, padding={padding:?}, \
input H×W={h}×{w}) — the average would divide by zero",
))
.bt());
}
let sum_weight = self.const_f32_like(
vec![1.0_f32; c * kh * kw],
Shape::from_dims(&[c, 1, kh, kw]),
)?;
let summed = self.conv2d(&sum_weight, None, stride, padding, c)?;
let ones_hw = self.const_f32_like(vec![1.0_f32; h * w], Shape::from_dims(&[1, 1, h, w]))?;
let count_weight =
self.const_f32_like(vec![1.0_f32; kh * kw], Shape::from_dims(&[1, 1, kh, kw]))?;
let count = ones_hw.conv2d(&count_weight, None, stride, padding, 1)?;
summed.broadcast_div(&count)
}
pub fn avg_pool2d_with_stride(
&self,
kernel: (usize, usize),
stride: (usize, usize),
) -> std::result::Result<Self, fuel_ir::Error> {
self.avg_pool2d(kernel, stride, (0, 0), true)
}
pub fn max_pool2d(
&self,
kernel: (usize, usize),
stride: (usize, usize),
padding: (usize, usize),
) -> std::result::Result<Self, fuel_ir::Error> {
self.max_pool2d_with_pad_value(kernel, stride, padding, f32::NEG_INFINITY)
}
pub fn max_pool2d_with_pad_value(
&self,
kernel: (usize, usize),
stride: (usize, usize),
padding: (usize, usize),
pad_value: f32,
) -> std::result::Result<Self, fuel_ir::Error> {
let dims = self.shape().dims().to_vec();
if dims.len() != 4 {
return Err(fuel_ir::Error::Msg(format!(
"max_pool2d: input must be rank 4 [N, C, H, W], got {dims:?}",
))
.bt());
}
let (n, c, h, w) = (dims[0], dims[1], dims[2], dims[3]);
let (kh, kw) = kernel;
let (sh, sw) = stride;
let (ph, pw) = padding;
if kh == 0 || kw == 0 {
return Err(
fuel_ir::Error::Msg("max_pool2d: kernel sizes must be positive".into()).bt(),
);
}
if sh == 0 || sw == 0 {
return Err(fuel_ir::Error::Msg("max_pool2d: strides must be positive".into()).bt());
}
let h_padded_min = h + 2 * ph;
let w_padded_min = w + 2 * pw;
if h_padded_min < kh || w_padded_min < kw {
return Err(fuel_ir::Error::Msg(format!(
"max_pool2d: padded input ({h_padded_min}×{w_padded_min}) smaller than kernel ({kh}×{kw})",
)).bt());
}
let h_out = (h_padded_min - kh) / sh + 1;
let w_out = (w_padded_min - kw) / sw + 1;
let h_total = h_out * sh + (kh - 1);
let w_total = w_out * sw + (kw - 1);
let extra_h = h_total.saturating_sub(h_padded_min);
let extra_w = w_total.saturating_sub(w_padded_min);
let padded = self
.pad_with_value(2, ph, ph + extra_h, pad_value)?
.pad_with_value(3, pw, pw + extra_w, pad_value)?;
let mut acc: Option<Tensor> = None;
for ky in 0..kh {
let row_slice = padded.slice(2, ky, h_out * sh).unwrap();
let row_reshaped = row_slice.reshape(vec![n, c, h_out, sh, w_total])?;
let row_tap = row_reshaped.slice(3, 0, 1).unwrap().squeeze(3)?;
for kx in 0..kw {
let col_slice = row_tap.slice(3, kx, w_out * sw).unwrap();
let col_reshaped = col_slice.reshape(vec![n, c, h_out, w_out, sw])?;
let win = col_reshaped.slice(4, 0, 1).unwrap().squeeze(4)?;
acc = Some(match acc {
None => win,
Some(a) => a.maximum(&win).unwrap(),
});
}
}
acc.ok_or_else(|| fuel_ir::Error::Msg("max_pool2d: empty kernel".into()).bt())
}
pub fn max_pool2d_with_stride(
&self,
kernel: (usize, usize),
stride: (usize, usize),
) -> std::result::Result<Self, fuel_ir::Error> {
self.max_pool2d(kernel, stride, (0, 0))
}
pub fn upsample_nearest2d(&self, scale: usize) -> std::result::Result<Self, fuel_ir::Error> {
if scale == 0 {
return Err(
fuel_ir::Error::Msg("upsample_nearest2d: scale must be positive".into()).bt(),
);
}
let dims = self.shape().dims().to_vec();
if dims.len() != 4 {
return Err(fuel_ir::Error::Msg(format!(
"upsample_nearest2d: input must be rank 4 [N, C, H, W], got {dims:?}",
))
.bt());
}
if scale == 1 {
return Ok(self.clone());
}
let (n, c, h, w) = (dims[0], dims[1], dims[2], dims[3]);
let expanded = self.reshape(vec![n, c, h, 1, w, 1])?;
let h_expanded = (0..scale)
.fold(None, |acc: Option<Tensor>, _| {
Some(match acc {
None => expanded.clone(),
Some(a) => a.concat(&expanded, 3).unwrap(),
})
})
.unwrap();
let h_then_w = (0..scale)
.fold(None, |acc: Option<Tensor>, _| {
Some(match acc {
None => h_expanded.clone(),
Some(a) => a.concat(&h_expanded, 5).unwrap(),
})
})
.unwrap();
h_then_w.reshape(vec![n, c, h * scale, w * scale])
}
pub fn upsample_nearest1d(&self, scale: usize) -> std::result::Result<Self, fuel_ir::Error> {
if scale == 0 {
return Err(
fuel_ir::Error::Msg("upsample_nearest1d: scale must be positive".into()).bt(),
);
}
let dims = self.shape().dims().to_vec();
if dims.len() != 3 {
return Err(fuel_ir::Error::Msg(format!(
"upsample_nearest1d: input must be rank 3 [N, C, T], got {dims:?}",
))
.bt());
}
if scale == 1 {
return Ok(self.clone());
}
let (n, c, t) = (dims[0], dims[1], dims[2]);
let expanded = self.reshape(vec![n, c, t, 1])?;
let replicated = (0..scale)
.fold(None, |acc: Option<Tensor>, _| {
Some(match acc {
None => expanded.clone(),
Some(a) => a.concat(&expanded, 3).unwrap(),
})
})
.unwrap();
replicated.reshape(vec![n, c, t * scale])
}
pub fn interpolate2d(
&self,
target_h: usize,
target_w: usize,
) -> std::result::Result<Self, fuel_ir::Error> {
let dims = self.shape().dims().to_vec();
if dims.len() != 4 {
return Err(fuel_ir::Error::Msg(format!(
"interpolate2d: input must be rank 4 [N, C, H, W], got {dims:?}",
))
.bt());
}
let h = dims[2];
let w = dims[3];
if h == 0 || w == 0 || target_h == 0 || target_w == 0 {
return Err(fuel_ir::Error::Msg(
"interpolate2d: input + target spatial dims must be positive".into(),
)
.bt());
}
if target_h == h && target_w == w {
return Ok(self.clone());
}
if target_h.is_multiple_of(h) && target_w.is_multiple_of(w) && target_h / h == target_w / w
{
return self.upsample_nearest2d(target_h / h);
}
let h_idx: Vec<u32> = (0..target_h)
.map(|oi| ((oi * h) / target_h).min(h - 1) as u32)
.collect();
let w_idx: Vec<u32> = (0..target_w)
.map(|oj| ((oj * w) / target_w).min(w - 1) as u32)
.collect();
let h_idx_tensor = self.const_u32_like(h_idx, fuel_ir::Shape::from_dims(&[target_h]))?;
let w_idx_tensor = self.const_u32_like(w_idx, fuel_ir::Shape::from_dims(&[target_w]))?;
let after_h = self.index_select(2_usize, &h_idx_tensor)?;
after_h.index_select(3_usize, &w_idx_tensor)
}
pub fn interpolate1d(&self, target_t: usize) -> std::result::Result<Self, fuel_ir::Error> {
let dims = self.shape().dims().to_vec();
if dims.len() != 3 {
return Err(fuel_ir::Error::Msg(format!(
"interpolate1d: input must be rank 3 [N, C, T], got {dims:?}",
))
.bt());
}
let t = dims[2];
if t == 0 {
return Err(
fuel_ir::Error::Msg("interpolate1d: input length must be positive".into()).bt(),
);
}
if !target_t.is_multiple_of(t) {
return Err(fuel_ir::Error::Msg(format!(
"interpolate1d: target {target_t} must be integer multiple of input {t}; non-integer ratios are future work",
)).bt());
}
self.upsample_nearest1d(target_t / t)
}
pub fn pad_with_zeros<D: Dim>(
&self,
dim: D,
left: usize,
right: usize,
) -> std::result::Result<Self, fuel_ir::Error> {
self.pad_with_value(dim, left, right, 0.0)
}
pub fn pad_with_value<D: Dim>(
&self,
dim: D,
before: usize,
after: usize,
value: f32,
) -> std::result::Result<Self, fuel_ir::Error> {
let shape = self.shape();
let dim = dim.to_index(&shape, "pad_with_value")?;
let rank = shape.dims().len();
if before == 0 && after == 0 {
return Ok(self.clone());
}
let mut padding: Vec<(usize, usize)> = vec![(0, 0); rank];
padding[dim] = (before, after);
self.pad(padding, fuel_graph::PadMode::Constant, value as f64)
}
pub fn meshgrid(
args: &[&Self],
xy_indexing: bool,
) -> std::result::Result<Vec<Self>, fuel_ir::Error> {
if args.len() < 2 {
return Err(fuel_ir::Error::Msg(
"meshgrid: requires at least two rank-1 tensors".into(),
)
.bt());
}
let ordered: Vec<&Self> = if xy_indexing {
args.iter().rev().copied().collect()
} else {
args.to_vec()
};
let mut lens = Vec::with_capacity(ordered.len());
for (i, t) in ordered.iter().enumerate() {
let dims = t.shape().dims().to_vec();
if dims.len() != 1 {
return Err(fuel_ir::Error::Msg(format!(
"meshgrid: input {i} must be rank 1, got shape {dims:?}",
))
.bt());
}
lens.push(dims[0]);
}
let mut grids = Vec::with_capacity(ordered.len());
for (idx, t) in ordered.iter().enumerate() {
let mut shape = vec![1_usize; ordered.len()];
shape[idx] = lens[idx];
let placed = t.reshape(shape)?;
let mut repeats = lens.clone();
repeats[idx] = 1;
let grid = placed.repeat(repeats)?;
grids.push(grid);
}
if xy_indexing {
grids.reverse();
}
Ok(grids)
}
pub fn repeat(&self, repeats: impl Into<Shape>) -> std::result::Result<Self, fuel_ir::Error> {
let repeats = repeats.into();
let repeats: Vec<usize> = repeats.dims().to_vec();
let self_rank = self.shape().dims().len();
let target_rank = repeats.len();
let mut work = if self_rank < target_rank {
let pad_count = target_rank - self_rank;
let mut new_shape: Vec<usize> = vec![1; pad_count];
new_shape.extend_from_slice(self.shape().dims());
self.reshape(new_shape)?
} else if self_rank > target_rank {
return Err(fuel_ir::Error::Msg(format!(
"repeat: repeats rank {target_rank} smaller than tensor rank {self_rank}",
))
.bt());
} else {
self.clone()
};
for (axis, &n) in repeats.iter().enumerate() {
if n == 0 {
return Err(fuel_ir::Error::Msg(format!(
"repeat: zero repeat count at axis {axis} not supported",
))
.bt());
}
if n == 1 {
continue;
}
let base = work.clone();
for _ in 1..n {
work = work.concat(&base, axis)?;
}
}
Ok(work)
}
}
fn first_zero_tap_output_index(
dim: usize,
kernel: usize,
stride: usize,
pad: usize,
) -> Option<usize> {
let out_dim = (dim + 2 * pad).saturating_sub(kernel) / stride + 1;
(0..out_dim).find(|&o| {
let start = o as isize * stride as isize - pad as isize;
let end = start + kernel as isize;
let valid_start = start.max(0);
let valid_end = end.min(dim as isize);
valid_end <= valid_start
})
}
impl Tensor {
pub fn from_safetensors_bytes(
bytes: &[u8],
dtype: safetensors::Dtype,
shape: &[usize],
device: &Device,
) -> crate::Result<Self> {
use safetensors::Dtype;
let shape_obj = Shape::from_dims(shape);
let elem_count = shape_obj.elem_count();
let check_len = |expected: usize| -> crate::Result<()> {
if bytes.len() != expected {
crate::bail!(
"from_safetensors_bytes: expected {expected} bytes for dtype {dtype:?} \
and shape {shape:?}, got {}",
bytes.len(),
);
}
Ok(())
};
match dtype {
Dtype::F32 => {
check_len(elem_count * 4)?;
let mut data = Vec::with_capacity(elem_count);
for chunk in bytes.as_chunks::<4>().0.iter() {
data.push(f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]));
}
Self::from_f32(data, shape_obj, device)
}
Dtype::F64 => {
check_len(elem_count * 8)?;
let mut data = Vec::with_capacity(elem_count);
for chunk in bytes.as_chunks::<8>().0.iter() {
let arr: [u8; 8] = *chunk;
data.push(f64::from_le_bytes(arr));
}
Self::from_f64(data, shape_obj, device)
}
Dtype::BF16 => {
check_len(elem_count * 2)?;
let mut data = Vec::with_capacity(elem_count);
for chunk in bytes.as_chunks::<2>().0.iter() {
let raw = u16::from_le_bytes([chunk[0], chunk[1]]);
data.push(half::bf16::from_bits(raw));
}
Self::from_bf16(data, shape_obj, device)
}
Dtype::F16 => {
check_len(elem_count * 2)?;
let mut data = Vec::with_capacity(elem_count);
for chunk in bytes.as_chunks::<2>().0.iter() {
let raw = u16::from_le_bytes([chunk[0], chunk[1]]);
data.push(half::f16::from_bits(raw));
}
Self::from_f16(data, shape_obj, device)
}
Dtype::U32 => {
check_len(elem_count * 4)?;
let mut data = Vec::with_capacity(elem_count);
for chunk in bytes.as_chunks::<4>().0.iter() {
data.push(u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]));
}
Self::from_u32(data, shape_obj, device)
}
other => crate::bail!(
"from_safetensors_bytes: unsupported dtype {other:?} — extend Tensor's \
safetensors loader to handle it",
),
}
}
pub fn from_safetensors_view(
view: &safetensors::tensor::TensorView<'_>,
device: &Device,
) -> crate::Result<Self> {
Self::from_safetensors_bytes(view.data(), view.dtype(), view.shape(), device)
}
}
pub type LayerNormPair = (Arc<[f32]>, Arc<[f32]>);
pub type ConvWeightBias = (Arc<[f32]>, Arc<[f32]>);
#[derive(Debug, Clone)]
pub enum WeightStorage {
F32(Arc<[f32]>),
BF16(Arc<[half::bf16]>),
Q4_0 {
words: Arc<[u32]>,
bytes_len: usize,
in_features: usize,
out_features: usize,
},
WithLoRA {
base: Box<WeightStorage>,
lora_a: Arc<[f32]>,
lora_b: Arc<[f32]>,
rank: usize,
alpha: f32,
in_features: usize,
out_features: usize,
},
}
impl WeightStorage {
pub fn elem_count(&self) -> usize {
match self {
Self::F32(a) => a.len(),
Self::BF16(a) => a.len(),
Self::Q4_0 {
in_features,
out_features,
..
} => *in_features * *out_features,
Self::WithLoRA {
in_features,
out_features,
..
} => *in_features * *out_features,
}
}
pub fn dtype(&self) -> fuel_ir::DType {
match self {
Self::F32(_) => fuel_ir::DType::F32,
Self::BF16(_) => fuel_ir::DType::BF16,
Self::Q4_0 { .. } => fuel_ir::DType::U32,
Self::WithLoRA { base, .. } => base.dtype(),
}
}
pub fn const_like(
&self,
anchor: &Tensor,
shape: Shape,
) -> std::result::Result<Tensor, fuel_ir::Error> {
match self {
Self::F32(a) => anchor.const_f32_like(a.clone(), shape),
Self::BF16(a) => anchor.const_bf16_like(a.clone(), shape),
Self::Q4_0 { words, .. } => {
let _ = shape; anchor.const_u32_like(Arc::clone(words), Shape::from_dims(&[words.len()]))
}
Self::WithLoRA { .. } => Err(fuel_ir::Error::Msg(
"WeightStorage::WithLoRA::const_like is not supported \
— the base + LoRA update must be applied via \
apply_linear to produce the right graph structure."
.into(),
)
.bt()),
}
}
pub fn apply_linear_with_bias(
&self,
x: &Tensor,
in_features: usize,
out_features: usize,
bias: std::sync::Arc<[f32]>,
) -> std::result::Result<Tensor, fuel_ir::Error> {
if bias.len() != out_features {
return Err(fuel_ir::Error::Msg(format!(
"apply_linear_with_bias: bias len ({}) != out_features ({out_features})",
bias.len(),
))
.bt());
}
let projected = self.apply_linear(x, in_features, out_features)?;
let bias_t = projected.const_f32_like(bias, Shape::from_dims(&[out_features]))?;
projected.broadcast_add(&bias_t)
}
pub fn apply_linear(
&self,
x: &Tensor,
in_features: usize,
out_features: usize,
) -> std::result::Result<Tensor, fuel_ir::Error> {
let x_dims = x.shape().dims().to_vec();
match x_dims.last() {
Some(&k) if k == in_features => {}
_ => {
return Err(fuel_ir::Error::Msg(format!(
"apply_linear({in_features} -> {out_features}): input shape {x_dims:?} \
must have trailing dim {in_features}",
))
.bt());
}
}
match self {
Self::F32(_) | Self::BF16(_) => {
let w = self.const_like(x, Shape::from_dims(&[in_features, out_features]))?;
x.matmul(&w)
}
Self::Q4_0 {
in_features: expected_in,
out_features: expected_out,
..
} => {
if *expected_in != in_features || *expected_out != out_features {
return Err(fuel_ir::Error::Msg(format!(
"WeightStorage::Q4_0 shape mismatch: stored \
[{expected_in}, {expected_out}], requested \
[{in_features}, {out_features}]",
))
.bt());
}
let w_bytes = self.const_like(x, Shape::from_dims(&[in_features, out_features]))?;
x.qmatmul(
&w_bytes,
fuel_graph::QuantType::Q4_0,
in_features,
out_features,
)
}
Self::WithLoRA {
base,
lora_a,
lora_b,
rank,
alpha,
in_features: expected_in,
out_features: expected_out,
} => {
if *expected_in != in_features || *expected_out != out_features {
return Err(fuel_ir::Error::Msg(format!(
"WeightStorage::WithLoRA shape mismatch: stored \
[{expected_in}, {expected_out}], requested \
[{in_features}, {out_features}]",
))
.bt());
}
let base_out = base.apply_linear(x, in_features, out_features)?;
let a_t =
x.const_f32_like(Arc::clone(lora_a), Shape::from_dims(&[in_features, *rank]))?;
let b_t =
x.const_f32_like(Arc::clone(lora_b), Shape::from_dims(&[*rank, out_features]))?;
let scale = *alpha as f64 / *rank as f64;
let lora_path = Tensor {
inner: x.matmul(&a_t)?.matmul(&b_t)?.inner.mul_scalar(scale),
};
base_out.add(&lora_path)
}
}
}
pub fn apply_linear_dyn_m(
&self,
x: &Tensor,
in_features: usize,
out_features: usize,
row_count: fuel_ir::DynScalar,
) -> std::result::Result<Tensor, fuel_ir::Error> {
match self {
Self::F32(_) => {
let w = self
.const_like(x, Shape::from_dims(&[in_features, out_features]))
.expect("apply_linear_dyn_m F32 arm: const_like cannot fail for F32");
x.matmul_dyn_m(&w, row_count)
}
other => Err(fuel_ir::Error::Msg(format!(
"apply_linear_dyn_m: sparse-MoE dispatch is F32-only today, got {:?} \
weight (cast the expert weights to F32)",
other.dtype(),
))
.bt()),
}
}
pub fn with_lora(
self,
lora_a: Arc<[f32]>,
lora_b: Arc<[f32]>,
rank: usize,
alpha: f32,
in_features: usize,
out_features: usize,
) -> Self {
assert_eq!(
lora_a.len(),
in_features * rank,
"lora_a length {} does not match in_features ({in_features}) × rank ({rank}) = {}",
lora_a.len(),
in_features * rank,
);
assert_eq!(
lora_b.len(),
rank * out_features,
"lora_b length {} does not match rank ({rank}) × out_features ({out_features}) = {}",
lora_b.len(),
rank * out_features,
);
assert!(
!matches!(self, Self::WithLoRA { .. }),
"with_lora: base is already WithLoRA (nested adapters unsupported)",
);
Self::WithLoRA {
base: Box::new(self),
lora_a,
lora_b,
rank,
alpha,
in_features,
out_features,
}
}
}
impl From<Arc<[f32]>> for WeightStorage {
fn from(a: Arc<[f32]>) -> Self {
Self::F32(a)
}
}
impl From<Vec<f32>> for WeightStorage {
fn from(v: Vec<f32>) -> Self {
Self::F32(Arc::from(v))
}
}
impl From<Arc<[half::bf16]>> for WeightStorage {
fn from(a: Arc<[half::bf16]>) -> Self {
Self::BF16(a)
}
}
impl From<Vec<half::bf16>> for WeightStorage {
fn from(v: Vec<half::bf16>) -> Self {
Self::BF16(Arc::from(v))
}
}
#[derive(Debug, Clone)]
pub struct LayerWeights {
pub attn_q: WeightStorage,
pub attn_q_bias: Option<Arc<[f32]>>,
pub attn_k: WeightStorage,
pub attn_k_bias: Option<Arc<[f32]>>,
pub attn_v: WeightStorage,
pub attn_v_bias: Option<Arc<[f32]>>,
pub attn_o: WeightStorage,
pub ffn_gate: WeightStorage,
pub ffn_up: WeightStorage,
pub ffn_down: WeightStorage,
pub attn_norm_gain: Arc<[f32]>,
pub ffn_norm_gain: Arc<[f32]>,
}
pub fn invalidate_decode_pair_if_stale<C>(
session: &mut Option<crate::inference_context::DecodeSession>,
captured: &mut Option<C>,
ctx: &mut InferenceContext,
seq: usize,
max_seq_len: Option<usize>,
cache_dtype: DType,
n_layers: usize,
shape_key: u64,
kv: &dyn crate::inference_context::KvRebindSource,
drop_session: impl FnOnce(
&mut Option<crate::inference_context::DecodeSession>,
&mut InferenceContext,
),
) -> SessionDisposition {
refresh_decode_session(
session,
ctx,
seq,
max_seq_len,
cache_dtype,
n_layers,
shape_key,
kv,
|| *captured = None,
drop_session,
)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SessionDisposition {
Absent,
Kept,
Rebound,
Dropped(Option<crate::inference_context::RebindRefusal>),
}
#[allow(clippy::too_many_arguments)]
pub fn refresh_decode_session(
session: &mut Option<crate::inference_context::DecodeSession>,
ctx: &mut InferenceContext,
seq: usize,
max_seq_len: Option<usize>,
cache_dtype: DType,
n_layers: usize,
shape_key: u64,
kv: &dyn crate::inference_context::KvRebindSource,
on_invalidate: impl FnOnce(),
drop_session: impl FnOnce(
&mut Option<crate::inference_context::DecodeSession>,
&mut InferenceContext,
),
) -> SessionDisposition {
use crate::inference_context::PlanValidity;
let verdict = match (session.as_ref(), max_seq_len) {
(None, _) => return SessionDisposition::Absent,
(Some(s), Some(msl)) => {
s.validity_for(seq, msl, n_layers, cache_dtype, shape_key, kv.alloc_id())
}
(Some(_), None) => PlanValidity::Stale,
};
if verdict == PlanValidity::Valid {
return SessionDisposition::Kept;
}
on_invalidate();
let refusal = if verdict == PlanValidity::AllocationChanged {
match session.as_mut().map(|s| s.rebind_kv(kv)) {
Some(Ok(())) => return SessionDisposition::Rebound,
Some(Err(r)) => Some(r),
None => None,
}
} else {
None
};
drop_session(session, ctx);
SessionDisposition::Dropped(refusal)
}
pub fn build_decode_causal_mask(cached_len: usize, seq: usize, max_seq_len: usize) -> Vec<f32> {
let mut mask_data = vec![0.0_f32; seq * max_seq_len];
for q_idx in 0..seq {
let abs_q = cached_len + q_idx;
for k_idx in (abs_q + 1)..max_seq_len {
mask_data[q_idx * max_seq_len + k_idx] = f32::NEG_INFINITY;
}
}
mask_data
}
pub fn build_decode_causal_mask_windowed(
cached_len: usize,
seq: usize,
max_seq_len: usize,
window: usize,
) -> Vec<f32> {
let mut mask_data = vec![0.0_f32; seq * max_seq_len];
for q_idx in 0..seq {
let abs_q = cached_len + q_idx;
for k_idx in 0..max_seq_len {
if k_idx > abs_q || k_idx + window <= abs_q {
mask_data[q_idx * max_seq_len + k_idx] = f32::NEG_INFINITY;
}
}
}
mask_data
}
#[cfg(test)]
mod decode_mask_tests {
use super::{build_decode_causal_mask, build_decode_causal_mask_windowed};
fn attendable(row: &[f32]) -> Vec<u8> {
row.iter().map(|&v| u8::from(v == 0.0)).collect()
}
#[test]
fn windowed_mask_excludes_positions_older_than_the_window() {
let m = build_decode_causal_mask_windowed(3, 1, 6, 2);
assert_eq!(
attendable(&m),
vec![0, 0, 1, 1, 0, 0],
"window=2 at absolute position 3 must attend to exactly {{2, 3}}",
);
}
#[test]
fn windowed_mask_offsets_each_row_by_cached_len() {
let m = build_decode_causal_mask_windowed(1, 2, 5, 2);
assert_eq!(
attendable(&m[0..5]),
vec![1, 1, 0, 0, 0],
"row 0 = abs pos 1"
);
assert_eq!(
attendable(&m[5..10]),
vec![0, 1, 1, 0, 0],
"row 1 = abs pos 2"
);
}
#[test]
fn window_wider_than_capacity_is_byte_identical_to_the_dense_mask() {
for (cached_len, seq, max_seq_len) in [(0, 1, 8), (3, 1, 8), (0, 4, 8), (2, 3, 8)] {
let dense = build_decode_causal_mask(cached_len, seq, max_seq_len);
let windowed =
build_decode_causal_mask_windowed(cached_len, seq, max_seq_len, max_seq_len);
assert_eq!(
dense, windowed,
"window == max_seq_len must equal dense at \
(cached_len={cached_len}, seq={seq}, max_seq_len={max_seq_len})",
);
}
}
#[test]
fn decode_row_matches_the_prefill_mask_row_at_the_same_position() {
let (n, window, max_seq_len) = (6_usize, 3_usize, 10_usize);
let mut prefill = vec![0.0_f32; n * n];
for i in 0..n {
for j in 0..n {
if j > i || j + window <= i {
prefill[i * n + j] = f32::NEG_INFINITY;
}
}
}
for p in 0..n {
let decode = build_decode_causal_mask_windowed(p, 1, max_seq_len, window);
assert_eq!(
attendable(&decode[0..n]),
attendable(&prefill[p * n..(p + 1) * n]),
"decode at cached_len={p} must match prefill row {p}",
);
assert!(
decode[n..].iter().all(|&v| v == f32::NEG_INFINITY),
"the unwritten tail beyond position {n} must stay masked",
);
}
}
}
#[allow(clippy::too_many_arguments)]
pub fn flash_window_bounds(window: Option<usize>) -> (Option<usize>, Option<usize>) {
match window {
None => (None, None),
Some(w) => (Some(w.saturating_sub(1)), Some(0)),
}
}
#[allow(clippy::too_many_arguments)]
pub fn offer_flash_decode_arm_for_region(
graph: &fuel_graph::SharedGraph,
q: fuel_graph::NodeId,
k: fuel_graph::NodeId,
v: fuel_graph::NodeId,
decomposed_out: fuel_graph::NodeId,
reconverge: fuel_graph::NodeId,
softmax_scale: f32,
attended_len_sym: fuel_ir::SymId,
attn_window: Option<usize>,
softcap: Option<f32>,
cap: fuel_dispatch::decode_flash::FlashArmCapability,
) -> crate::Result<Option<fuel_graph::NodeId>> {
use fuel_dispatch::decode_flash::{DecodeFlashSpec, offer_decode_flash_arm};
let (window_size_left, window_size_right) = flash_window_bounds(attn_window);
let spec = DecodeFlashSpec {
q,
k,
v,
alibi: None,
softmax_scale,
causal: true,
window_size_left,
window_size_right,
softcap,
k_len: fuel_ir::DynScalar::Sym(attended_len_sym),
decomposed_out,
reconverge,
};
let mut g = graph.write().map_err(|_| {
fuel_ir::Error::Msg("graph lock poisoned during flash-arm offer".into()).bt()
})?;
offer_decode_flash_arm(&mut g, &spec, cap)
}
pub fn apply_affine_rms_norm(x: &Tensor, gain: &Arc<[f32]>, dim: usize, eps: f64) -> Tensor {
assert_eq!(
gain.len(),
dim,
"apply_affine_rms_norm: gain length must equal dim"
);
let normalized = x.rms_norm_last_dim(eps).unwrap();
let gain_t = x
.const_like_dtype(gain, Shape::from_dims(&[dim]), x.dtype())
.expect("apply_affine_rms_norm: activation dtype must be F32 or BF16");
normalized.broadcast_mul(&gain_t).unwrap()
}
pub fn load_tensor_as_f32(
st: &fuel_loaders::safetensors::MmapedSafetensors,
name: &str,
) -> crate::Result<Vec<f32>> {
use safetensors::Dtype;
let view = st.get(name)?;
let bytes = view.data();
match view.dtype() {
Dtype::F32 => {
let mut out = Vec::with_capacity(bytes.len() / 4);
for chunk in bytes.as_chunks::<4>().0.iter() {
out.push(f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]));
}
Ok(out)
}
Dtype::F64 => {
let mut out = Vec::with_capacity(bytes.len() / 8);
for chunk in bytes.as_chunks::<8>().0.iter() {
let arr: [u8; 8] = *chunk;
out.push(f64::from_le_bytes(arr) as f32);
}
Ok(out)
}
Dtype::BF16 => {
let mut out = Vec::with_capacity(bytes.len() / 2);
for chunk in bytes.as_chunks::<2>().0.iter() {
let raw = u16::from_le_bytes([chunk[0], chunk[1]]);
out.push(half::bf16::from_bits(raw).to_f32());
}
Ok(out)
}
Dtype::F16 => {
let mut out = Vec::with_capacity(bytes.len() / 2);
for chunk in bytes.as_chunks::<2>().0.iter() {
let raw = u16::from_le_bytes([chunk[0], chunk[1]]);
out.push(half::f16::from_bits(raw).to_f32());
}
Ok(out)
}
other => {
crate::bail!("load_tensor_as_f32: unsupported dtype {other:?} for tensor {name:?}",)
}
}
}
pub fn load_transposed_matrix(
st: &fuel_loaders::safetensors::MmapedSafetensors,
name: &str,
out_features: usize,
in_features: usize,
) -> crate::Result<Vec<f32>> {
let flat = load_tensor_as_f32(st, name)?;
if flat.len() != out_features * in_features {
crate::bail!(
"load_transposed_matrix: tensor {name:?} has {} elements, expected {} ({out_features} × {in_features})",
flat.len(),
out_features * in_features,
);
}
let mut out = vec![0.0_f32; out_features * in_features];
for i in 0..out_features {
for j in 0..in_features {
out[j * out_features + i] = flat[i * in_features + j];
}
}
Ok(out)
}
pub fn load_transposed_matrix_preserve_dtype(
st: &fuel_loaders::safetensors::MmapedSafetensors,
name: &str,
out_features: usize,
in_features: usize,
) -> crate::Result<WeightStorage> {
use safetensors::Dtype;
let view = st.get(name)?;
let bytes = view.data();
let expected = out_features * in_features;
match view.dtype() {
Dtype::BF16 => {
if bytes.len() != expected * 2 {
crate::bail!(
"load_transposed_matrix_preserve_dtype: bf16 tensor {name:?} has {} bytes, expected {}",
bytes.len(),
expected * 2,
);
}
let mut out = vec![half::bf16::ZERO; expected];
for i in 0..out_features {
for j in 0..in_features {
let src_off = (i * in_features + j) * 2;
let bits = u16::from_le_bytes([bytes[src_off], bytes[src_off + 1]]);
out[j * out_features + i] = half::bf16::from_bits(bits);
}
}
Ok(WeightStorage::BF16(Arc::from(out)))
}
_ => {
let flat = load_transposed_matrix(st, name, out_features, in_features)?;
Ok(WeightStorage::F32(Arc::from(flat)))
}
}
}
#[derive(Debug, Clone, Copy, Default)]
pub enum SamplingStrategy {
#[default]
Greedy,
Temperature { temp: f32, seed: u64 },
}
pub fn sample_logits(logits: &[f32], strategy: SamplingStrategy, rng_state: &mut u64) -> u32 {
match strategy {
SamplingStrategy::Greedy => {
let (i, _) = logits
.iter()
.enumerate()
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
.expect("sample_logits: empty logits");
i as u32
}
SamplingStrategy::Temperature { temp, .. } => {
let inv_temp = if temp == 0.0 { 1.0 } else { 1.0 / temp };
let scaled: Vec<f32> = logits.iter().map(|&x| x * inv_temp).collect();
let max = scaled.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let exp: Vec<f32> = scaled.iter().map(|&x| (x - max).exp()).collect();
let sum: f32 = exp.iter().sum();
let probs: Vec<f32> = exp.iter().map(|&x| x / sum).collect();
sample_multinomial(&probs, rng_state)
}
}
}
pub fn spec_argmax(logits: &[f32]) -> u32 {
let mut best = 0;
let mut best_v = logits[0];
for (i, &v) in logits.iter().enumerate().skip(1) {
if v > best_v {
best_v = v;
best = i;
}
}
best as u32
}
pub fn spec_softmax_temp(logits: &[f32], temp: f32) -> Vec<f32> {
let inv_t = if temp == 0.0 { 1.0 } else { 1.0 / temp };
let max = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let exp: Vec<f32> = logits.iter().map(|&x| ((x - max) * inv_t).exp()).collect();
let sum: f32 = exp.iter().sum();
exp.iter().map(|&x| x / sum).collect()
}
pub fn spec_next_u01(state: &mut u64) -> f32 {
*state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
(*state >> 32) as f32 / u32::MAX as f32
}
pub fn spec_sample_cat(probs: &[f32], state: &mut u64) -> u32 {
let u = spec_next_u01(state);
let mut cum = 0.0_f32;
for (i, &p) in probs.iter().enumerate() {
cum += p;
if u <= cum {
return i as u32;
}
}
(probs.len() - 1) as u32
}
pub fn sample_multinomial(probs: &[f32], rng_state: &mut u64) -> u32 {
*rng_state = rng_state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
let u = (*rng_state >> 32) as f32 / u32::MAX as f32;
let mut cumulative = 0.0_f32;
for (i, &p) in probs.iter().enumerate() {
cumulative += p;
if u <= cumulative {
return i as u32;
}
}
(probs.len() - 1) as u32
}
#[cfg(test)]
mod lora_tests {
use super::*;
use crate::Shape;
#[test]
fn with_lora_matches_manual_base_plus_lora() {
let in_f = 4;
let out_f = 3;
let rank = 2;
let alpha = 8.0_f32;
let anchor =
Tensor::from_f32(vec![0.0_f32; 1], Shape::from_dims(&[1]), &Device::cpu()).unwrap();
let base_vec: Vec<f32> = (0..in_f * out_f).map(|i| (i as f32) * 0.1).collect();
let lora_a_vec: Vec<f32> = (0..in_f * rank).map(|i| (i as f32) * 0.05).collect();
let lora_b_vec: Vec<f32> = (0..rank * out_f).map(|i| (i as f32) * 0.02).collect();
let ws = WeightStorage::F32(Arc::from(base_vec.clone())).with_lora(
Arc::from(lora_a_vec.clone()),
Arc::from(lora_b_vec.clone()),
rank,
alpha,
in_f,
out_f,
);
let batch = 2;
let x_data: Vec<f32> = (0..batch * in_f).map(|i| (i as f32) * 0.1 + 0.5).collect();
let x = anchor
.const_f32_like(x_data.clone(), Shape::from_dims(&[batch, in_f]))
.unwrap();
let y = ws.apply_linear(&x, in_f, out_f).unwrap();
let got = y.realize_f32().to_vec();
let mut expected = vec![0.0_f32; batch * out_f];
for b in 0..batch {
for j in 0..out_f {
let mut acc = 0.0_f32;
for k in 0..in_f {
acc += x_data[b * in_f + k] * base_vec[k * out_f + j];
}
let scale = alpha as f64 / rank as f64;
for r in 0..rank {
let mut xar = 0.0_f32;
for k in 0..in_f {
xar += x_data[b * in_f + k] * lora_a_vec[k * rank + r];
}
acc += (xar * lora_b_vec[r * out_f + j]) * scale as f32;
}
expected[b * out_f + j] = acc;
}
}
for (i, (&e, &g)) in expected.iter().zip(got.iter()).enumerate() {
let diff = (e - g).abs();
assert!(
diff <= 1e-4,
"LoRA mismatch at {i}: expected {e}, got {g} (diff {diff})",
);
}
}
#[test]
#[should_panic(expected = "lora_a length")]
fn with_lora_rejects_mismatched_a_shape() {
let ws = WeightStorage::F32(Arc::from(vec![0.0_f32; 12])); let bad_a = Arc::from(vec![0.0_f32; 3]); let b = Arc::from(vec![0.0_f32; 6]); let _ = ws.with_lora(bad_a, b, 2, 8.0, 4, 3);
}
#[test]
fn apply_linear_errors_on_trailing_dim_mismatch() {
let ws = WeightStorage::F32(Arc::from(vec![0.0_f32; 12])); let x = Tensor::from_f32(
vec![0.0_f32; 10],
Shape::from_dims(&[2, 5]), &Device::cpu(),
)
.unwrap();
let err = ws
.apply_linear(&x, 4, 3)
.expect_err("must not panic, must Err");
let msg = err.to_string();
assert!(
msg.contains("trailing dim 4"),
"error should name the expected trailing dim, got: {msg}",
);
}
#[test]
fn apply_linear_errors_on_q4_0_stored_shape_mismatch() {
let ws = WeightStorage::Q4_0 {
words: Arc::from(vec![0_u32; 8]),
bytes_len: 18, in_features: 4,
out_features: 4,
};
let x =
Tensor::from_f32(vec![0.0_f32; 8], Shape::from_dims(&[2, 4]), &Device::cpu()).unwrap();
let err = ws
.apply_linear(&x, 4, 7)
.expect_err("must not panic, must Err");
let msg = err.to_string();
assert!(
msg.contains("Q4_0 shape mismatch"),
"error should name the stored/requested disagreement, got: {msg}",
);
}
#[test]
fn apply_linear_with_bias_errors_on_wrong_bias_length() {
let ws = WeightStorage::F32(Arc::from(vec![0.0_f32; 12])); let x =
Tensor::from_f32(vec![0.0_f32; 8], Shape::from_dims(&[2, 4]), &Device::cpu()).unwrap();
let bias: Arc<[f32]> = Arc::from(vec![0.0_f32; 2]); let err = ws
.apply_linear_with_bias(&x, 4, 3, bias)
.expect_err("wrong-length bias must Err in EVERY profile, not just debug");
assert!(
err.to_string().contains("bias len"),
"error should name the bias length, got: {err}",
);
}
}
#[cfg(test)]
mod safetensors_bridge_tests {
use super::*;
#[test]
fn from_safetensors_bytes_round_trip_f32() {
let original: Vec<f32> = vec![1.0, -2.5, 3.25, 0.0];
let mut bytes = Vec::with_capacity(original.len() * 4);
for &v in &original {
bytes.extend_from_slice(&v.to_le_bytes());
}
let t =
Tensor::from_safetensors_bytes(&bytes, safetensors::Dtype::F32, &[4], &Device::cpu())
.unwrap();
assert_eq!(t.shape().dims(), &[4]);
assert_eq!(t.dtype(), DType::F32);
assert_eq!(t.realize_f32(), original);
}
#[test]
fn from_safetensors_bytes_round_trip_bf16() {
let original_f32: Vec<f32> = vec![0.5, -1.0, 2.0, 4.0];
let bf16_vec: Vec<half::bf16> = original_f32
.iter()
.map(|&v| half::bf16::from_f32(v))
.collect();
let mut bytes = Vec::with_capacity(bf16_vec.len() * 2);
for b in &bf16_vec {
bytes.extend_from_slice(&b.to_bits().to_le_bytes());
}
let t =
Tensor::from_safetensors_bytes(&bytes, safetensors::Dtype::BF16, &[4], &Device::cpu())
.unwrap();
assert_eq!(t.dtype(), DType::BF16);
let realized = t.realize_bf16();
assert_eq!(realized, bf16_vec);
}
#[test]
fn from_safetensors_bytes_rejects_wrong_byte_count() {
let bad_bytes = vec![0_u8; 4];
let result = Tensor::from_safetensors_bytes(
&bad_bytes,
safetensors::Dtype::F32,
&[3],
&Device::cpu(),
);
assert!(result.is_err());
}
}
#[cfg(test)]
mod phase_a1_wrapper_tests {
use super::*;
fn cpu_f32(data: Vec<f32>, shape: &[usize]) -> Tensor {
Tensor::from_f32(data, shape.to_vec(), &Device::cpu()).unwrap()
}
#[test]
fn unsqueeze_adds_size_one_dim() {
let t = cpu_f32(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]);
let out = t.unsqueeze(0_usize).unwrap();
assert_eq!(out.shape().dims(), &[1, 2, 2]);
assert_eq!(out.realize_f32(), vec![1.0, 2.0, 3.0, 4.0]);
}
#[test]
fn unsqueeze_errors_out_of_bounds() {
let t = cpu_f32(vec![1.0, 2.0], &[2]);
assert!(t.unsqueeze(0_usize).is_ok());
assert!(t.unsqueeze(1_usize).is_ok());
assert!(t.unsqueeze(2_usize).is_err());
}
#[test]
fn try_reshape_errors_on_size_mismatch() {
let t = cpu_f32(vec![1.0; 6], &[2, 3]);
assert!(t.reshape(vec![3, 2]).is_ok());
assert!(t.reshape(vec![2, 2]).is_err());
}
#[test]
fn try_permute_validates_axes() {
let t = cpu_f32(vec![0.0; 24], &[2, 3, 4]);
assert!(t.permute((2_usize, 0_usize, 1_usize)).is_ok());
assert!(t.permute([2_usize, 0, 1]).is_ok());
assert!(t.permute([0_usize, 1]).is_err()); assert!(t.permute([0_usize, 0, 1]).is_err()); }
#[test]
fn try_transpose_requires_rank_two_plus() {
let scalar = cpu_f32(vec![1.0], &[1]);
let _ = scalar.transpose();
}
#[test]
fn triu_tril_shape_preserved() {
let t = cpu_f32(vec![1.0; 9], &[3, 3]);
let upper = t.triu(0).unwrap();
let lower = t.tril(0).unwrap();
assert_eq!(upper.shape().dims(), &[3, 3]);
assert_eq!(lower.shape().dims(), &[3, 3]);
assert_eq!(
lower.realize_f32(),
vec![1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0]
);
}
#[test]
fn triu_rejects_rank_one() {
let t = cpu_f32(vec![1.0, 2.0], &[2]);
assert!(t.triu(0).is_err());
}
#[test]
fn log_softmax_last_dim_shape_preserved() {
let t = cpu_f32(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]);
let out = t.log_softmax_last_dim().unwrap();
assert_eq!(out.shape().dims(), &[2, 2]);
for v in out.realize_f32() {
assert!(v <= 0.0 + 1e-6, "log_softmax produced positive value: {v}");
}
}
#[test]
fn softmax_general_axis_matches_hand_computed_rank3_middle_axis() {
let data: Vec<f32> = vec![
-1.0, 0.0, 0.0, -3.0, -2.0, -1.0, -2.0, -1.0, -1.0, -1.0, 0.0, 0.0, ];
let t = cpu_f32(data, &[2, 3, 2]);
let out = t.softmax(1_usize).unwrap();
assert_eq!(out.shape().dims(), &[2, 3, 2]);
let v = out.realize_f32();
let lane_softmax = |xs: [f32; 3]| -> [f32; 3] {
let m = xs.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let exps = [(xs[0] - m).exp(), (xs[1] - m).exp(), (xs[2] - m).exp()];
let s = exps[0] + exps[1] + exps[2];
[exps[0] / s, exps[1] / s, exps[2] / s]
};
let lanes: [[f32; 3]; 4] = [
[-1.0, 0.0, -2.0], [0.0, -3.0, -1.0], [-2.0, -1.0, 0.0], [-1.0, -1.0, 0.0], ];
let refs: Vec<[f32; 3]> = lanes.iter().map(|l| lane_softmax(*l)).collect();
for b in 0..2 {
for c in 0..2 {
let lane_ix = b * 2 + c;
for r in 0..3 {
let got = v[b * 6 + r * 2 + c];
let want = refs[lane_ix][r];
assert!(
(got - want).abs() < 1e-6,
"softmax mismatch at (b={b}, r={r}, c={c}): got {got}, want {want}",
);
}
let sum: f32 = (0..3).map(|r| v[b * 6 + r * 2 + c]).sum();
assert!((sum - 1.0).abs() < 1e-6, "lane (b={b},c={c}) sums to {sum}");
}
}
}
#[test]
fn softmax_last_axis_matches_softmax_last_dim() {
let data: Vec<f32> = vec![
1.0, 2.0, -1.0, 0.5, 0.0, -2.0, 3.0, 1.5, 4.0, -1.0, 2.0, 0.0, -3.0, 0.25, 0.75, 1.0,
];
let t = cpu_f32(data, &[2, 2, 4]);
let via_general = t.softmax(2_usize).unwrap();
let via_fused = t.softmax_last_dim().unwrap();
assert_eq!(via_general.shape().dims(), via_fused.shape().dims());
let g = via_general.realize_f32();
let f = via_fused.realize_f32();
assert_eq!(g.len(), f.len());
for (i, (a, b)) in g.iter().zip(f.iter()).enumerate() {
assert!(
(a - b).abs() < 1e-6,
"softmax (general axis=2) vs softmax_last_dim diverge at {i}: {a} vs {b}",
);
}
}
#[test]
fn log_softmax_general_axis_matches_hand_computed_rank3_middle_axis() {
let data: Vec<f32> = vec![
-1.0, 0.0, 0.0, -3.0, -2.0, -1.0, -2.0, -1.0, -1.0, -1.0, 0.0, 0.0,
];
let t = cpu_f32(data, &[2, 3, 2]);
let out = t.log_softmax(1_usize).unwrap();
assert_eq!(out.shape().dims(), &[2, 3, 2]);
let v = out.realize_f32();
let lane_log_softmax = |xs: [f32; 3]| -> [f32; 3] {
let m = xs.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let shifted = [xs[0] - m, xs[1] - m, xs[2] - m];
let lse = (shifted[0].exp() + shifted[1].exp() + shifted[2].exp()).ln();
[shifted[0] - lse, shifted[1] - lse, shifted[2] - lse]
};
let lanes: [[f32; 3]; 4] = [
[-1.0, 0.0, -2.0],
[0.0, -3.0, -1.0],
[-2.0, -1.0, 0.0],
[-1.0, -1.0, 0.0],
];
let refs: Vec<[f32; 3]> = lanes.iter().map(|l| lane_log_softmax(*l)).collect();
for b in 0..2 {
for c in 0..2 {
let lane_ix = b * 2 + c;
for r in 0..3 {
let got = v[b * 6 + r * 2 + c];
let want = refs[lane_ix][r];
assert!(
(got - want).abs() < 1e-6,
"log_softmax mismatch at (b={b}, r={r}, c={c}): got {got}, want {want}",
);
assert!(got <= 1e-6, "log_softmax produced positive value: {got}");
}
}
}
}
#[test]
fn log_softmax_last_axis_matches_log_softmax_last_dim() {
let data: Vec<f32> = vec![
1.0, 2.0, -1.0, 0.5, 0.0, -2.0, 3.0, 1.5, 4.0, -1.0, 2.0, 0.0, -3.0, 0.25, 0.75, 1.0,
];
let t = cpu_f32(data, &[2, 2, 4]);
let via_general = t.log_softmax(2_usize).unwrap();
let via_fused = t.log_softmax_last_dim().unwrap();
assert_eq!(via_general.shape().dims(), via_fused.shape().dims());
let g = via_general.realize_f32();
let f = via_fused.realize_f32();
assert_eq!(g.len(), f.len());
for (i, (a, b)) in g.iter().zip(f.iter()).enumerate() {
assert!(
(a - b).abs() < 1e-6,
"log_softmax (general axis=2) vs log_softmax_last_dim diverge at {i}: {a} vs {b}",
);
}
}
#[test]
fn argmin_dim_drops_reduced_axis() {
let t = cpu_f32(vec![3.0, 1.0, 2.0, 0.5, 5.0, 4.0], &[2, 3]);
let out = t.argmin_dim(1_usize).unwrap();
assert_eq!(out.shape().dims(), &[2]);
assert_eq!(out.dtype(), DType::U32);
assert_eq!(out.realize_u32(), vec![1, 0]);
}
#[test]
fn masked_fill_smoke() {
let t = cpu_f32(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]);
let probe = t
.const_f32_like(vec![0.0, 1.0, 1.0, 0.0], vec![2, 2])
.unwrap();
let threshold = t.const_f32_like(vec![0.5; 4], vec![2, 2]).unwrap();
let mask = probe.gt(&threshold).unwrap(); let out = t.masked_fill(&mask, fuel_ir::Scalar::F32(-9.0)).unwrap();
assert_eq!(out.realize_f32(), vec![1.0, -9.0, -9.0, 4.0]);
}
#[test]
fn bool_is_distinguishable_from_u8_end_to_end() {
let t = cpu_f32(vec![0.0, 5.0, 3.0, 0.0], &[4]);
let thr = t.const_f32_like(vec![0.5; 4], vec![4]).unwrap();
let mask = t.gt(&thr).unwrap();
assert_eq!(mask.dtype(), DType::Bool, "gt yields Bool, not U8");
assert_eq!(mask.realize_u8(), vec![0, 1, 1, 0], "mask bytes are 0/1");
let u8_vals = t
.const_f32_like(vec![0.0, 5.0, 3.0, 0.0], vec![4])
.unwrap()
.to_dtype(DType::U8)
.unwrap();
assert_eq!(u8_vals.dtype(), DType::U8);
assert_eq!(
u8_vals.realize_u8(),
vec![0, 5, 3, 0],
"negative control: the U8 source must really hold 5 and 3, else the \
conversion assertion below cannot discriminate"
);
let as_bool = u8_vals.to_dtype(DType::Bool).unwrap();
assert_eq!(as_bool.dtype(), DType::Bool);
assert_eq!(
as_bool.realize_u8(),
vec![0, 1, 1, 0],
"U8->Bool must CONVERT (5 -> 1), not reinterpret (5 -> 5)"
);
assert_eq!(
as_bool.to_dtype(DType::F32).unwrap().realize_f32(),
vec![0.0, 1.0, 1.0, 0.0],
"Bool->F32 yields 0.0/1.0"
);
assert_eq!(
mask.to_dtype(DType::U8).unwrap().realize_u8(),
vec![0, 1, 1, 0]
);
let err = t.masked_fill(&u8_vals, fuel_ir::Scalar::F32(-9.0));
assert!(
err.is_err(),
"a U8 mask must be REJECTED by masked_fill — Bool and U8 are \
byte-identical, so accepting it would 'work' and erase the distinction"
);
}
#[test]
fn float_to_bool_lowering_preserves_ieee_semantics() {
let t = cpu_f32(
vec![
0.0,
-0.0,
5.0,
-3.0,
f32::NAN,
f32::INFINITY,
f32::NEG_INFINITY,
],
&[7],
);
let b = t.to_dtype(DType::Bool).unwrap();
assert_eq!(b.dtype(), DType::Bool, "to_dtype(Bool) still yields Bool");
assert_eq!(
b.realize_u8(),
vec![0, 0, 1, 1, 1, 1, 1],
"0 and -0.0 are false; every other value including NaN and both \
infinities is true",
);
}
#[test]
fn index_add_smoke() {
let base = cpu_f32(vec![1.0, 1.0, 1.0, 1.0], &[2, 2]);
let src = base
.const_f32_like(vec![10.0, 20.0, 30.0, 40.0], vec![2, 2])
.unwrap();
let indices = base.const_u32_like(vec![0_u32, 0_u32], vec![2]).unwrap();
let out = base.index_add(0, &indices, &src).unwrap();
assert_eq!(out.shape().dims(), &[2, 2]);
let v = out.realize_f32();
assert_eq!(v[0], 41.0); assert_eq!(v[1], 61.0); assert_eq!(v[2], 1.0);
assert_eq!(v[3], 1.0);
}
#[test]
fn scatter_add_smoke() {
let base = cpu_f32(vec![0.0, 0.0, 0.0, 0.0], &[2, 2]);
let src = base
.const_f32_like(vec![5.0, 6.0, 7.0, 8.0], vec![2, 2])
.unwrap();
let indices = base
.const_u32_like(vec![0_u32, 1_u32, 1_u32, 0_u32], vec![2, 2])
.unwrap();
let out = base.scatter_add(0, &indices, &src).unwrap();
assert_eq!(out.shape().dims(), &[2, 2]);
}
#[test]
fn inplace_activations_compile_and_run() {
let t = cpu_f32(vec![-1.0, 0.5, -3.0, 2.0], &[4]);
let r = cpu_f32(vec![-1.0, 0.5, -3.0, 2.0], &[4]).relu_inplace();
let _ = cpu_f32(vec![-1.0, 0.5, -3.0, 2.0], &[4]).silu_inplace();
let _ = cpu_f32(vec![-1.0, 0.5, -3.0, 2.0], &[4]).gelu_tanh_inplace();
let _ = cpu_f32(vec![-1.0, 0.5, -3.0, 2.0], &[4]).tanh_inplace();
let _ = cpu_f32(vec![-1.0, 0.5, -3.0, 2.0], &[4]).sigmoid_inplace();
let _ = t.affine_inplace(2.0, 1.0);
let v = r.realize_f32();
assert_eq!(v, vec![0.0, 0.5, 0.0, 2.0]);
}
#[test]
fn const_f64_like_round_trips() {
let anchor = cpu_f32(vec![0.0], &[1]);
let t = anchor.const_f64_like(vec![1.5, 2.5, 3.5], vec![3]).unwrap();
assert_eq!(t.shape().dims(), &[3]);
assert_eq!(t.dtype(), DType::F64);
assert_eq!(t.realize_f64(), vec![1.5, 2.5, 3.5]);
}
#[test]
fn const_i64_like_round_trips() {
let anchor = cpu_f32(vec![0.0], &[1]);
let t = anchor.const_i64_like(vec![-1_i64, 2, -3], vec![3]).unwrap();
assert_eq!(t.shape().dims(), &[3]);
assert_eq!(t.dtype(), DType::I64);
}
#[test]
fn on_device_smoke_cpu() {
let t = cpu_f32(vec![1.0, 2.0], &[2]);
let pinned = t.on_device(&Device::cpu());
assert_eq!(pinned.realize_f32(), vec![1.0, 2.0]);
}
#[test]
fn copy_to_device_same_device_round_trips() {
let t = cpu_f32(vec![1.0, 2.0, 3.0], &[3]);
let copied = t.copy_to_device(&Device::cpu());
assert_eq!(copied.realize_f32(), vec![1.0, 2.0, 3.0]);
}
}
#[cfg(test)]
mod phase_a2_composite_tests {
use super::*;
fn cpu_f32(data: Vec<f32>, shape: &[usize]) -> Tensor {
Tensor::from_f32(data, shape.to_vec(), &Device::cpu()).unwrap()
}
#[test]
fn transpose_last_two_swaps_last_two_dims() {
let t = cpu_f32(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]);
let out = t.transpose_last_two().unwrap();
assert_eq!(out.shape().dims(), &[3, 2]);
assert_eq!(out.realize_f32(), vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0]);
}
#[test]
fn t_is_alias_of_transpose_last_two() {
let t = cpu_f32(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]);
assert_eq!(
t.t().unwrap().realize_f32(),
t.transpose_last_two().unwrap().realize_f32()
);
}
#[test]
fn transpose_dims_swaps_arbitrary_axes() {
let t = cpu_f32((0..24).map(|i| i as f32).collect(), &[2, 3, 4]);
let out = t.transpose_dims(0, 2).unwrap();
assert_eq!(out.shape().dims(), &[4, 3, 2]);
}
#[test]
fn transpose_dims_identity_when_same_axis() {
let t = cpu_f32(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]);
let out = t.transpose_dims(0, 0).unwrap();
assert_eq!(out.realize_f32(), vec![1.0, 2.0, 3.0, 4.0]);
}
#[test]
fn transpose_dims_errors_out_of_bounds() {
let t = cpu_f32(vec![0.0; 6], &[2, 3]);
assert!(t.transpose_dims(0, 2).is_err());
assert!(t.transpose_dims(5, 0).is_err());
}
#[test]
fn flatten_merges_middle_dims() {
let t = cpu_f32(vec![0.0; 24], &[2, 3, 4]);
let out = t.flatten(0, 1).unwrap();
assert_eq!(out.shape().dims(), &[6, 4]);
}
#[test]
fn flatten_to_merges_leading_dims() {
let t = cpu_f32(vec![0.0; 24], &[2, 3, 4]);
let out = t.flatten_to(1).unwrap();
assert_eq!(out.shape().dims(), &[6, 4]);
}
#[test]
fn flatten_from_merges_trailing_dims() {
let t = cpu_f32(vec![0.0; 24], &[2, 3, 4]);
let out = t.flatten_from(1).unwrap();
assert_eq!(out.shape().dims(), &[2, 12]);
}
#[test]
fn flatten_all_produces_rank_one() {
let t = cpu_f32(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]);
let out = t.flatten_all().unwrap();
assert_eq!(out.shape().dims(), &[6]);
assert_eq!(out.realize_f32(), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
}
#[test]
fn rope_batched_matches_single_position_rope_row_by_row() {
let lcg = |n: usize, seed: u32| -> Vec<f32> {
let mut s = seed;
(0..n)
.map(|_| {
s = s.wrapping_mul(1103515245).wrapping_add(12345);
((s >> 16) as u16 as f32 / 65535.0 - 0.5) * 0.5
})
.collect()
};
let (b, heads, hd) = (3usize, 4usize, 8usize);
let base = 10000.0f64;
let positions = [0usize, 7, 3];
let x_data = lcg(b * heads * hd, 11);
let x = cpu_f32(x_data.clone(), &[b, heads, 1, hd]);
let (cos, sin) = x.rope_tables_const_batched(base, &positions, hd);
let got = x.rope_batched(&cos, &sin).unwrap().realize_f32();
let per_row = heads * hd;
for (row, &pos) in positions.iter().enumerate() {
let row_data = x_data[row * per_row..(row + 1) * per_row].to_vec();
let xr = cpu_f32(row_data, &[1, heads, 1, hd]);
let (c, s) = xr.rope_tables_const(base, pos, 1, hd);
let want = xr
.rope_with_tables_decomposed(&c, &s)
.unwrap()
.realize_f32();
assert_eq!(
&got[row * per_row..(row + 1) * per_row],
&want[..],
"row {row} (position {pos}) != decomposed RoPE at that position",
);
}
let shared = [positions[0]; 3];
let (c0, s0) = x.rope_tables_const_batched(base, &shared, hd);
let wrong = x.rope_batched(&c0, &s0).unwrap().realize_f32();
for (row, &pos) in positions.iter().enumerate().skip(1) {
assert_ne!(
&wrong[row * per_row..(row + 1) * per_row],
&got[row * per_row..(row + 1) * per_row],
"CONTROL: row {row} at shared position 0 matched its own position \
{pos} — per-row RoPE not applied, or comparison blind to it",
);
}
}
#[test]
fn flatten_errors_on_bad_range() {
let t = cpu_f32(vec![0.0; 6], &[2, 3]);
assert!(t.flatten(0, 5).is_err());
assert!(t.flatten(2, 1).is_err()); }
#[test]
fn split_heads_then_merge_heads_round_trip() {
let x = cpu_f32((0..12).map(|i| i as f32).collect(), &[1, 2, 6]);
let split = x.split_heads(2, 3).unwrap();
assert_eq!(split.shape().dims(), &[1, 2, 2, 3]);
let merged = split.merge_heads().unwrap();
assert_eq!(merged.shape().dims(), &[1, 2, 6]);
let m = merged.realize_f32();
let original = x.realize_f32();
for (a, b) in m.iter().zip(original.iter()) {
assert!((a - b).abs() < 1e-7, "{a} vs {b}");
}
}
#[test]
fn rope_tables_const_shapes_and_anchoring() {
let anchor = cpu_f32(vec![0.0_f32; 8], &[1, 4, 2]);
let (cos, sin) = anchor.rope_tables_const(10_000.0, 0, 4, 8);
assert_eq!(cos.shape().dims(), &[4, 8]);
assert_eq!(sin.shape().dims(), &[4, 8]);
let cv = cos.realize_f32();
let sv = sin.realize_f32();
for i in 0..8 {
assert!((cv[i] - 1.0).abs() < 1e-5, "cos[0,{i}] = {}", cv[i]);
assert!(sv[i].abs() < 1e-5, "sin[0,{i}] = {}", sv[i]);
}
}
#[test]
fn embed_tokens_shape_and_lookup() {
let vocab_size = 5;
let hidden = 3;
let table: Vec<f32> = (0..vocab_size)
.flat_map(|i| vec![i as f32, i as f32 + 0.5, i as f32 + 1.0])
.collect();
let tokens = vec![1_u32, 3, 0];
let out = Tensor::embed_tokens(
std::sync::Arc::from(table),
vocab_size,
hidden,
&tokens,
&crate::Device::cpu(),
)
.unwrap();
assert_eq!(out.shape().dims(), &[1, 3, hidden]);
let v = out.realize_f32();
let want = [
1.0_f32, 1.5, 2.0, 3.0, 3.5, 4.0, 0.0, 0.5, 1.0, ];
for (i, (&got, &exp)) in v.iter().zip(want.iter()).enumerate() {
assert!((got - exp).abs() < 1e-5, "row {i}: got={got} want={exp}");
}
}
#[test]
fn embed_tokens_anchored_lives_on_receiver_graph() {
let vocab_size = 4;
let hidden = 2;
let table: Vec<f32> = (0..vocab_size)
.flat_map(|i| vec![i as f32, i as f32 * 2.0])
.collect();
let table_arc: std::sync::Arc<[f32]> = std::sync::Arc::from(table);
let tokens = vec![2_u32, 1];
let anchor = cpu_f32(vec![0.0_f32], &[1]);
let embedded = anchor
.embed_tokens_anchored(
std::sync::Arc::clone(&table_arc),
vocab_size,
hidden,
&tokens,
)
.unwrap();
assert_eq!(embedded.shape().dims(), &[1, 2, hidden]);
let one_scaled = anchor
.const_f32_like(std::sync::Arc::from(vec![1.0_f32]), Shape::from_dims(&[1]))
.unwrap();
let _ = embedded
.add(
&one_scaled
.reshape(Shape::from_dims(&[1, 1, 1]))
.unwrap()
.broadcast_to(Shape::from_dims(&[1, 2, hidden]))
.unwrap(),
)
.unwrap();
let v = embedded.realize_f32();
let want = [
2.0_f32, 4.0, 1.0, 2.0, ];
for (i, (&got, &exp)) in v.iter().zip(want.iter()).enumerate() {
assert!((got - exp).abs() < 1e-5, "row {i}: got={got} want={exp}");
}
}
#[test]
fn embed_tokens_empty_returns_error() {
let r = Tensor::embed_tokens(
std::sync::Arc::from(vec![0.0_f32]),
1,
1,
&[],
&crate::Device::cpu(),
);
assert!(r.is_err());
}
#[test]
fn softcap_matches_tanh_form() {
let x = cpu_f32(vec![0.0_f32, 5.0, -5.0, 30.0], &[1, 4]);
let capped = x.softcap(10.0).realize_f32();
let expect = [0.0_f32, 4.6212, -4.6212, 9.9505];
for (i, (&got, &want)) in capped.iter().zip(expect.iter()).enumerate() {
assert!(
(got - want).abs() < 1e-3,
"softcap[{i}] got={got} want={want}"
);
}
}
#[test]
fn softcap_optional_none_returns_input_unchanged() {
let x = cpu_f32(vec![1.0_f32, -2.0, 30.0], &[1, 3]);
let out = x.softcap_optional(None).realize_f32();
let expect = [1.0_f32, -2.0, 30.0];
for (g, w) in out.iter().zip(expect.iter()) {
assert!((g - w).abs() < 1e-6);
}
let out = x.softcap_optional(Some(0.0)).realize_f32();
for (g, w) in out.iter().zip(expect.iter()) {
assert!((g - w).abs() < 1e-6);
}
}
#[test]
fn rope_partial_full_dim_matches_rope_with_tables() {
let qk = cpu_f32(vec![1.0_f32; 1 * 1 * 2 * 4], &[1, 1, 2, 4]);
let (cos, sin) = qk.rope_tables_const(10_000.0, 0, 2, 4);
let via_partial = qk.rope_partial(&cos, &sin, 4).unwrap();
let via_full = qk.rope_with_tables(&cos, &sin).unwrap();
let a = via_partial.realize_f32();
let b = via_full.realize_f32();
for (i, (av, bv)) in a.iter().zip(b.iter()).enumerate() {
assert!((av - bv).abs() < 1e-6, "{i}: {av} vs {bv}");
}
}
#[test]
fn rope_partial_pass_through_suffix_unchanged() {
let qk = cpu_f32(
vec![
1.0_f32, 2.0, 3.0, 4.0,
],
&[1, 1, 1, 4],
);
let (cos, sin) = qk.rope_tables_const(10_000.0, 0, 1, 2);
let out = qk.rope_partial(&cos, &sin, 2).unwrap().realize_f32();
for (i, &v) in out.iter().enumerate() {
let expect = [1.0_f32, 2.0, 3.0, 4.0][i];
assert!((v - expect).abs() < 1e-6, "{i}: {v} != {expect}");
}
}
#[test]
fn add_optional_trailing_bias_none_returns_input_unchanged() {
let a = cpu_f32(vec![1.0_f32, 2.0, 3.0], &[1, 3]);
let original = a.realize_f32();
let out = a.add_optional_trailing_bias(None).unwrap();
assert_eq!(out.realize_f32(), original);
}
#[test]
fn add_optional_trailing_bias_some_applies_add() {
let a = cpu_f32(vec![1.0_f32, 2.0, 3.0], &[1, 3]);
let bias = std::sync::Arc::<[f32]>::from(vec![10.0_f32, 20.0, 30.0]);
let out = a.add_optional_trailing_bias(Some(&bias)).unwrap();
assert_eq!(out.realize_f32(), vec![11.0, 22.0, 33.0]);
}
#[test]
fn add_trailing_bias_broadcasts_across_leading_dims() {
let x = cpu_f32(vec![1.0_f32, 2.0, 3.0, 10.0, 20.0, 30.0], &[2, 3]);
let bias = std::sync::Arc::<[f32]>::from(vec![100.0_f32, 200.0, 300.0]);
let out = x.add_trailing_bias(bias).unwrap();
assert_eq!(out.shape().dims(), &[2, 3]);
let v = out.realize_f32();
assert_eq!(v, vec![101.0, 202.0, 303.0, 110.0, 220.0, 330.0]);
}
#[test]
fn rms_norm_affine_with_offset_adds_offset_to_each_gain() {
let x = cpu_f32(vec![1.0_f32, 2.0, 3.0], &[1, 3]);
let gain_raw: [f32; 3] = [-0.5, 0.0, 0.5];
let via_offset = x.rms_norm_affine_with_offset(&gain_raw, 1.0, 1e-6).unwrap();
let gain_shifted = std::sync::Arc::<[f32]>::from(vec![0.5_f32, 1.0, 1.5]);
let via_plain = x.rms_norm_affine(gain_shifted, 1e-6).unwrap();
let a = via_offset.realize_f32();
let b = via_plain.realize_f32();
for (x, y) in a.iter().zip(b.iter()) {
assert!((x - y).abs() < 1e-6, "{x} vs {y}");
}
}
#[test]
fn global_avg_pool_2d_averages_spatial_dims() {
let x = cpu_f32(
vec![
1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 10.0, 20.0, 30.0, 40.0, 50.0, 60.0,
],
&[1, 2, 2, 3],
);
let out = x.global_avg_pool_2d().unwrap();
assert_eq!(out.shape().dims(), &[1, 2]);
let v = out.realize_f32();
assert!((v[0] - 3.5).abs() < 1e-5);
assert!((v[1] - 35.0).abs() < 1e-4);
}
#[test]
fn channel_affine_4d_applies_per_channel_gain_and_bias() {
let x = cpu_f32(
vec![
1.0_f32, 2.0, 3.0, 4.0, 10.0, 20.0, 30.0, 40.0, ],
&[1, 2, 2, 2],
);
let gain = std::sync::Arc::<[f32]>::from(vec![2.0_f32, 0.5]);
let bias = std::sync::Arc::<[f32]>::from(vec![1.0_f32, -10.0]);
let out = x.channel_affine_4d(gain, bias).unwrap();
let v = out.realize_f32();
assert_eq!(&v[0..4], &[3.0, 5.0, 7.0, 9.0]);
assert_eq!(&v[4..8], &[-5.0, 0.0, 5.0, 10.0]);
}
#[test]
fn additive_causal_mask_has_strict_lower_triangle() {
let anchor = cpu_f32(vec![0.0_f32], &[1]);
let mask = Tensor::additive_causal_mask_like(&anchor, 4);
assert_eq!(mask.shape().dims(), &[4, 4]);
let v = mask.realize_f32();
for i in 0..4 {
for j in 0..4 {
let got = v[i * 4 + j];
if j > i {
assert!(
got.is_infinite() && got.is_sign_negative(),
"above-diag (i={i}, j={j}) should be -inf, got {got}"
);
} else {
assert_eq!(
got, 0.0,
"on/below-diag (i={i}, j={j}) should be 0, got {got}"
);
}
}
}
}
#[test]
fn layer_norm_affine_unit_gain_zero_bias_matches_layer_norm_last_dim() {
let a = cpu_f32(vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]);
let gain = std::sync::Arc::<[f32]>::from(vec![1.0_f32; 3]);
let bias = std::sync::Arc::<[f32]>::from(vec![0.0_f32; 3]);
let out_affine = a.layer_norm_affine(gain, bias, 1e-5).unwrap();
let out_plain = a.layer_norm_last_dim(1e-5).unwrap();
let va = out_affine.realize_f32();
let vp = out_plain.realize_f32();
for (a, b) in va.iter().zip(vp.iter()) {
assert!((a - b).abs() < 1e-5, "{a} vs {b}");
}
}
#[test]
fn layer_norm_affine_applies_gain_and_bias() {
let a = cpu_f32(vec![1.0_f32, 2.0, 3.0], &[1, 3]);
let gain = std::sync::Arc::<[f32]>::from(vec![2.0_f32, 0.5, 1.5]);
let bias = std::sync::Arc::<[f32]>::from(vec![10.0_f32, -5.0, 0.0]);
let out = a.layer_norm_affine(gain, bias, 1e-5).unwrap();
let v = out.realize_f32();
let mean = 2.0_f32;
let var = ((1.0 - mean).powi(2) + (2.0 - mean).powi(2) + (3.0 - mean).powi(2)) / 3.0;
let den = (var + 1e-5_f32).sqrt();
let expected = [
((1.0 - mean) / den) * 2.0 + 10.0,
((2.0 - mean) / den) * 0.5 + (-5.0),
((3.0 - mean) / den) * 1.5 + 0.0,
];
for (got, want) in v.iter().zip(expected.iter()) {
assert!((got - want).abs() < 1e-4, "{got} vs {want}");
}
}
#[test]
fn l2_normalize_last_dim_unit_norm_per_row() {
let a = cpu_f32(vec![3.0, 4.0, 0.0, 1.0, 2.0, 2.0], &[2, 3]);
let out = a.l2_normalize(1_usize, 1e-12).unwrap();
assert_eq!(out.shape().dims(), &[2, 3]);
let v = out.realize_f32();
let row0_norm = (v[0] * v[0] + v[1] * v[1] + v[2] * v[2]).sqrt();
let row1_norm = (v[3] * v[3] + v[4] * v[4] + v[5] * v[5]).sqrt();
assert!((row0_norm - 1.0).abs() < 1e-5, "row 0 norm = {row0_norm}");
assert!((row1_norm - 1.0).abs() < 1e-5, "row 1 norm = {row1_norm}");
assert!((v[0] - 0.6).abs() < 1e-5);
assert!((v[1] - 0.8).abs() < 1e-5);
assert!(v[2].abs() < 1e-5);
}
#[test]
fn l2_normalize_eps_zero_works_when_nonzero() {
let a = cpu_f32(vec![1.0_f32, 0.0], &[2]);
let out = a.l2_normalize(0_usize, 0.0).unwrap();
let v = out.realize_f32();
assert!((v[0] - 1.0).abs() < 1e-6);
assert!(v[1].abs() < 1e-6);
}
#[test]
fn repeat_interleave_last_dim_matches_torch_semantics() {
let a = cpu_f32(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]);
let out = a.repeat_interleave(1_usize, 2).unwrap();
assert_eq!(out.shape().dims(), &[2, 6]);
assert_eq!(
out.realize_f32(),
vec![1.0, 1.0, 2.0, 2.0, 3.0, 3.0, 4.0, 4.0, 5.0, 5.0, 6.0, 6.0]
);
}
#[test]
fn repeat_interleave_middle_dim() {
let a = cpu_f32((0..8).map(|i| i as f32).collect(), &[2, 2, 2]);
let out = a.repeat_interleave(1_usize, 3).unwrap();
assert_eq!(out.shape().dims(), &[2, 6, 2]);
let v = out.realize_f32();
assert_eq!(&v[0..6], &[0.0, 1.0, 0.0, 1.0, 0.0, 1.0]);
assert_eq!(&v[6..12], &[2.0, 3.0, 2.0, 3.0, 2.0, 3.0]);
}
#[test]
fn repeat_interleave_repeats_one_is_noop() {
let a = cpu_f32(vec![1.0, 2.0, 3.0], &[3]);
let out = a.repeat_interleave(0_usize, 1).unwrap();
assert_eq!(out.shape().dims(), &[3]);
assert_eq!(out.realize_f32(), vec![1.0, 2.0, 3.0]);
}
#[test]
fn repeat_interleave_repeats_zero_errors() {
let a = cpu_f32(vec![1.0, 2.0, 3.0], &[3]);
assert!(a.repeat_interleave(0_usize, 0).is_err());
}
#[test]
fn stack_adds_leading_dim() {
let a = cpu_f32(vec![1.0, 2.0, 3.0], &[3]);
let b = a.const_f32_like(vec![4.0, 5.0, 6.0], vec![3]).unwrap();
let out = Tensor::stack(&[&a, &b], 0).unwrap();
assert_eq!(out.shape().dims(), &[2, 3]);
assert_eq!(out.realize_f32(), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
}
#[test]
fn stack_adds_trailing_dim() {
let a = cpu_f32(vec![1.0, 2.0, 3.0], &[3]);
let b = a.const_f32_like(vec![4.0, 5.0, 6.0], vec![3]).unwrap();
let out = Tensor::stack(&[&a, &b], 1).unwrap();
assert_eq!(out.shape().dims(), &[3, 2]);
assert_eq!(out.realize_f32(), vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0]);
}
#[test]
fn stack_rejects_mismatched_shapes() {
let a = cpu_f32(vec![1.0, 2.0], &[2]);
let b = a.const_f32_like(vec![3.0, 4.0, 5.0], vec![3]).unwrap();
assert!(Tensor::stack(&[&a, &b], 0).is_err());
}
#[test]
fn stack_rejects_empty_input() {
let result = Tensor::stack(&[], 0);
assert!(result.is_err());
}
#[test]
fn repeat_tiles_along_each_dim() {
let t = cpu_f32(vec![1.0, 2.0], &[2]);
let out = t.repeat(vec![3]).unwrap();
assert_eq!(out.shape().dims(), &[6]);
assert_eq!(out.realize_f32(), vec![1.0, 2.0, 1.0, 2.0, 1.0, 2.0]);
}
#[test]
fn repeat_extends_rank_when_needed() {
let t = cpu_f32(vec![1.0, 2.0], &[2]);
let out = t.repeat(vec![3, 2]).unwrap();
assert_eq!(out.shape().dims(), &[3, 4]);
}
#[test]
fn repeat_identity_with_all_ones() {
let t = cpu_f32(vec![1.0, 2.0, 3.0], &[3]);
let out = t.repeat(vec![1]).unwrap();
assert_eq!(out.realize_f32(), vec![1.0, 2.0, 3.0]);
}
}
#[cfg(test)]
mod phase_a3_keepdim_tests {
use super::*;
fn cpu_f32(data: Vec<f32>, shape: &[usize]) -> Tensor {
Tensor::from_f32(data, shape.to_vec(), &Device::cpu()).unwrap()
}
#[test]
fn sum_keepdim_preserves_dim_as_one() {
let t = cpu_f32(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]);
let out = t.sum_keepdim(1_usize).unwrap();
assert_eq!(out.shape().dims(), &[2, 1]);
assert_eq!(out.realize_f32(), vec![3.0, 7.0]);
}
#[test]
fn mean_keepdim_preserves_dim_as_one() {
let t = cpu_f32(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]);
let out = t.mean_keepdim(0_usize).unwrap();
assert_eq!(out.shape().dims(), &[1, 2]);
assert_eq!(out.realize_f32(), vec![2.0, 3.0]);
}
#[test]
fn max_keepdim_preserves_dim_as_one() {
let t = cpu_f32(vec![1.0, 3.0, 2.0, 4.0], &[2, 2]);
let out = t.max_keepdim(1_usize).unwrap();
assert_eq!(out.shape().dims(), &[2, 1]);
assert_eq!(out.realize_f32(), vec![3.0, 4.0]);
}
#[test]
fn min_keepdim_preserves_dim_as_one() {
let t = cpu_f32(vec![1.0, 3.0, 2.0, 4.0], &[2, 2]);
let out = t.min_keepdim(1_usize).unwrap();
assert_eq!(out.shape().dims(), &[2, 1]);
assert_eq!(out.realize_f32(), vec![1.0, 2.0]);
}
#[test]
fn var_matches_unbiased_formula_with_correction_1() {
let t = cpu_f32(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]);
let out = t.var(1, 1.0).unwrap();
assert_eq!(out.shape().dims(), &[2]);
let v = out.realize_f32();
assert!((v[0] - 1.0).abs() < 1e-5, "var row 0 = {} != 1.0", v[0]);
assert!((v[1] - 1.0).abs() < 1e-5, "var row 1 = {} != 1.0", v[1]);
}
#[test]
fn var_population_with_correction_0() {
let t = cpu_f32(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]);
let out = t.var(1, 0.0).unwrap();
let v = out.realize_f32();
let want = 2.0 / 3.0;
assert!(
(v[0] - want).abs() < 1e-5,
"population var row 0 = {} != {want}",
v[0]
);
assert!(
(v[1] - want).abs() < 1e-5,
"population var row 1 = {} != {want}",
v[1]
);
}
#[test]
fn var_keepdim_preserves_dim() {
let t = cpu_f32(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]);
let out = t.var_keepdim(1, 1.0).unwrap();
assert_eq!(out.shape().dims(), &[2, 1]);
}
#[test]
fn var_errors_out_of_bounds() {
let t = cpu_f32(vec![1.0, 2.0], &[2]);
assert!(t.var(3, 1.0).is_err());
assert!(t.var_keepdim(3, 1.0).is_err());
}
}
#[cfg(test)]
mod phase_a4_composite_tests {
use super::*;
fn cpu_f32(data: Vec<f32>, shape: &[usize]) -> Tensor {
Tensor::from_f32(data, shape.to_vec(), &Device::cpu()).unwrap()
}
#[test]
fn affine_applies_mul_then_add() {
let t = cpu_f32(vec![1.0, 2.0, 3.0], &[3]);
let out = t.affine(2.0, 10.0);
assert_eq!(out.realize_f32(), vec![12.0, 14.0, 16.0]);
}
#[test]
fn scale_and_shift_alias_of_affine() {
let t = cpu_f32(vec![1.0, 2.0, 3.0], &[3]);
assert_eq!(
t.scale_and_shift(2.0, 10.0).realize_f32(),
t.affine(2.0, 10.0).realize_f32(),
);
}
#[test]
fn elu_matches_reference_values() {
let t = cpu_f32(vec![1.0, 0.0, -1.0, -2.0], &[4]);
let out = t.elu(1.0);
let v = out.realize_f32();
assert!((v[0] - 1.0).abs() < 1e-5);
assert!(v[1].abs() < 1e-5);
assert!((v[2] - ((-1.0_f32).exp() - 1.0)).abs() < 1e-5);
assert!((v[3] - ((-2.0_f32).exp() - 1.0)).abs() < 1e-5);
}
#[test]
fn dot_of_rank_one_vectors() {
let a = cpu_f32(vec![1.0, 2.0, 3.0], &[3]);
let b = a.const_f32_like(vec![4.0, 5.0, 6.0], vec![3]).unwrap();
let out = a.dot(&b).unwrap();
assert_eq!(out.shape().elem_count(), 1);
let v = out.realize_f32();
assert_eq!(v[0], 32.0); }
#[test]
fn dot_rejects_non_rank_one() {
let a = cpu_f32(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]);
let b = a
.const_f32_like(vec![1.0, 2.0, 3.0, 4.0], vec![2, 2])
.unwrap();
assert!(a.dot(&b).is_err());
}
#[test]
fn dot_rejects_length_mismatch() {
let a = cpu_f32(vec![1.0, 2.0], &[2]);
let b = a.const_f32_like(vec![1.0, 2.0, 3.0], vec![3]).unwrap();
assert!(a.dot(&b).is_err());
}
#[test]
fn mv_matrix_times_vector() {
let m = cpu_f32(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]);
let v = m.const_f32_like(vec![1.0, 1.0, 1.0], vec![3]).unwrap();
let out = m.mv(&v).unwrap();
assert_eq!(out.shape().dims(), &[2]);
assert_eq!(out.realize_f32(), vec![6.0, 15.0]);
}
#[test]
fn matvec_is_mv_alias() {
let m = cpu_f32(vec![1.0, 0.0, 0.0, 1.0], &[2, 2]);
let v = m.const_f32_like(vec![3.0, 4.0], vec![2]).unwrap();
let a = m.mv(&v).unwrap().realize_f32();
let b = m.matvec(&v).unwrap().realize_f32();
assert_eq!(a, b);
}
#[test]
fn mv_rejects_shape_mismatch() {
let m = cpu_f32(vec![1.0; 6], &[2, 3]);
let v = m.const_f32_like(vec![1.0, 1.0], vec![2]).unwrap();
assert!(m.mv(&v).is_err());
}
#[test]
fn broadcast_matmul_passes_through_to_matmul() {
let a = cpu_f32(vec![1.0, 0.0, 0.0, 1.0], &[2, 2]);
let b = a
.const_f32_like(vec![5.0, 6.0, 7.0, 8.0], vec![2, 2])
.unwrap();
let out = a.broadcast_matmul(&b).unwrap();
assert_eq!(out.realize_f32(), vec![5.0, 6.0, 7.0, 8.0]);
}
}
#[cfg(test)]
mod phase_a5_factory_tests {
use super::*;
fn cpu_f32(data: Vec<f32>, shape: &[usize]) -> Tensor {
Tensor::from_f32(data, shape.to_vec(), &Device::cpu()).unwrap()
}
#[test]
fn ones_like_matches_shape_dtype_graph() {
let t = cpu_f32(vec![5.0, 6.0, 7.0, 8.0], &[2, 2]);
let ones = t.ones_like().unwrap();
assert_eq!(ones.shape().dims(), t.shape().dims());
assert_eq!(ones.dtype(), t.dtype());
assert_eq!(ones.realize_f32(), vec![1.0; 4]);
}
#[test]
fn zeros_like_matches_shape_dtype_graph() {
let t = cpu_f32(vec![5.0, 6.0, 7.0], &[3]);
let zeros = t.zeros_like().unwrap();
assert_eq!(zeros.realize_f32(), vec![0.0; 3]);
}
#[test]
fn static_ones_f32() {
let t = Tensor::ones(vec![2, 3], DType::F32, &Device::cpu()).unwrap();
assert_eq!(t.shape().dims(), &[2, 3]);
assert_eq!(t.realize_f32(), vec![1.0; 6]);
}
#[test]
fn static_zeros_f64() {
let t = Tensor::zeros(vec![4], DType::F64, &Device::cpu()).unwrap();
assert_eq!(t.dtype(), DType::F64);
assert_eq!(t.realize_f64(), vec![0.0; 4]);
}
#[test]
fn full_with_f32_scalar() {
let t = Tensor::full(vec![5], fuel_ir::Scalar::F32(2.5), &Device::cpu()).unwrap();
assert_eq!(t.realize_f32(), vec![2.5; 5]);
}
#[test]
fn eye_identity_matrix() {
let t = Tensor::eye(3, DType::F32, &Device::cpu());
assert_eq!(t.shape().dims(), &[3, 3]);
assert_eq!(
t.realize_f32(),
vec![1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0],
);
}
#[test]
fn tril2_lower_triangular_ones() {
let t = Tensor::tril2(3, DType::F32, &Device::cpu());
assert_eq!(
t.realize_f32(),
vec![1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0],
);
}
#[test]
fn triu2_upper_triangular_ones() {
let t = Tensor::triu2(3, DType::F32, &Device::cpu());
assert_eq!(
t.realize_f32(),
vec![1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0],
);
}
#[test]
fn meshgrid_ij_indexing_two_inputs() {
let x = Tensor::from_f32(vec![1.0_f32, 2.0, 3.0], vec![3], &Device::cpu()).unwrap();
let y = x.const_f32_like(vec![4.0_f32, 5.0], vec![2]).unwrap();
let grids = Tensor::meshgrid(&[&x, &y], false).unwrap();
assert_eq!(grids.len(), 2);
assert_eq!(grids[0].shape().dims(), &[3, 2]);
assert_eq!(grids[1].shape().dims(), &[3, 2]);
assert_eq!(grids[0].realize_f32(), vec![1.0, 1.0, 2.0, 2.0, 3.0, 3.0]);
assert_eq!(grids[1].realize_f32(), vec![4.0, 5.0, 4.0, 5.0, 4.0, 5.0]);
}
#[test]
fn meshgrid_xy_indexing_swaps_first_two() {
let x = Tensor::from_f32(vec![1.0_f32, 2.0, 3.0], vec![3], &Device::cpu()).unwrap();
let y = x.const_f32_like(vec![4.0_f32, 5.0], vec![2]).unwrap();
let grids = Tensor::meshgrid(&[&x, &y], true).unwrap();
assert_eq!(grids[0].shape().dims(), &[2, 3]);
assert_eq!(grids[1].shape().dims(), &[2, 3]);
assert_eq!(grids[0].realize_f32(), vec![1.0, 2.0, 3.0, 1.0, 2.0, 3.0]);
assert_eq!(grids[1].realize_f32(), vec![4.0, 4.0, 4.0, 5.0, 5.0, 5.0]);
}
#[test]
fn meshgrid_rejects_single_input() {
let x = Tensor::from_f32(vec![1.0_f32, 2.0], vec![2], &Device::cpu()).unwrap();
assert!(Tensor::meshgrid(&[&x], false).is_err());
}
#[test]
fn meshgrid_rejects_non_rank_one() {
let x = Tensor::from_f32(vec![1.0; 4], vec![2, 2], &Device::cpu()).unwrap();
let y = x.const_f32_like(vec![1.0, 2.0], vec![2]).unwrap();
assert!(Tensor::meshgrid(&[&x, &y], false).is_err());
}
#[test]
fn narrow_is_slice_alias() {
let t = cpu_f32(vec![1.0, 2.0, 3.0, 4.0, 5.0], &[5]);
let a = t.narrow(0, 1, 3).unwrap().realize_f32();
let b = t.slice(0, 1, 3).unwrap().realize_f32();
assert_eq!(a, b);
assert_eq!(a, vec![2.0, 3.0, 4.0]);
}
#[test]
fn chunk_splits_evenly() {
let t = cpu_f32((1..=6).map(|i| i as f32).collect(), &[6]);
let parts = t.chunk(3, 0).unwrap();
assert_eq!(parts.len(), 3);
assert_eq!(parts[0].realize_f32(), vec![1.0, 2.0]);
assert_eq!(parts[1].realize_f32(), vec![3.0, 4.0]);
assert_eq!(parts[2].realize_f32(), vec![5.0, 6.0]);
}
#[test]
fn chunk_distributes_remainder_to_leading() {
let t = cpu_f32((1..=7).map(|i| i as f32).collect(), &[7]);
let parts = t.chunk(3, 0).unwrap();
assert_eq!(parts.len(), 3);
assert_eq!(parts[0].realize_f32(), vec![1.0, 2.0, 3.0]);
assert_eq!(parts[1].realize_f32(), vec![4.0, 5.0]);
assert_eq!(parts[2].realize_f32(), vec![6.0, 7.0]);
}
#[test]
fn chunk_returns_singletons_when_size_less_than_chunks() {
let t = cpu_f32(vec![1.0, 2.0], &[2]);
let parts = t.chunk(5, 0).unwrap();
assert_eq!(parts.len(), 2);
assert_eq!(parts[0].realize_f32(), vec![1.0]);
assert_eq!(parts[1].realize_f32(), vec![2.0]);
}
#[test]
fn get_at_first_dim() {
let t = cpu_f32(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let row1 = t.get(1).unwrap();
assert_eq!(row1.shape().dims(), &[2]);
assert_eq!(row1.realize_f32(), vec![3.0, 4.0]);
}
#[test]
fn get_on_dim_arbitrary_axis() {
let t = cpu_f32(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let col0 = t.get_on_dim(1, 0).unwrap();
assert_eq!(col0.shape().dims(), &[3]);
assert_eq!(col0.realize_f32(), vec![1.0, 3.0, 5.0]);
}
#[test]
fn sum_dims_multi_dim_reduces_to_smaller() {
let t = cpu_f32(vec![1.0; 24], &[2, 3, 4]);
let s = t.sum_dims([0, 2_usize]).unwrap();
assert_eq!(s.shape().dims(), &[3]);
assert_eq!(s.realize_f32(), vec![8.0; 3]);
}
#[test]
fn mean_dims_multi_dim() {
let t = cpu_f32(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]);
let m = t.mean_dims([0, 1_usize]).unwrap();
assert_eq!(m.shape().dims(), &[] as &[usize]);
assert_eq!(m.realize_f32(), vec![2.5]);
}
#[test]
fn sum_dims_keepdim_preserves_rank() {
let t = cpu_f32(vec![1.0; 24], &[2, 3, 4]);
let s = t.sum_dims_keepdim(&[0, 2]).unwrap();
assert_eq!(s.shape().dims(), &[1, 3, 1]);
}
#[test]
fn mean_dims_keepdim_preserves_rank() {
let t = cpu_f32(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]);
let m = t.mean_dims_keepdim(&[0, 1]).unwrap();
assert_eq!(m.shape().dims(), &[1, 1]);
assert_eq!(m.realize_f32(), vec![2.5]);
}
#[test]
fn rand_like_matches_shape_dtype() {
let t = cpu_f32(vec![0.0; 6], &[2, 3]);
let r = t.rand_like(-1.0, 1.0).unwrap();
assert_eq!(r.shape().dims(), t.shape().dims());
assert_eq!(r.dtype(), t.dtype());
for v in r.realize_f32() {
assert!((-1.0..1.0).contains(&v), "sample {v} out of [-1, 1)");
}
}
#[test]
fn randn_like_matches_shape_dtype() {
let t = cpu_f32(vec![0.0; 4], &[4]);
let r = t.randn_like(0.0, 1.0).unwrap();
assert_eq!(r.shape().dims(), &[4]);
assert_eq!(r.dtype(), DType::F32);
for v in r.realize_f32() {
assert!(v.is_finite());
}
}
#[test]
fn static_rand_f32() {
let t = Tensor::rand(vec![100], 0.0, 1.0, DType::F32, &Device::cpu()).unwrap();
let v = t.realize_f32();
let mean: f32 = v.iter().sum::<f32>() / v.len() as f32;
assert!((0.3..0.7).contains(&mean), "mean {mean} too far from 0.5");
}
#[test]
fn static_randn_f64() {
let t = Tensor::randn(vec![1000], 0.0, 1.0, DType::F64, &Device::cpu()).unwrap();
let v = t.realize_f64();
let mean: f64 = v.iter().sum::<f64>() / v.len() as f64;
assert!(mean.abs() < 0.2, "mean {mean} too far from 0");
}
#[test]
fn arange_int_step() {
let t = Tensor::arange(0.0, 5.0, &Device::cpu());
assert_eq!(t.shape().dims(), &[5]);
assert_eq!(t.realize_f32(), vec![0.0, 1.0, 2.0, 3.0, 4.0]);
}
#[test]
fn arange_step_fractional() {
let t = Tensor::arange_step(2.0, 4.0, 0.5, &Device::cpu());
assert_eq!(t.realize_f32(), vec![2.0, 2.5, 3.0, 3.5]);
}
#[test]
fn arange_step_negative_descends() {
let t = Tensor::arange_step(5.0, 0.0, -1.0, &Device::cpu());
assert_eq!(t.realize_f32(), vec![5.0, 4.0, 3.0, 2.0, 1.0]);
}
#[test]
fn linspace_includes_endpoints() {
let t = Tensor::linspace(0.0, 1.0, 5, &Device::cpu());
assert_eq!(t.shape().dims(), &[5]);
let v = t.realize_f32();
assert!((v[0] - 0.0).abs() < 1e-6);
assert!((v[4] - 1.0).abs() < 1e-6);
assert!((v[2] - 0.5).abs() < 1e-6);
}
#[test]
fn linspace_n_one_returns_start() {
let t = Tensor::linspace(7.0, 99.0, 1, &Device::cpu());
assert_eq!(t.realize_f32(), vec![7.0]);
}
#[test]
fn norm_is_sqrt_sum_sq() {
let t = Tensor::from_f32(vec![3.0_f32, 4.0], vec![2], &Device::cpu()).unwrap();
let n = t.norm();
assert!((n.realize_f32()[0] - 5.0).abs() < 1e-6);
}
#[test]
fn pad_with_zeros_left_and_right() {
let t = cpu_f32(vec![1.0, 2.0, 3.0], &[3]);
let p = t.pad_with_zeros(0, 2, 1).unwrap();
assert_eq!(p.shape().dims(), &[6]);
assert_eq!(p.realize_f32(), vec![0.0, 0.0, 1.0, 2.0, 3.0, 0.0]);
}
#[test]
fn pad_with_zeros_identity_when_both_zero() {
let t = cpu_f32(vec![1.0, 2.0, 3.0], &[3]);
let p = t.pad_with_zeros(0, 0, 0).unwrap();
assert_eq!(p.realize_f32(), vec![1.0, 2.0, 3.0]);
}
#[test]
fn pad_with_zeros_rejects_bad_dim() {
let t = cpu_f32(vec![1.0, 2.0], &[2]);
assert!(t.pad_with_zeros(3, 1, 1).is_err());
}
#[test]
fn pad_with_value_zero_matches_pad_with_zeros() {
let t = cpu_f32(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]);
let via_zeros = t.pad_with_zeros(0, 1, 2).unwrap();
let via_value = t.pad_with_value(0, 1, 2, 0.0).unwrap();
assert_eq!(via_zeros.shape().dims(), via_value.shape().dims());
assert_eq!(via_zeros.shape().dims(), &[5, 2]);
assert_eq!(via_zeros.realize_f32(), via_value.realize_f32());
assert_eq!(
via_value.realize_f32(),
vec![0.0, 0.0, 1.0, 2.0, 3.0, 4.0, 0.0, 0.0, 0.0, 0.0],
);
}
#[test]
fn pad_with_value_neg_inf_then_max_reduce_ignores_pad() {
let t = cpu_f32(vec![-1.0, -2.0, -3.0], &[3]);
let padded = t.pad_with_value(0, 2, 2, f32::NEG_INFINITY).unwrap();
assert_eq!(padded.shape().dims(), &[7]);
let v = padded.realize_f32();
assert!(v[0].is_infinite() && v[0].is_sign_negative());
assert!(v[1].is_infinite() && v[1].is_sign_negative());
assert_eq!(v[2], -1.0);
assert_eq!(v[3], -2.0);
assert_eq!(v[4], -3.0);
assert!(v[5].is_infinite() && v[5].is_sign_negative());
assert!(v[6].is_infinite() && v[6].is_sign_negative());
let m = padded.max_all().realize_f32();
assert_eq!(m, vec![-1.0]);
}
#[test]
fn pad_with_value_identity_when_both_zero() {
let t = cpu_f32(vec![1.0, 2.0, 3.0], &[3]);
let p = t.pad_with_value(0, 0, 0, f32::NEG_INFINITY).unwrap();
assert_eq!(p.realize_f32(), vec![1.0, 2.0, 3.0]);
}
#[test]
fn pad_with_value_rejects_bad_dim() {
let t = cpu_f32(vec![1.0, 2.0], &[2]);
assert!(t.pad_with_value(3, 1, 1, 0.0).is_err());
}
#[test]
fn max_pool2d_with_pad_value_neg_inf_on_negative_interior() {
let x = cpu_f32(
vec![-1.0, -2.0, -3.0, -4.0, -5.0, -6.0, -7.0, -8.0, -9.0],
&[1, 1, 3, 3],
);
let out = x
.max_pool2d_with_pad_value((3, 3), (1, 1), (1, 1), f32::NEG_INFINITY)
.unwrap();
assert_eq!(out.shape().dims(), &[1, 1, 3, 3]);
let v = out.realize_f32();
assert_eq!(v[0], -1.0);
assert_eq!(v[4], -1.0);
assert_eq!(v[8], -5.0);
let default_pad = x.max_pool2d((3, 3), (1, 1), (1, 1)).unwrap();
let vd = default_pad.realize_f32();
assert_eq!(
vd[0], -1.0,
"max_pool2d's default must now match NEG_INFINITY padding, not 0.0"
);
assert_eq!(vd[8], -5.0);
let explicit_zero_pad = x
.max_pool2d_with_pad_value((3, 3), (1, 1), (1, 1), 0.0)
.unwrap();
let vz = explicit_zero_pad.realize_f32();
assert_eq!(
vz[0], 0.0,
"explicit 0.0 padding must still be reachable via max_pool2d_with_pad_value"
);
assert_eq!(vz[8], 0.0);
}
#[test]
fn max_pool2d_with_pad_value_zero_matches_max_pool2d() {
let data: Vec<f32> = (0..16).map(|i| i as f32).collect();
let x = cpu_f32(data, &[1, 1, 4, 4]);
let a = x.max_pool2d((2, 2), (2, 2), (0, 0)).unwrap();
let b = x
.max_pool2d_with_pad_value((2, 2), (2, 2), (0, 0), 0.0)
.unwrap();
assert_eq!(a.shape().dims(), b.shape().dims());
assert_eq!(a.realize_f32(), b.realize_f32());
}
#[test]
fn conv1d_identity_kernel_passes_input_through() {
let x = cpu_f32(vec![1.0, 2.0, 3.0, 4.0, 5.0], &[1, 1, 5]);
let w = x.const_f32_like(vec![1.0], vec![1, 1, 1]).unwrap();
let out = x.conv1d(&w, None, 1, 0, 1).unwrap();
assert_eq!(out.shape().dims(), &[1, 1, 5]);
assert_eq!(out.realize_f32(), vec![1.0, 2.0, 3.0, 4.0, 5.0]);
}
#[test]
fn conv1d_sum_kernel_two_wide() {
let x = cpu_f32(vec![1.0, 2.0, 3.0, 4.0], &[1, 1, 4]);
let w = x.const_f32_like(vec![1.0, 1.0], vec![1, 1, 2]).unwrap();
let out = x.conv1d(&w, None, 1, 0, 1).unwrap();
assert_eq!(out.shape().dims(), &[1, 1, 3]);
assert_eq!(out.realize_f32(), vec![3.0, 5.0, 7.0]);
}
#[test]
fn conv1d_with_bias_applies_correctly() {
let x = cpu_f32(vec![1.0, 1.0, 1.0], &[1, 1, 3]);
let w = x.const_f32_like(vec![1.0], vec![1, 1, 1]).unwrap();
let bias = x.const_f32_like(vec![10.0], vec![1]).unwrap();
let out = x.conv1d(&w, Some(&bias), 1, 0, 1).unwrap();
assert_eq!(out.realize_f32(), vec![11.0, 11.0, 11.0]);
}
#[test]
fn conv1d_stride_two_halves_output() {
let x = cpu_f32(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[1, 1, 6]);
let w = x.const_f32_like(vec![1.0, 1.0], vec![1, 1, 2]).unwrap();
let out = x.conv1d(&w, None, 2, 0, 1).unwrap();
assert_eq!(out.shape().dims(), &[1, 1, 3]);
assert_eq!(out.realize_f32(), vec![3.0, 7.0, 11.0]);
}
#[test]
fn conv1d_padding_one_preserves_length() {
let x = cpu_f32(vec![1.0, 2.0, 3.0, 4.0], &[1, 1, 4]);
let w = x
.const_f32_like(vec![1.0, 1.0, 1.0], vec![1, 1, 3])
.unwrap();
let out = x.conv1d(&w, None, 1, 1, 1).unwrap();
assert_eq!(out.shape().dims(), &[1, 1, 4]);
assert_eq!(out.realize_f32(), vec![3.0, 6.0, 9.0, 7.0]);
}
#[test]
fn conv1d_multi_channel_output() {
let x = cpu_f32(vec![1.0, 2.0, 3.0], &[1, 1, 3]);
let w = x.const_f32_like(vec![2.0, -1.0], vec![2, 1, 1]).unwrap();
let out = x.conv1d(&w, None, 1, 0, 1).unwrap();
assert_eq!(out.shape().dims(), &[1, 2, 3]);
assert_eq!(out.realize_f32(), vec![2.0, 4.0, 6.0, -1.0, -2.0, -3.0]);
}
#[test]
fn conv1d_rejects_rank_two_input() {
let x = cpu_f32(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]);
let w = x.const_f32_like(vec![1.0], vec![1, 1, 1]).unwrap();
assert!(x.conv1d(&w, None, 1, 0, 1).is_err());
}
#[test]
fn conv1d_rejects_rank_two_weight() {
let x = cpu_f32(vec![1.0; 4], &[1, 1, 4]);
let w = x.const_f32_like(vec![1.0], vec![1, 1]).unwrap();
assert!(x.conv1d(&w, None, 1, 0, 1).is_err());
}
#[test]
fn conv1d_rejects_zero_groups_or_stride() {
let x = cpu_f32(vec![1.0; 4], &[1, 1, 4]);
let w = x.const_f32_like(vec![1.0], vec![1, 1, 1]).unwrap();
assert!(x.conv1d(&w, None, 0, 0, 1).is_err());
assert!(x.conv1d(&w, None, 1, 0, 0).is_err());
}
#[test]
fn conv1d_with_algo_ignores_algo_param() {
let x = cpu_f32(vec![1.0, 2.0, 3.0, 4.0], &[1, 1, 4]);
let w = x.const_f32_like(vec![1.0, 1.0], vec![1, 1, 2]).unwrap();
let a = x.conv1d_with_algo(&w, None, 1, 0, 1, "unused").unwrap();
let b = x.conv1d(&w, None, 1, 0, 1).unwrap();
assert_eq!(a.realize_f32(), b.realize_f32());
}
#[test]
fn avg_pool2d_2x2_stride2() {
let data: Vec<f32> = (0..16).map(|i| i as f32).collect();
let x = cpu_f32(data, &[1, 1, 4, 4]);
let out = x.avg_pool2d((2, 2), (2, 2), (0, 0), true).unwrap();
assert_eq!(out.shape().dims(), &[1, 1, 2, 2]);
let v = out.realize_f32();
assert!((v[0] - 2.5).abs() < 1e-5);
assert!((v[1] - 4.5).abs() < 1e-5);
assert!((v[2] - 10.5).abs() < 1e-5);
assert!((v[3] - 12.5).abs() < 1e-5);
}
#[test]
fn avg_pool2d_3x3_stride1_padding1_preserves_size() {
let x = cpu_f32(vec![1.0; 16], &[1, 1, 4, 4]);
let out = x.avg_pool2d((3, 3), (1, 1), (1, 1), true).unwrap();
assert_eq!(out.shape().dims(), &[1, 1, 4, 4]);
}
#[test]
fn avg_pool2d_excludes_padding_keeps_uniform_data_exactly_at_its_true_value() {
let x = cpu_f32(vec![1.0; 16], &[1, 1, 4, 4]);
let included = x.avg_pool2d((3, 3), (1, 1), (1, 1), true).unwrap();
let excluded = x.avg_pool2d((3, 3), (1, 1), (1, 1), false).unwrap();
assert_eq!(included.shape().dims(), excluded.shape().dims());
let ve = excluded.realize_f32();
for (i, &v) in ve.iter().enumerate() {
assert!(
(v - 1.0).abs() < 1e-5,
"count_include_pad=false must be exactly 1.0 at every position \
(all real taps are 1.0) — position {i} got {v}",
);
}
let vi = included.realize_f32();
assert!(
(vi[0] - 1.0).abs() > 0.1,
"count_include_pad=true must NOT match at a corner (4 of 9 taps \
real → 4/9, far from 1.0) — got {}, modes should disagree here",
vi[0]
);
}
#[test]
fn avg_pool2d_excludes_padding_from_divisor_discriminating() {
let x = cpu_f32(vec![10.0, 10.0, 10.0, 10.0], &[1, 1, 2, 2]);
let included = x.avg_pool2d((2, 2), (1, 1), (1, 1), true).unwrap();
let excluded = x.avg_pool2d((2, 2), (1, 1), (1, 1), false).unwrap();
let vi = included.realize_f32();
let ve = excluded.realize_f32();
assert!(
(vi[0] - 2.5).abs() < 1e-5,
"count_include_pad=true: top-left = 10/4 = 2.5, got {}",
vi[0]
);
assert!(
(ve[0] - 10.0).abs() < 1e-5,
"count_include_pad=false: top-left = 10/1 (ONE real tap) = 10.0, got {}",
ve[0]
);
}
#[test]
fn avg_pool2d_rejects_zero_valid_tap_window() {
let x = cpu_f32(vec![5.0], &[1, 1, 1, 1]);
let err = x
.avg_pool2d((1, 1), (1, 1), (2, 2), false)
.expect_err("a window with zero real taps must be rejected, not silently NaN");
let msg = err.to_string();
assert!(
msg.contains("divide by zero") || msg.contains("no valid"),
"error should explain the zero-valid-tap cause, got: {msg}"
);
let ok = x.avg_pool2d((1, 1), (1, 1), (2, 2), true);
assert!(
ok.is_ok(),
"count_include_pad=true must still accept this shape (uniform divisor, never zero)"
);
}
#[test]
fn avg_pool2d_multi_channel() {
let x = cpu_f32(vec![0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0], &[1, 2, 2, 2]);
let out = x.avg_pool2d((2, 2), (2, 2), (0, 0), true).unwrap();
assert_eq!(out.shape().dims(), &[1, 2, 1, 1]);
assert_eq!(out.realize_f32(), vec![0.0, 1.0]);
}
#[test]
fn max_pool2d_2x2_stride2() {
let data: Vec<f32> = (0..16).map(|i| i as f32).collect();
let x = cpu_f32(data, &[1, 1, 4, 4]);
let out = x.max_pool2d((2, 2), (2, 2), (0, 0)).unwrap();
assert_eq!(out.shape().dims(), &[1, 1, 2, 2]);
assert_eq!(out.realize_f32(), vec![5.0, 7.0, 13.0, 15.0]);
}
#[test]
fn max_pool2d_3x3_stride1_padding1() {
let x = cpu_f32(
vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0],
&[1, 1, 3, 3],
);
let out = x.max_pool2d((3, 3), (1, 1), (1, 1)).unwrap();
assert_eq!(out.shape().dims(), &[1, 1, 3, 3]);
let v = out.realize_f32();
assert!((v[4] - 9.0).abs() < 1e-5);
}
#[test]
fn upsample_nearest2d_2x() {
let x = cpu_f32(vec![1.0, 2.0, 3.0, 4.0], &[1, 1, 2, 2]);
let out = x.upsample_nearest2d(2).unwrap();
assert_eq!(out.shape().dims(), &[1, 1, 4, 4]);
assert_eq!(
out.realize_f32(),
vec![
1.0, 1.0, 2.0, 2.0, 1.0, 1.0, 2.0, 2.0, 3.0, 3.0, 4.0, 4.0, 3.0, 3.0, 4.0, 4.0,
],
);
}
#[test]
fn upsample_nearest2d_3x() {
let x = cpu_f32(vec![5.0], &[1, 1, 1, 1]);
let out = x.upsample_nearest2d(3).unwrap();
assert_eq!(out.shape().dims(), &[1, 1, 3, 3]);
assert_eq!(out.realize_f32(), vec![5.0; 9]);
}
#[test]
fn upsample_nearest2d_identity_scale_one() {
let x = cpu_f32(vec![1.0, 2.0, 3.0, 4.0], &[1, 1, 2, 2]);
let out = x.upsample_nearest2d(1).unwrap();
assert_eq!(out.realize_f32(), vec![1.0, 2.0, 3.0, 4.0]);
}
#[test]
fn upsample_nearest1d_2x() {
let x = cpu_f32(vec![1.0, 2.0, 3.0], &[1, 1, 3]);
let out = x.upsample_nearest1d(2).unwrap();
assert_eq!(out.shape().dims(), &[1, 1, 6]);
assert_eq!(out.realize_f32(), vec![1.0, 1.0, 2.0, 2.0, 3.0, 3.0]);
}
#[test]
fn interpolate2d_integer_multiple() {
let x = cpu_f32(vec![1.0, 2.0, 3.0, 4.0], &[1, 1, 2, 2]);
let out = x.interpolate2d(4, 4).unwrap();
assert_eq!(out.shape().dims(), &[1, 1, 4, 4]);
}
#[test]
fn interpolate2d_accepts_non_integer_ratio() {
let x = cpu_f32(vec![1.0; 4], &[1, 1, 2, 2]);
let out = x.interpolate2d(3, 4).unwrap();
assert_eq!(out.shape().dims(), &[1, 1, 3, 4]);
}
#[test]
fn interpolate1d_integer_multiple() {
let x = cpu_f32(vec![1.0, 2.0], &[1, 1, 2]);
let out = x.interpolate1d(6).unwrap();
assert_eq!(out.shape().dims(), &[1, 1, 6]);
assert_eq!(out.realize_f32(), vec![1.0, 1.0, 1.0, 2.0, 2.0, 2.0]);
}
#[test]
fn pool_rejects_bad_rank() {
let x = cpu_f32(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]);
assert!(x.avg_pool2d((2, 2), (2, 2), (0, 0), true).is_err());
assert!(x.max_pool2d((2, 2), (2, 2), (0, 0)).is_err());
}
#[test]
fn pool_rejects_zero_kernel() {
let x = cpu_f32(vec![1.0; 16], &[1, 1, 4, 4]);
assert!(x.avg_pool2d((0, 2), (1, 1), (0, 0), true).is_err());
assert!(x.max_pool2d((2, 0), (1, 1), (0, 0)).is_err());
}
#[test]
fn numel_matches_elem_count() {
let t = cpu_f32(vec![1.0; 12], &[3, 4]);
assert_eq!(t.numel(), t.elem_count());
assert_eq!(t.numel(), 12);
}
#[test]
fn dim_returns_specific_axis_size() {
let t = cpu_f32(vec![0.0; 24], &[2, 3, 4]);
assert_eq!(t.dim(0).unwrap(), 2);
assert_eq!(t.dim(1).unwrap(), 3);
assert_eq!(t.dim(2).unwrap(), 4);
assert!(t.dim(3).is_err());
}
#[test]
fn to_dtype_switches_dtype() {
let t = cpu_f32(vec![1.0, 2.0], &[2]);
let b = t.to_dtype(DType::F64).unwrap();
assert_eq!(b.dtype(), DType::F64);
assert_eq!(b.realize_f64(), vec![1.0, 2.0]);
}
#[test]
fn to_dtype_same_dtype_is_noop() {
let t = cpu_f32(vec![1.0_f32], &[1]);
let b = t.to_dtype(DType::F32).unwrap();
assert_eq!(b.dtype(), DType::F32);
}
#[test]
fn detach_is_identity_on_lazy() {
let t = cpu_f32(vec![1.0, 2.0, 3.0], &[3]);
let d = t.detach();
assert_eq!(d.realize_f32(), t.realize_f32());
}
#[test]
fn track_op_is_true_on_lazy() {
let t = cpu_f32(vec![0.0], &[1]);
assert!(t.track_op());
}
#[test]
fn try_permute_accepts_tuple_syntax() {
let t = cpu_f32(vec![0.0; 24], &[2, 3, 4]);
let out = t.permute((2_usize, 0_usize, 1_usize)).unwrap();
assert_eq!(out.shape().dims(), &[4, 2, 3]);
}
#[test]
fn dim_arg_methods_accept_negative_indexing() {
use fuel_ir::D;
let t = cpu_f32(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]);
let sum_last = t.shape().dims().to_vec(); assert_eq!(sum_last, vec![2, 3]);
let cumsum_last = t.cumsum(D::Minus1).unwrap();
assert_eq!(cumsum_last.shape().dims(), &[2, 3]);
}
#[test]
fn unsqueeze_accepts_dim_trait() {
use fuel_ir::D;
let t = cpu_f32(vec![1.0, 2.0, 3.0], &[3]);
let out = t.unsqueeze(D::Minus1).unwrap();
assert_eq!(out.shape().dims().len(), 2);
}
}
pub struct TokenDataHost {
pub token_ids: fuel_ir::HostBuffer,
pub rope_cos: fuel_ir::HostBuffer,
pub rope_sin: fuel_ir::HostBuffer,
pub mask: fuel_ir::HostBuffer,
pub offset: Option<fuel_ir::HostBuffer>,
}
#[cfg(feature = "cuda")]
pub struct TokenDataBytes {
pub token_ids: Vec<u8>,
pub rope_cos: Vec<u8>,
pub rope_sin: Vec<u8>,
pub mask: Vec<u8>,
pub offset: Option<Vec<u8>>,
}
#[cfg(feature = "cuda")]
pub fn captured_output_to_f32(
output: &Arc<std::sync::RwLock<fuel_memory::Storage>>,
) -> fuel_ir::error::Result<Vec<f32>> {
use fuel_memory::BackendStorage;
let guard = output.read().map_err(|_| {
fuel_ir::Error::Msg("captured decode output storage lock poisoned".into()).bt()
})?;
let bytes: Vec<u8> = match &guard.inner {
BackendStorage::Cpu(c) => c.bytes().to_vec(),
BackendStorage::Cuda(c) => c.to_cpu_bytes()?,
#[allow(unreachable_patterns)]
other => {
return Err(fuel_ir::Error::Msg(format!(
"captured decode output is BackendStorage::{:?}, expected Cpu or Cuda",
std::mem::discriminant(other),
))
.bt());
}
};
Ok(bytemuck::cast_slice::<u8, f32>(&bytes).to_vec())
}