diff --git a/fusor-ml/core/src/compute_graph/resolve/execution.rs b/fusor-ml/core/src/compute_graph/resolve/execution.rs index 26c7cfde3..2f4c67d76 100644 --- a/fusor-ml/core/src/compute_graph/resolve/execution.rs +++ b/fusor-ml/core/src/compute_graph/resolve/execution.rs @@ -636,6 +636,10 @@ impl Resolver { if !self.execution_graph.contains_node(node_idx) { return; } + let inner_idx = self.execution_graph[node_idx].inner_idx; + if self.targets.contains(&inner_idx) { + return; + } if self .execution_graph .neighbors_directed(node_idx, petgraph::Direction::Outgoing) diff --git a/fusor-ml/core/src/compute_graph/resolve/run.rs b/fusor-ml/core/src/compute_graph/resolve/run.rs index 12cc100c3..fed7c10a2 100644 --- a/fusor-ml/core/src/compute_graph/resolve/run.rs +++ b/fusor-ml/core/src/compute_graph/resolve/run.rs @@ -534,7 +534,24 @@ impl Resolver { let mut dispatch_index = 0usize; let mut command_index = 0usize; let dispatches_per_pass = dispatches_per_pass(total_kernels); + let dispatches_per_submit = dispatches_per_submit(total_kernels, device.backend()); + let wait_after_chunk_submit = device.backend() == wgpu::Backend::Metal; + let mut dispatches_in_submit = 0usize; + let mut encoder_has_commands = false; while command_index < commands.len() { + if encoder_has_commands && dispatches_in_submit >= dispatches_per_submit { + let next_encoder = resolver_command_encoder(&device); + let ready_encoder = std::mem::replace(&mut command_encoder, next_encoder); + submit_resolver_encoder( + &device, + ready_encoder, + wait_after_chunk_submit, + host_trace, + &mut host_profile, + ); + encoder_has_commands = false; + dispatches_in_submit = 0; + } match &commands[command_index] { CommandRecord::CopyBuffer(copy) => { command_encoder.copy_buffer_to_buffer( @@ -544,6 +561,7 @@ impl Resolver { copy.destination_offset, copy.size, ); + encoder_has_commands = true; command_index += 1; } CommandRecord::Dispatch(_) => { @@ -568,6 +586,8 @@ impl Resolver { record.dispatch.run(&mut pass); } dispatch_index += 1; + dispatches_in_submit += 1; + encoder_has_commands = true; command_index += 1; continue; } @@ -582,6 +602,9 @@ impl Resolver { if pass_dispatches >= dispatches_per_pass { break; } + if dispatches_in_submit >= dispatches_per_submit { + break; + } let CommandRecord::Dispatch(record) = &commands[command_index] else { break; }; @@ -595,8 +618,10 @@ impl Resolver { pass.write_timestamp(query_set, (dispatch_index * 2 + 1) as u32); } dispatch_index += 1; + dispatches_in_submit += 1; command_index += 1; pass_dispatches += 1; + encoder_has_commands = true; } } } @@ -625,11 +650,13 @@ impl Resolver { let tail_result = tail(&data, &mut command_encoder); // Submit any remaining commands. - let submit_start = host_trace.then(Instant::now); - device.wgpu_queue().submit(Some(command_encoder.finish())); - if let Some(start) = submit_start { - host_profile.submit += start.elapsed(); - } + submit_resolver_encoder( + &device, + command_encoder, + false, + host_trace, + &mut host_profile, + ); #[cfg(not(target_arch = "wasm32"))] { if let Some((_, _, readback_buffer, raw_query_size)) = &query_resources { @@ -726,3 +753,45 @@ fn dispatches_per_pass(total_kernels: usize) -> usize { if total_kernels >= 1024 { 1 } else { usize::MAX } } + +fn dispatches_per_submit(total_kernels: usize, backend: wgpu::Backend) -> usize { + if let Ok(value) = std::env::var("FUSOR_RESOLVE_DISPATCHES_PER_SUBMIT") + && let Ok(parsed) = value.parse::() + && parsed > 0 + { + return parsed; + } + + // Chunked submits exist to bound in-flight memory on giant training + // graphs; small inference graphs must stay a single submit. + if backend == wgpu::Backend::Metal && total_kernels >= 1024 { + 256 + } else { + usize::MAX + } +} + +fn resolver_command_encoder(device: &crate::Device) -> wgpu::CommandEncoder { + device + .wgpu_device() + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("Resolver Encoder"), + }) +} + +fn submit_resolver_encoder( + device: &crate::Device, + command_encoder: wgpu::CommandEncoder, + wait: bool, + host_trace: bool, + host_profile: &mut ResolveHostProfile, +) { + let submit_start = host_trace.then(Instant::now); + device.wgpu_queue().submit(Some(command_encoder.finish())); + if wait { + device.poll_wait(); + } + if let Some(start) = submit_start { + host_profile.submit += start.elapsed(); + } +} diff --git a/fusor-ml/core/src/nary_direct.rs b/fusor-ml/core/src/nary_direct.rs index 245507df8..0d2de145b 100644 --- a/fusor-ml/core/src/nary_direct.rs +++ b/fusor-ml/core/src/nary_direct.rs @@ -849,6 +849,31 @@ fn emit_function(function: &NaryFunction, values: &mut [(ValueTile, DataTypeEnum values[1].0.clone(), function.output_type, ), + NaryOp::Less => values[0].0.clone().compare( + tile_ir::TileCompareOp::Lt, + values[1].0.clone(), + function.output_type, + ), + NaryOp::Equal => values[0].0.clone().compare( + tile_ir::TileCompareOp::Eq, + values[1].0.clone(), + function.output_type, + ), + NaryOp::NotEqual => values[0].0.clone().compare( + tile_ir::TileCompareOp::Ne, + values[1].0.clone(), + function.output_type, + ), + NaryOp::Greater => values[0].0.clone().compare( + tile_ir::TileCompareOp::Gt, + values[1].0.clone(), + function.output_type, + ), + NaryOp::GreaterEqual => values[0].0.clone().compare( + tile_ir::TileCompareOp::Ge, + values[1].0.clone(), + function.output_type, + ), NaryOp::AddConst(scalar) => values[0].0.clone().binary( tile_ir::TileBinaryOp::Add, tile_literal(scalar).cast_to(values[0].1), diff --git a/fusor-ml/core/src/nary_wise.rs b/fusor-ml/core/src/nary_wise.rs index 859799931..488b8e3d6 100644 --- a/fusor-ml/core/src/nary_wise.rs +++ b/fusor-ml/core/src/nary_wise.rs @@ -86,6 +86,16 @@ pub(crate) enum NaryOp { /// index-dependent masking (e.g. composed causal attention compares the /// kv position against the query position). LessEqual, + /// Binary `a < b` comparison producing 1/0 in the output type. + Less, + /// Binary `a == b` comparison producing 1/0 in the output type. + Equal, + /// Binary `a != b` comparison producing 1/0 in the output type. + NotEqual, + /// Binary `a > b` comparison producing 1/0 in the output type. + Greater, + /// Binary `a >= b` comparison producing 1/0 in the output type. + GreaterEqual, AddConst(NaryScalar), SubConst(NaryScalar), RSubConst(NaryScalar), diff --git a/fusor-ml/core/src/pair_wise.rs b/fusor-ml/core/src/pair_wise.rs index 32066a2f6..99b6038de 100644 --- a/fusor-ml/core/src/pair_wise.rs +++ b/fusor-ml/core/src/pair_wise.rs @@ -1,7 +1,7 @@ use std::ops::{Add, Div, Mul, Sub}; use crate::{ - Tensor, + DataType, Tensor, nary_wise::{NaryFunction, NaryOp}, }; @@ -81,3 +81,51 @@ macro_rules! impl_pairwise_method { } impl_pairwise_method!(pow, NaryOp::Pow, "pow", pow_, |a, b| a.pow(&b)); + +/// Emit a tensor-tensor comparison method producing 1/0 in the output type +/// `D`, mirroring the scalar comparisons in `element_wise`. +macro_rules! impl_pairwise_cmp { + ($(#[$meta:meta])* $method:ident, $nary_op:expr, $op_name:literal) => { + impl Tensor { + $(#[$meta])* + pub fn $method(&self, other: &Self) -> Tensor { + assert_eq!(self.datatype(), other.datatype()); + self.binary_nary( + other, + NaryFunction::binary( + Some($op_name.to_string()), + $nary_op, + self.datatype(), + other.datatype(), + D::DATA_TYPE, + ), + ) + } + } + }; +} + +impl_pairwise_cmp!( + /// Element-wise `self == other` returning 1 for true and 0 for false. + eq_tensor, NaryOp::Equal, "eq" +); +impl_pairwise_cmp!( + /// Element-wise `self != other` returning 1 for true and 0 for false. + ne_tensor, NaryOp::NotEqual, "ne" +); +impl_pairwise_cmp!( + /// Element-wise `self < other` returning 1 for true and 0 for false. + lt_tensor, NaryOp::Less, "lt" +); +impl_pairwise_cmp!( + /// Element-wise `self <= other` returning 1 for true and 0 for false. + lte_tensor, NaryOp::LessEqual, "lte" +); +impl_pairwise_cmp!( + /// Element-wise `self > other` returning 1 for true and 0 for false. + gt_tensor, NaryOp::Greater, "gt" +); +impl_pairwise_cmp!( + /// Element-wise `self >= other` returning 1 for true and 0 for false. + gte_tensor, NaryOp::GreaterEqual, "gte" +); diff --git a/fusor-ml/fusor/src/autograd/composite.rs b/fusor-ml/fusor/src/autograd/composite.rs new file mode 100644 index 000000000..c07a62554 --- /dev/null +++ b/fusor-ml/fusor/src/autograd/composite.rs @@ -0,0 +1,871 @@ +use crate::MaskKind; +use fusor_types::SlidingWindow; + +use super::*; + +impl Tensor { + fn softmax_composite(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + let input_shape = self.shape(); + let max_values = self.max_keepdim_any::(axis); + let shifted = self.sub(&max_values.broadcast_as(input_shape)); + let exp_values = shifted.exp(); + let normalization = exp_values + .sum_keepdim_any::(axis) + .broadcast_as(input_shape); + exp_values.div(&normalization) + } + + fn rms_norm_composite( + &self, + weight: &Tensor, + bias: Option<&Tensor>, + eps: f32, + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.layer_norm_composite::(weight, bias, eps, false) + } + + fn layer_norm_composite( + &self, + weight: &Tensor, + bias: Option<&Tensor>, + eps: f32, + remove_mean: bool, + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + let centered = if remove_mean { + let mean = self.mean_keepdim_any::(R - 1); + self.sub(&mean.broadcast_as(self.shape())) + } else { + self.clone() + }; + let variance = centered.sqr().mean_keepdim_any::(R - 1); + let std = variance.add_scalar(eps).sqrt(); + let normalized = centered.div(&std.broadcast_as(self.shape())); + let scaled = normalized.mul(&weight.broadcast_as(self.shape())); + if let Some(bias) = bias { + scaled.add(&bias.broadcast_as(self.shape())) + } else { + scaled + } + } + + pub fn softmax(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.softmax_composite::(axis) + } + + pub fn softmax_last_dim(&self) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.softmax::(R - 1) + } + + pub fn softmax_slow(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.softmax::(axis) + } + + pub fn softmax_slow_last_dim(&self) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.softmax_slow::(R - 1) + } + + pub fn layer_norm( + &self, + weight: &Tensor, + bias: Option<&Tensor>, + eps: f32, + remove_mean: bool, + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.layer_norm_composite::(weight, bias, eps, remove_mean) + } + + pub fn rms_norm(&self, weight: &Tensor, eps: f32) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.rms_norm_composite::(weight, None, eps) + } + + pub fn softmax_last_dim_fused(&self) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::gpu::Tensor: crate::gpu::LastRankInner, + { + let value = self + .value + .softmax_last_dim_fused::() + .to_concrete(); + self.replay_unary("softmax_last_dim_fused", value, |input| { + input.softmax_last_dim::() + }) + } + + pub fn rms_norm_fused( + &self, + weight: &Tensor, + bias: Option<&Tensor>, + eps: f32, + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + as crate::gpu::LastRankInner>::LastRank: + crate::gpu::NextRankInner>, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::MulOp: crate::cpu::SimdBinaryOp, + crate::DivOp: crate::cpu::SimdBinaryOp, + crate::AddOp: crate::cpu::SimdBinaryOp, + crate::SqrtOp: crate::cpu::SimdUnaryOp, + (crate::gpu::Tensor, crate::gpu::Tensor): crate::gpu::MaxRank, + { + let value = self.value.rms_norm_fused::( + &weight.value, + bias.as_ref().map(|bias| &bias.value), + eps, + ); + if let Some(bias) = bias { + self.replay_ternary( + weight, + bias, + "rms_norm_fused", + value, + move |input, weight, bias| { + input.rms_norm_composite::(&weight, Some(&bias), eps) + }, + ) + } else { + self.replay_binary(weight, "rms_norm_fused", value, move |input, weight| { + input.rms_norm_composite::(&weight, None, eps) + }) + } + } + + pub fn rms_norm_fused_no_bias( + &self, + weight: &Tensor, + eps: f32, + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + as crate::gpu::LastRankInner>::LastRank: + crate::gpu::NextRankInner>, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::MulOp: crate::cpu::SimdBinaryOp, + crate::DivOp: crate::cpu::SimdBinaryOp, + crate::AddOp: crate::cpu::SimdBinaryOp, + crate::SqrtOp: crate::cpu::SimdUnaryOp, + (crate::gpu::Tensor, crate::gpu::Tensor): crate::gpu::MaxRank, + { + self.rms_norm_fused::(weight, None, eps) + } + + pub fn rms_norm_residual_fused( + &self, + residual: &Self, + weight: &Tensor, + bias: Option<&Tensor>, + eps: f32, + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + as crate::gpu::LastRankInner>::LastRank: + crate::gpu::NextRankInner>, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::MulOp: crate::cpu::SimdBinaryOp, + crate::DivOp: crate::cpu::SimdBinaryOp, + crate::AddOp: crate::cpu::SimdBinaryOp, + crate::SqrtOp: crate::cpu::SimdUnaryOp, + (crate::gpu::Tensor, crate::gpu::Tensor): crate::gpu::MaxRank, + { + let value = self.value.rms_norm_residual_fused::( + &residual.value, + &weight.value, + bias.as_ref().map(|bias| &bias.value), + eps, + ); + match bias { + None => self.replay_ternary( + residual, + weight, + "rms_norm_residual_fused", + value, + move |input, residual, weight| { + input + .add(&residual) + .rms_norm_composite::(&weight, None, eps) + }, + ), + Some(bias) => self.replay_quaternary( + residual, + weight, + bias, + "rms_norm_residual_fused", + value, + move |input, residual, weight, bias| { + input + .add(&residual) + .rms_norm_composite::(&weight, Some(&bias), eps) + }, + ), + } + } + + pub fn layer_norm_last_dim_fused( + &self, + weight: &Tensor, + bias: Option<&Tensor>, + eps: f32, + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + as crate::gpu::LastRankInner>::LastRank: + crate::gpu::NextRankInner>, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::AddOp: crate::cpu::SimdBinaryOp, + crate::SubOp: crate::cpu::SimdBinaryOp, + crate::MulOp: crate::cpu::SimdBinaryOp, + crate::DivOp: crate::cpu::SimdBinaryOp, + crate::SqrtOp: crate::cpu::SimdUnaryOp, + { + let value = self.value.layer_norm_last_dim_fused::( + &weight.value, + bias.as_ref().map(|bias| &bias.value), + eps, + ); + let mut param_shape = [1usize; R]; + param_shape[R - 1] = self.shape()[R - 1]; + if let Some(bias) = bias { + self.replay_ternary( + weight, + bias, + "layer_norm_last_dim_fused", + value, + move |input, weight, bias| { + input.layer_norm_composite::( + &weight.reshape(param_shape), + Some(&bias.reshape(param_shape)), + eps, + true, + ) + }, + ) + } else { + self.replay_binary( + weight, + "layer_norm_last_dim_fused", + value, + move |input, weight| { + input.layer_norm_composite::( + &weight.reshape(param_shape), + None, + eps, + true, + ) + }, + ) + } + } + + fn pad_spatial(&self, padding: [usize; DIFF]) -> Self { + let mut padded = self.clone(); + for (i, padding) in padding.into_iter().enumerate() { + padded = padded.pad_axis(R - DIFF + i, padding); + } + padded + } + + fn conv_output_shape( + input_shape: [usize; R], + out_channels: usize, + kernel: [usize; DIFF], + padding: [usize; DIFF], + strides: [usize; DIFF], + ) -> [usize; R] { + let spatial_start = R - DIFF; + let mut output_shape = input_shape; + output_shape[1] = out_channels; + for i in 0..DIFF { + let padded_len = input_shape[spatial_start + i] + 2 * padding[i]; + output_shape[spatial_start + i] = (padded_len - kernel[i]) / strides[i] + 1; + } + output_shape + } + + /// Pad + sliding-window view + flatten to one matmul row per output + /// location: `(batch * out_spatial, in_channels * kernel_size)`. + fn conv_windows_flat( + &self, + kernel: [usize; DIFF], + padding: [usize; DIFF], + strides: [usize; DIFF], + ) -> Tensor<2> + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank, + { + let input_shape = self.shape(); + let spatial_start = R - DIFF; + let output_shape = Self::conv_output_shape(input_shape, 0, kernel, padding, strides); + let windows: [SlidingWindow; DIFF] = std::array::from_fn(|i| { + SlidingWindow::new(spatial_start + i, kernel[i], strides[i]) + }); + let windows: Tensor = self.pad_spatial(padding).sliding_window_view(windows); + let permutation: [usize; R2] = std::array::from_fn(|index| { + if index == 0 { + 0 + } else if index <= DIFF { + index + 1 + } else if index == DIFF + 1 { + 1 + } else { + index + } + }); + let out_spatial_size: usize = output_shape[spatial_start..].iter().product(); + let kernel_size: usize = kernel.iter().product(); + windows.permute(permutation).reshape([ + input_shape[0] * out_spatial_size, + input_shape[1] * kernel_size, + ]) + } + + /// Reshape the `(batch * out_spatial, out_channels)` matmul output back to + /// `(batch, out_channels, ...out_spatial)` and add the broadcast bias. + fn conv_reassemble( + output: Tensor<2>, + bias: Option<&Tensor<1>>, + output_shape: [usize; R], + ) -> Self { + let out_channels = output_shape[1]; + let output: Tensor = output.reshape(std::array::from_fn(|axis| { + if axis == 0 { + output_shape[0] + } else if axis <= DIFF { + output_shape[axis + 1] + } else { + out_channels + } + })); + let permutation: [usize; R] = std::array::from_fn(|index| { + if index == 0 { + 0 + } else if index == 1 { + DIFF + 1 + } else { + index - 1 + } + }); + let output = output.permute(permutation); + if let Some(bias) = bias { + let bias_shape: [usize; R] = + std::array::from_fn(|axis| if axis == 1 { out_channels } else { 1 }); + output.add(&bias.reshape(bias_shape).broadcast_as(output_shape)) + } else { + output + } + } + + fn conv_composite( + &self, + weight: &Tensor, + bias: Option<&Tensor<1>>, + padding: [usize; DIFF], + strides: [usize; DIFF], + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank, + { + assert_eq!( + R, + 2 + DIFF, + "Conv expects (batch, channels, ...spatial) format where R = 2 + DIFF" + ); + let input_shape = self.shape(); + let weight_shape = weight.shape(); + let spatial_start = R - DIFF; + let in_channels = input_shape[1]; + let out_channels = weight_shape[0]; + assert_eq!( + weight_shape[1], in_channels, + "Weight in_channels must match input in_channels" + ); + + let kernel: [usize; DIFF] = std::array::from_fn(|i| weight_shape[spatial_start + i]); + let kernel_size: usize = kernel.iter().product(); + let windows_flat = self.conv_windows_flat::(kernel, padding, strides); + let weight_t = weight + .reshape([out_channels, in_channels * kernel_size]) + .transpose(0, 1); + let output = windows_flat.mat_mul_internal(&weight_t); + + let output_shape = + Self::conv_output_shape(input_shape, out_channels, kernel, padding, strides); + Self::conv_reassemble::(output, bias, output_shape) + } + + pub fn conv( + &self, + weight: &Tensor, + bias: Option<&Tensor<1>>, + padding: [usize; DIFF], + strides: [usize; DIFF], + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank, + { + let value = self.value.conv::( + &weight.value, + bias.map(|bias| &bias.value), + padding, + strides, + ); + match bias { + None => self.replay_binary(weight, "conv", value, move |input, weight| { + input.conv_composite::(&weight, None, padding, strides) + }), + Some(bias) => self.replay_ternary( + weight, + bias, + "conv", + value, + move |input, weight, bias| { + input.conv_composite::( + &weight, + Some(&bias), + padding, + strides, + ) + }, + ), + } + } + + fn grouped_conv_composite( + &self, + weight: &Tensor, + bias: Option<&Tensor<1>>, + padding: [usize; DIFF], + strides: [usize; DIFF], + groups: usize, + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank, + { + assert_eq!(R, 2 + DIFF); + let input_shape = self.shape(); + let weight_shape = weight.shape(); + let spatial_start = R - DIFF; + let batch = input_shape[0]; + let in_channels = input_shape[1]; + let out_channels = weight_shape[0]; + assert_eq!(in_channels % groups, 0); + assert_eq!(out_channels % groups, 0); + let in_ch_per_group = in_channels / groups; + let out_ch_per_group = out_channels / groups; + assert_eq!(weight_shape[1], in_ch_per_group); + + let kernel: [usize; DIFF] = std::array::from_fn(|i| weight_shape[spatial_start + i]); + let kernel_size: usize = kernel.iter().product(); + let output_shape = + Self::conv_output_shape(input_shape, out_channels, kernel, padding, strides); + let out_spatial_size: usize = output_shape[spatial_start..].iter().product(); + + let windows_grouped = self + .conv_windows_flat::(kernel, padding, strides) + .reshape([ + batch * out_spatial_size, + groups, + in_ch_per_group * kernel_size, + ]) + .transpose(0, 1); + let weight_grouped_t = weight + .reshape([groups, out_ch_per_group, in_ch_per_group * kernel_size]) + .transpose(1, 2); + let output = windows_grouped + .mat_mul_internal(&weight_grouped_t) + .transpose(0, 1) + .reshape([batch * out_spatial_size, out_channels]); + Self::conv_reassemble::(output, bias, output_shape) + } + + pub fn grouped_conv( + &self, + weight: &Tensor, + bias: Option<&Tensor<1>>, + padding: [usize; DIFF], + strides: [usize; DIFF], + groups: usize, + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank, + { + let value = self.value.grouped_conv::( + &weight.value, + bias.map(|bias| &bias.value), + padding, + strides, + groups, + ); + match bias { + None => self.replay_binary(weight, "grouped_conv", value, move |input, weight| { + input.grouped_conv_composite::( + &weight, None, padding, strides, groups, + ) + }), + Some(bias) => self.replay_ternary( + weight, + bias, + "grouped_conv", + value, + move |input, weight, bias| { + input.grouped_conv_composite::( + &weight, + Some(&bias), + padding, + strides, + groups, + ) + }, + ), + } + } +} + +impl Tensor<4> { + fn rotate_half(&self) -> Tensor<4> { + let [batch, heads, sequence_length, embed] = self.shape(); + let half = embed / 2; + let first_half = self.narrow(3, 0, half); + let second_half = self.narrow(3, half, embed - half).mul_scalar(-1.0); + let graph = self.graph(); + let device = self.device(); + let zeros = Tensor::zeros(&graph, &device, [batch, heads, sequence_length, embed]); + let combined = zeros.slice_assign( + [0..batch, 0..heads, 0..sequence_length, 0..half], + &second_half, + ); + combined.slice_assign( + [0..batch, 0..heads, 0..sequence_length, half..embed], + &first_half, + ) + } + + fn rope_interleaved_composite(&self, cos: &Tensor<2>, sin: &Tensor<2>) -> Tensor<4> { + assert_same_graph(self, cos); + assert_same_graph(self, sin); + + let [batch, heads, sequence_length, embed] = self.shape(); + let half = embed / 2; + let cos = cos + .narrow(0, 0, sequence_length) + .reshape([sequence_length, half, 1]) + .broadcast_as([batch, heads, sequence_length, half, 1]); + let sin = sin + .narrow(0, 0, sequence_length) + .reshape([sequence_length, half, 1]) + .broadcast_as([batch, heads, sequence_length, half, 1]); + let x = self.reshape([batch, heads, sequence_length, half, 2]); + let x0 = x.narrow(4, 0, 1); + let x1 = x.narrow(4, 1, 1); + let y0 = x0.mul(&cos).sub(&x1.mul(&sin)); + let y1 = x0.mul(&sin).add(&x1.mul(&cos)); + let graph = self.graph(); + let device = self.device(); + let zeros = Tensor::zeros(&graph, &device, [batch, heads, sequence_length, half, 2]); + let combined = + zeros.slice_assign([0..batch, 0..heads, 0..sequence_length, 0..half, 0..1], &y0); + combined + .slice_assign([0..batch, 0..heads, 0..sequence_length, 0..half, 1..2], &y1) + .flatten_last_n::<1, 4>() + } + + pub fn rope(&self, cos: &Tensor<2>, sin: &Tensor<2>) -> Tensor<4> { + assert_same_graph(self, cos); + assert_same_graph(self, sin); + + let [batch, heads, sequence_length, embed] = self.shape(); + let half = embed / 2; + let graph = self.graph(); + let device = self.device(); + let cos_base = cos.narrow(0, 0, sequence_length); + let sin_base = sin.narrow(0, 0, sequence_length); + let cos = Tensor::zeros(&graph, &device, [sequence_length, embed]) + .slice_assign([0..sequence_length, 0..half], &cos_base) + .slice_assign([0..sequence_length, half..embed], &cos_base) + .unsqueeze_dims::<2, 4>([0, 1]) + .broadcast_as([batch, heads, sequence_length, embed]); + let sin = Tensor::zeros(&graph, &device, [sequence_length, embed]) + .slice_assign([0..sequence_length, 0..half], &sin_base) + .slice_assign([0..sequence_length, half..embed], &sin_base) + .unsqueeze_dims::<2, 4>([0, 1]) + .broadcast_as([batch, heads, sequence_length, embed]); + let rotated = self.rotate_half(); + self.mul(&cos).add(&rotated.mul(&sin)) + } + + pub fn rope_interleaved(&self, cos: &Tensor<2>, sin: &Tensor<2>) -> Tensor<4> { + self.rope_interleaved_composite(cos, sin) + } + + pub fn flash_attention( + &self, + k: &Tensor<4>, + v: &Tensor<4>, + scale: f32, + mask: Option<(&RawTensor<2, f32>, MaskKind)>, + ) -> Tensor<4> { + let value = self.value.flash_attention(&k.value, &v.value, scale, mask); + let mask_value = mask.map(|(mask, kind)| (mask.clone(), kind)); + self.replay_ternary(k, v, "flash_attention", value, move |q, k, v| { + q.flash_attention_composite(&k, &v, scale, mask_value.as_ref()) + }) + } + + pub fn rope_fused(&self, cos: &Tensor<2>, sin: &Tensor<2>) -> Tensor<4> { + assert_same_graph(self, cos); + assert_same_graph(self, sin); + + let value = self.value.rope_fused(&cos.value, &sin.value).to_concrete(); + self.replay_ternary(cos, sin, "rope_fused", value, |input, cos, sin| { + input.rope_interleaved_composite(&cos, &sin) + }) + } + + pub fn rope_normal_fused(&self, cos: &Tensor<2>, sin: &Tensor<2>) -> Tensor<4> { + assert_same_graph(self, cos); + assert_same_graph(self, sin); + + let value = self + .value + .rope_normal_fused(&cos.value, &sin.value) + .to_concrete(); + self.replay_ternary(cos, sin, "rope_normal_fused", value, |input, cos, sin| { + input.rope(&cos, &sin) + }) + } + + pub fn rope_pair_fused( + &self, + k: &Self, + cos: &Tensor<2>, + sin: &Tensor<2>, + ) -> (Tensor<4>, Tensor<4>) { + let (q_value, k_value) = self.value.rope_pair_fused(&k.value, &cos.value, &sin.value); + ( + self.replay_ternary( + cos, + sin, + "rope_pair_fused", + q_value.to_concrete(), + |input, cos, sin| input.rope_interleaved_composite(&cos, &sin), + ), + k.replay_ternary( + cos, + sin, + "rope_pair_fused", + k_value.to_concrete(), + |input, cos, sin| input.rope_interleaved_composite(&cos, &sin), + ), + ) + } + + pub fn rope_normal_pair_fused( + &self, + k: &Self, + cos: &Tensor<2>, + sin: &Tensor<2>, + ) -> (Tensor<4>, Tensor<4>) { + let (q_value, k_value) = self + .value + .rope_normal_pair_fused(&k.value, &cos.value, &sin.value); + ( + self.replay_ternary( + cos, + sin, + "rope_normal_pair_fused", + q_value.to_concrete(), + |input, cos, sin| input.rope(&cos, &sin), + ), + k.replay_ternary( + cos, + sin, + "rope_normal_pair_fused", + k_value.to_concrete(), + |input, cos, sin| input.rope(&cos, &sin), + ), + ) + } + + fn rope_cache_tables( + &self, + cache: &crate::RopeCache, + start_pos: usize, + ) -> (Tensor<2>, Tensor<2>) { + let seq_len = self.shape()[2]; + let graph = self.graph(); + let table = |table: RawTensor<2, f32>| { + Tensor::constant_from_raw(&graph, table.narrow(0, start_pos, seq_len).to_concrete()) + }; + (table(cache.cos().clone()), table(cache.sin().clone())) + } + + /// Autograd companion of [`crate::RopeCache::forward`]: applies normal + /// RoPE to `self` (q) and `k` from the cache's tables at `start_pos`. + pub fn rope_cache_forward( + &self, + k: &Self, + cache: &crate::RopeCache, + start_pos: usize, + ) -> (Tensor<4>, Tensor<4>) { + let (cos, sin) = self.rope_cache_tables(cache, start_pos); + self.rope_normal_pair_fused(k, &cos, &sin) + } + + /// Autograd companion of [`crate::RopeCache::forward_interleaved`]. + pub fn rope_cache_forward_interleaved( + &self, + k: &Self, + cache: &crate::RopeCache, + start_pos: usize, + ) -> (Tensor<4>, Tensor<4>) { + let (cos, sin) = self.rope_cache_tables(cache, start_pos); + self.rope_pair_fused(k, &cos, &sin) + } + + pub fn upsample_nearest2d(&self, scale_h: usize, scale_w: usize) -> Tensor<4> { + let value = self.value.upsample_nearest2d(scale_h, scale_w); + self.replay_unary("upsample_nearest2d", value, move |input| { + let [b, c, h, w] = input.shape(); + input + .reshape([b, c, h, 1, w, 1]) + .broadcast_as([b, c, h, scale_h, w, scale_w]) + .reshape([b, c, h * scale_h, w * scale_w]) + }) + } + + pub(super) fn flash_attention_composite( + &self, + k: &Tensor<4>, + v: &Tensor<4>, + scale: f32, + mask: Option<&(RawTensor<2, f32>, MaskKind)>, + ) -> Tensor<4> { + let q_shape = self.shape(); + let k_shape = k.shape(); + let batch = q_shape[0]; + let num_heads = q_shape[1]; + let q_seq_len = q_shape[2]; + let head_dim = q_shape[3]; + let num_kv_heads = k_shape[1]; + let kv_seq_len = k_shape[2]; + assert!( + num_heads.is_multiple_of(num_kv_heads), + "Number of Q heads ({num_heads}) must be divisible by number of K/V heads ({num_kv_heads})" + ); + + let num_key_value_groups = num_heads / num_kv_heads; + let (k_expanded, v_expanded) = if num_key_value_groups > 1 { + let k_broadcast = k + .reshape([batch, num_kv_heads, 1, kv_seq_len, head_dim]) + .broadcast_as([ + batch, + num_kv_heads, + num_key_value_groups, + kv_seq_len, + head_dim, + ]); + let v_broadcast = v + .reshape([batch, num_kv_heads, 1, kv_seq_len, head_dim]) + .broadcast_as([ + batch, + num_kv_heads, + num_key_value_groups, + kv_seq_len, + head_dim, + ]); + ( + k_broadcast.reshape([batch, num_heads, kv_seq_len, head_dim]), + v_broadcast.reshape([batch, num_heads, kv_seq_len, head_dim]), + ) + } else { + (k.clone(), v.clone()) + }; + + let scores = self + .mat_mul_internal(&k_expanded.transpose(2, 3)) + .div_scalar(scale.recip()); + let masked_scores = if let Some((mask, kind)) = mask { + let mask_tensor = Tensor::constant_from_raw(&self.graph(), mask.clone()); + let mask_4d = match kind { + // Causal falls back to QKMask semantics here: the provided mask is the + // [seq_len, seq_len] (== [q_seq_len, kv_seq_len]) additive causal mask. + MaskKind::QKMask | MaskKind::Causal => { + assert_eq!(mask_tensor.shape(), [q_seq_len, kv_seq_len]); + mask_tensor.reshape([1, 1, q_seq_len, kv_seq_len]) + } + MaskKind::BatchKeyMask => { + assert_eq!(mask_tensor.shape(), [batch, kv_seq_len]); + mask_tensor.reshape([batch, 1, 1, kv_seq_len]) + } + }; + scores.add(&mask_4d.broadcast_as([batch, num_heads, q_seq_len, kv_seq_len])) + } else { + scores + }; + masked_scores + .softmax_last_dim::<3>() + .mat_mul_internal(&v_expanded) + } +} diff --git a/fusor-ml/fusor/src/autograd/elementwise.rs b/fusor-ml/fusor/src/autograd/elementwise.rs new file mode 100644 index 000000000..c2d9f3537 --- /dev/null +++ b/fusor-ml/fusor/src/autograd/elementwise.rs @@ -0,0 +1,673 @@ +use super::*; + +impl Tensor { + pub fn add(&self, rhs: &Self) -> Self { + self.binary_op( + rhs, + (self.value.clone() + rhs.value.clone()).to_concrete(), + |grad, _, _| vec![grad.clone().to_concrete(), grad.to_concrete()], + ) + } + + pub fn add_(&self, second: &Tensor) -> Tensor { + let out_shape: [usize; R3] = + crate::composite::broadcast_shapes(&self.shape(), &second.shape()); + let lhs = self.broadcast_as(out_shape); + let rhs = second.broadcast_as(out_shape); + lhs.add(&rhs) + } + + pub fn sub(&self, rhs: &Self) -> Self { + self.binary_op( + rhs, + (self.value.clone() - rhs.value.clone()).to_concrete(), + |grad, _, _| vec![grad.clone().to_concrete(), (-grad).to_concrete()], + ) + } + + pub fn sub_(&self, second: &Tensor) -> Tensor { + let out_shape: [usize; R3] = + crate::composite::broadcast_shapes(&self.shape(), &second.shape()); + let lhs = self.broadcast_as(out_shape); + let rhs = second.broadcast_as(out_shape); + lhs.sub(&rhs) + } + + pub fn mul(&self, rhs: &Self) -> Self { + self.binary_op( + rhs, + (self.value.clone() * rhs.value.clone()).to_concrete(), + |grad, lhs, rhs| { + vec![ + (grad.clone() * rhs).to_concrete(), + (grad * lhs).to_concrete(), + ] + }, + ) + } + + pub fn mul_(&self, second: &Tensor) -> Tensor { + let out_shape: [usize; R3] = + crate::composite::broadcast_shapes(&self.shape(), &second.shape()); + let lhs = self.broadcast_as(out_shape); + let rhs = second.broadcast_as(out_shape); + lhs.mul(&rhs) + } + + pub fn div(&self, rhs: &Self) -> Self { + self.binary_op( + rhs, + (self.value.clone() / rhs.value.clone()).to_concrete(), + |grad, lhs, rhs| { + let lhs_grad = (grad.clone() / rhs.clone()).to_concrete(); + let rhs_grad = (-((grad * lhs) / rhs.sqr().to_concrete())).to_concrete(); + vec![lhs_grad, rhs_grad] + }, + ) + } + + pub fn div_(&self, second: &Tensor) -> Tensor { + let out_shape: [usize; R3] = + crate::composite::broadcast_shapes(&self.shape(), &second.shape()); + let lhs = self.broadcast_as(out_shape); + let rhs = second.broadcast_as(out_shape); + lhs.div(&rhs) + } + + pub fn pow(&self, rhs: &Self) -> Self { + self.binary_op( + rhs, + self.value.pow(&rhs.value).to_concrete(), + |grad, lhs, rhs| { + let rhs_minus_one = rhs.sub_scalar(1.0).to_concrete(); + let lhs_power = lhs.pow(&rhs_minus_one).to_concrete(); + let lhs_grad = + ((grad.clone() * rhs.clone()).to_concrete() * lhs_power).to_concrete(); + let rhs_grad = ((grad * lhs.pow(&rhs).to_concrete()).to_concrete() + * lhs.log().to_concrete()) + .to_concrete(); + vec![lhs_grad, rhs_grad] + }, + ) + } + + pub fn pow_(&self, second: &Tensor) -> Tensor { + let out_shape: [usize; R3] = + crate::composite::broadcast_shapes(&self.shape(), &second.shape()); + let lhs = self.broadcast_as(out_shape); + let rhs = second.broadcast_as(out_shape); + lhs.pow(&rhs) + } + + pub fn pow_elementwise(&self, exponent: f32) -> Self { + let input = self.value.clone(); + self.unary_from_value( + self.value.pow_elementwise(exponent).to_concrete(), + move |grad, _| { + let power = input.pow_elementwise(exponent - 1.0).to_concrete(); + (grad * power) + .to_concrete() + .mul_scalar(exponent) + .to_concrete() + }, + ) + } + + pub fn pow_scalar(&self, exponent: f32) -> Self { + self.pow_elementwise(exponent) + } + + pub fn add_scalar(&self, scalar: f32) -> Self { + self.unary_from_value(self.value.add_scalar(scalar), move |grad, _| grad) + } + + pub fn sub_scalar(&self, scalar: f32) -> Self { + self.unary_from_value(self.value.sub_scalar(scalar), move |grad, _| grad) + } + + pub fn mul_scalar(&self, scalar: f32) -> Self { + self.unary_from_value( + self.value.mul_scalar(scalar).to_concrete(), + move |grad, _| grad.mul_scalar(scalar).to_concrete(), + ) + } + + pub fn div_scalar(&self, scalar: f32) -> Self { + self.unary_from_value( + self.value.div_scalar(scalar).to_concrete(), + move |grad, _| grad.div_scalar(scalar).to_concrete(), + ) + } + + pub fn neg(&self) -> Self { + self.unary_from_value((-self.value.clone()).to_concrete(), move |grad, _| { + (-grad).to_concrete() + }) + } + + pub fn sqr(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.sqr().to_concrete(), move |grad, _| { + ((grad * input.clone()).to_concrete().mul_scalar(2.0)).to_concrete() + }) + } + + pub fn abs(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.abs().to_concrete(), move |grad, _| { + let positive = input.mt(0.0).to_concrete(); + let negative = input.lt(0.0).to_concrete(); + ((grad.clone() * positive).to_concrete() - (grad * negative).to_concrete()) + .to_concrete() + }) + } + + pub fn acos(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.acos().to_concrete(), move |grad, _| { + let denom = (RawTensor::splat(&input.device(), 1.0, input.shape()) + - input.sqr().to_concrete()) + .to_concrete() + .sqrt() + .to_concrete(); + (-(grad / denom).to_concrete()).to_concrete() + }) + } + + pub fn acosh(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.acosh().to_concrete(), move |grad, _| { + let lower = input.add_scalar(-1.0).to_concrete().sqrt().to_concrete(); + let upper = input.add_scalar(1.0).to_concrete().sqrt().to_concrete(); + (grad / (lower * upper).to_concrete()).to_concrete() + }) + } + + pub fn approximate_exp(&self) -> Self { + self.unary_from_value( + self.value.approximate_exp().to_concrete(), + move |grad, out| (grad * out).to_concrete(), + ) + } + + pub fn asin(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.asin().to_concrete(), move |grad, _| { + let denom = (RawTensor::splat(&input.device(), 1.0, input.shape()) + - input.sqr().to_concrete()) + .to_concrete() + .sqrt() + .to_concrete(); + (grad / denom).to_concrete() + }) + } + + pub fn asinh(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.asinh().to_concrete(), move |grad, _| { + let denom = input + .sqr() + .add_scalar(1.0) + .to_concrete() + .sqrt() + .to_concrete(); + (grad / denom).to_concrete() + }) + } + + pub fn atan(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.atan().to_concrete(), move |grad, _| { + let denom = input.sqr().add_scalar(1.0).to_concrete(); + (grad / denom).to_concrete() + }) + } + + pub fn atanh(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.atanh().to_concrete(), move |grad, _| { + let denom = (RawTensor::splat(&input.device(), 1.0, input.shape()) + - input.sqr().to_concrete()) + .to_concrete(); + (grad / denom).to_concrete() + }) + } + + pub fn cos(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.cos().to_concrete(), move |grad, _| { + (-(grad * input.sin().to_concrete()).to_concrete()).to_concrete() + }) + } + + pub fn cosh(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.cosh().to_concrete(), move |grad, _| { + (grad * input.sinh().to_concrete()).to_concrete() + }) + } + + pub fn exp2(&self) -> Self { + self.unary_from_value(self.value.exp2().to_concrete(), move |grad, out| { + (grad * out) + .to_concrete() + .mul_scalar(std::f32::consts::LN_2) + .to_concrete() + }) + } + + pub fn less_approximate_exp(&self) -> Self { + self.unary_from_value( + self.value.less_approximate_exp().to_concrete(), + move |grad, out| (grad * out).to_concrete(), + ) + } + + pub fn log2(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.log2().to_concrete(), move |grad, _| { + (grad / input.clone()) + .to_concrete() + .div_scalar(std::f32::consts::LN_2) + .to_concrete() + }) + } + + pub fn sin(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.sin().to_concrete(), move |grad, _| { + (grad * input.cos().to_concrete()).to_concrete() + }) + } + + pub fn sinh(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.sinh().to_concrete(), move |grad, _| { + (grad * input.cosh().to_concrete()).to_concrete() + }) + } + + pub fn tan(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.tan().to_concrete(), move |grad, _| { + let cos = input.cos().to_concrete(); + (grad / (cos.clone() * cos).to_concrete()).to_concrete() + }) + } + + pub fn tanh_exact(&self) -> Self { + self.unary_from_value(self.value.tanh_exact().to_concrete(), move |grad, out| { + let one_minus_sq = (RawTensor::splat(&out.device(), 1.0, out.shape()) + - out.sqr().to_concrete()) + .to_concrete(); + (grad * one_minus_sq).to_concrete() + }) + } + + pub fn cast(&self) -> crate::Tensor + where + f32: crate::CastTo + crate::CastTensor, + D2: crate::SimdElement + crate::DataType + Default, + { + self.value.cast() + } + + pub fn to_concrete(&self) -> Self { + self.unary_from_value(self.value.to_concrete(), move |grad, _| grad) + } + + pub fn relu(&self) -> Self { + self.max_elementwise(0.0) + } + + pub fn clamp(&self, min: f32, max: f32) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.clamp(min, max).to_concrete(), move |grad, _| { + let lower = input.mt(min).to_concrete(); + let upper = input.lt(max).to_concrete(); + ((grad * lower).to_concrete() * upper).to_concrete() + }) + } + + pub fn eq(&self, rhs: f32) -> Self { + self.unary_from_value(self.value.eq(rhs).to_concrete(), move |_, out| { + RawTensor::zeros(&out.device(), out.shape()) + }) + } + + pub fn eq_scalar(&self, rhs: f32) -> Self { + self.eq(rhs) + } + + pub fn eq_tensor(&self, rhs: &Self) -> Self { + assert_same_graph(self, rhs); + self.binary_op( + rhs, + self.value.eq_tensor(&rhs.value).to_concrete(), + move |_, lhs, rhs| { + vec![ + RawTensor::zeros(&lhs.device(), lhs.shape()), + RawTensor::zeros(&rhs.device(), rhs.shape()), + ] + }, + ) + } + + pub fn gt_scalar(&self, rhs: f32) -> Self { + self.unary_from_value(self.value.gt_scalar(rhs).to_concrete(), move |_, out| { + RawTensor::zeros(&out.device(), out.shape()) + }) + } + + pub fn gt_tensor(&self, rhs: &Self) -> Self { + assert_same_graph(self, rhs); + self.binary_op( + rhs, + self.value.gt_tensor(&rhs.value).to_concrete(), + move |_, lhs, rhs| { + vec![ + RawTensor::zeros(&lhs.device(), lhs.shape()), + RawTensor::zeros(&rhs.device(), rhs.shape()), + ] + }, + ) + } + + pub fn gte_scalar(&self, rhs: f32) -> Self { + self.unary_from_value(self.value.gte_scalar(rhs).to_concrete(), move |_, out| { + RawTensor::zeros(&out.device(), out.shape()) + }) + } + + pub fn gte_tensor(&self, rhs: &Self) -> Self { + assert_same_graph(self, rhs); + self.binary_op( + rhs, + self.value.gte_tensor(&rhs.value).to_concrete(), + move |_, lhs, rhs| { + vec![ + RawTensor::zeros(&lhs.device(), lhs.shape()), + RawTensor::zeros(&rhs.device(), rhs.shape()), + ] + }, + ) + } + + pub fn lt(&self, rhs: f32) -> Self { + self.unary_from_value(self.value.lt(rhs).to_concrete(), move |_, out| { + RawTensor::zeros(&out.device(), out.shape()) + }) + } + + pub fn lt_scalar(&self, rhs: f32) -> Self { + self.lt(rhs) + } + + pub fn lt_tensor(&self, rhs: &Self) -> Self { + assert_same_graph(self, rhs); + self.binary_op( + rhs, + self.value.lt_tensor(&rhs.value).to_concrete(), + move |_, lhs, rhs| { + vec![ + RawTensor::zeros(&lhs.device(), lhs.shape()), + RawTensor::zeros(&rhs.device(), rhs.shape()), + ] + }, + ) + } + + pub fn lte(&self, rhs: f32) -> Self { + self.unary_from_value(self.value.lte(rhs).to_concrete(), move |_, out| { + RawTensor::zeros(&out.device(), out.shape()) + }) + } + + pub fn lte_scalar(&self, rhs: f32) -> Self { + self.lte(rhs) + } + + pub fn lte_tensor(&self, rhs: &Self) -> Self { + assert_same_graph(self, rhs); + self.binary_op( + rhs, + self.value.lte_tensor(&rhs.value).to_concrete(), + move |_, lhs, rhs| { + vec![ + RawTensor::zeros(&lhs.device(), lhs.shape()), + RawTensor::zeros(&rhs.device(), rhs.shape()), + ] + }, + ) + } + + pub fn max_elementwise(&self, rhs: f32) -> Self { + let input = self.value.clone(); + self.unary_from_value( + self.value.max_elementwise(rhs).to_concrete(), + move |grad, _| (grad * input.mt(rhs).to_concrete()).to_concrete(), + ) + } + + pub fn max_scalar(&self, rhs: f32) -> Self { + self.max_elementwise(rhs) + } + + pub fn min_elementwise(&self, rhs: f32) -> Self { + let input = self.value.clone(); + self.unary_from_value( + self.value.min_elementwise(rhs).to_concrete(), + move |grad, _| (grad * input.lt(rhs).to_concrete()).to_concrete(), + ) + } + + pub fn min_scalar(&self, rhs: f32) -> Self { + self.min_elementwise(rhs) + } + + pub fn mt(&self, rhs: f32) -> Self { + self.gt_scalar(rhs) + } + + pub fn mte(&self, rhs: f32) -> Self { + self.gte_scalar(rhs) + } + + pub fn ne(&self, rhs: f32) -> Self { + self.unary_from_value(self.value.ne(rhs).to_concrete(), move |_, out| { + RawTensor::zeros(&out.device(), out.shape()) + }) + } + + pub fn ne_scalar(&self, rhs: f32) -> Self { + self.ne(rhs) + } + + pub fn ne_tensor(&self, rhs: &Self) -> Self { + assert_same_graph(self, rhs); + self.binary_op( + rhs, + self.value.ne_tensor(&rhs.value).to_concrete(), + move |_, lhs, rhs| { + vec![ + RawTensor::zeros(&lhs.device(), lhs.shape()), + RawTensor::zeros(&rhs.device(), rhs.shape()), + ] + }, + ) + } + + pub fn sigmoid(&self) -> Self { + self.mul_scalar(-1.0).exp().add_scalar(1.0).pow_scalar(-1.0) + } + + pub fn silu(&self) -> Self { + let denom = self.mul_scalar(-1.0).exp().add_scalar(1.0); + self.div(&denom) + } + + pub fn gelu(&self) -> Self { + let cubic = self.sqr().mul(self); + let inner = self + .add(&cubic.mul_scalar(0.044_715)) + .mul_scalar((2.0 / std::f32::consts::PI).sqrt()); + let gate = inner.tanh().add_scalar(1.0); + self.mul(&gate).mul_scalar(0.5) + } + + pub fn tanh(&self) -> Self { + self.unary_from_value(self.value.tanh().to_concrete(), move |grad, out| { + let one_minus_sq = (RawTensor::splat(&out.device(), 1.0, out.shape()) + - out.sqr().to_concrete()) + .to_concrete(); + (grad * one_minus_sq).to_concrete() + }) + } + + pub fn exp(&self) -> Self { + self.unary_from_value(self.value.exp().to_concrete(), move |grad, out| { + (grad * out).to_concrete() + }) + } + + pub fn where_cond(&self, on_true: &Self, on_false: &Self) -> Self { + assert_same_graph(self, on_true); + assert_same_graph(self, on_false); + + let value = self + .value + .where_cond(&on_true.value, &on_false.value) + .to_concrete(); + let condition_id = self.handle.id; + let true_id = on_true.handle.id; + let false_id = on_false.handle.id; + let condition = self.value.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "where_cond")?; + let zeros = RawTensor::zeros(&condition.device(), condition.shape()); + let ones = RawTensor::splat(&condition.device(), 1.0, condition.shape()); + let true_mask = condition.where_cond(&ones, &zeros).to_concrete(); + let false_mask = condition.where_cond(&zeros, &ones).to_concrete(); + Ok(vec![ + BackwardTarget { + node: condition_id, + gradient: Box::new(zeros), + }, + BackwardTarget { + node: true_id, + gradient: Box::new((gradient.clone() * true_mask).to_concrete()), + }, + BackwardTarget { + node: false_id, + gradient: Box::new((gradient * false_mask).to_concrete()), + }, + ]) + }); + self.emit_op( + value, + vec![ + self.handle.clone(), + on_true.handle.clone(), + on_false.handle.clone(), + ], + Some(backward), + ) + } + + pub fn log(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.log().to_concrete(), move |grad, _| { + (grad / input.clone()).to_concrete() + }) + } + + pub fn sqrt(&self) -> Self { + self.unary_from_value(self.value.sqrt().to_concrete(), move |grad, out| { + let denom = out.mul_scalar(2.0).to_concrete(); + (grad / denom).to_concrete() + }) + } +} + +macro_rules! impl_autograd_pairwise_op { + ($trait:ident, $method:ident) => { + impl std::ops::$trait> for Tensor { + type Output = Tensor; + + fn $method(self, rhs: Tensor) -> Tensor { + Tensor::$method(&self, &rhs) + } + } + + impl std::ops::$trait<&Tensor> for Tensor { + type Output = Tensor; + + fn $method(self, rhs: &Tensor) -> Tensor { + Tensor::$method(&self, rhs) + } + } + + impl std::ops::$trait> for &Tensor { + type Output = Tensor; + + fn $method(self, rhs: Tensor) -> Tensor { + Tensor::$method(self, &rhs) + } + } + + impl std::ops::$trait<&Tensor> for &Tensor { + type Output = Tensor; + + fn $method(self, rhs: &Tensor) -> Tensor { + Tensor::$method(self, rhs) + } + } + }; +} + +impl_autograd_pairwise_op!(Add, add); +impl_autograd_pairwise_op!(Sub, sub); +impl_autograd_pairwise_op!(Mul, mul); +impl_autograd_pairwise_op!(Div, div); + +macro_rules! impl_autograd_scalar_op { + ($trait:ident, $method:ident, $scalar_method:ident) => { + impl std::ops::$trait for Tensor { + type Output = Tensor; + + fn $method(self, rhs: f32) -> Tensor { + Tensor::$scalar_method(&self, rhs) + } + } + + impl std::ops::$trait for &Tensor { + type Output = Tensor; + + fn $method(self, rhs: f32) -> Tensor { + Tensor::$scalar_method(self, rhs) + } + } + }; +} + +impl_autograd_scalar_op!(Mul, mul, mul_scalar); +impl_autograd_scalar_op!(Add, add, add_scalar); +impl_autograd_scalar_op!(Sub, sub, sub_scalar); +impl_autograd_scalar_op!(Div, div, div_scalar); + +impl std::ops::Neg for Tensor { + type Output = Tensor; + + fn neg(self) -> Tensor { + Tensor::neg(&self) + } +} + +impl std::ops::Neg for &Tensor { + type Output = Tensor; + + fn neg(self) -> Tensor { + Tensor::neg(self) + } +} + diff --git a/fusor-ml/fusor/src/autograd/indexing.rs b/fusor-ml/fusor/src/autograd/indexing.rs new file mode 100644 index 000000000..e3b131dba --- /dev/null +++ b/fusor-ml/fusor/src/autograd/indexing.rs @@ -0,0 +1,147 @@ +use std::ops::Range; + +use crate::composite::index::IndexOp; + +use super::*; + +impl Tensor { + pub fn index_select(&self, dimension: usize, indices: &RawTensor<1, u32>) -> Self { + let input_shape = self.shape(); + assert!(dimension < R, "index_select dimension out of bounds"); + + let value = self.value.index_select(dimension, indices).to_concrete(); + let input_id = self.handle.id; + let indices = indices.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "index_select")?; + let one_hot = one_hot_matrix(&indices, input_shape[dimension]); + // transpose+reshape only commute through a copy, so the moved axis + // is materialized once on each side of the matmul; dimension 0 + // needs neither. + let moved = if dimension == 0 { + gradient + } else { + gradient.transpose(0, dimension).to_concrete() + }; + let moved_shape = moved.shape(); + let rest = moved_shape[1..].iter().product::(); + let flat = moved.reshape([moved_shape[0], rest]); + let scattered = one_hot.transpose(0, 1).mat_mul(&flat); + let mut unmoved_shape = moved_shape; + unmoved_shape[0] = input_shape[dimension]; + let scattered = scattered.reshape(unmoved_shape); + let input_gradient = if dimension == 0 { + scattered.to_concrete() + } else { + scattered.to_concrete().transpose(0, dimension).to_concrete() + }; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(input_gradient), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + fn index_ops(&self, ops: [IndexOp; R]) -> Tensor + where + crate::gpu::Tensor: crate::gpu::SmallerRank<1, OUT, f32>, + { + let shape = self.shape(); + let slices: [Range; R] = + std::array::from_fn(|axis| ops[axis].to_range(shape[axis])); + let dim = crate::composite::index::removed_dim(ops.map(|op| op.removes_dim())); + self.slice(slices).squeeze_dims::<1, OUT>([dim]) + } +} + +impl Tensor<2> { + pub fn i(&self, index: (I1, I2)) -> Tensor<1> + where + I1: Into, + I2: Into, + { + self.index_ops([index.0.into(), index.1.into()]) + } + + pub fn gather_last(&self, indices: &RawTensor<1, u32>) -> Tensor<1> { + let shape = self.shape(); + assert_eq!( + shape[0], + indices.shape()[0], + "gather_last expects one index per row" + ); + let width = shape[1]; + let device = self.device(); + let row_offsets = (0..shape[0]) + .map(|row| (row * width) as u32) + .collect::>(); + let row_offsets: RawTensor<1, u32> = + RawTensor::from_slice(&device, [shape[0]], &row_offsets); + let linear_indices = (row_offsets + indices.clone()).to_concrete(); + let flat = self.value.reshape([shape[0] * width]).to_concrete(); + let value = flat.index_select(0, &linear_indices).to_concrete(); + let input_id = self.handle.id; + let indices = indices.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::<1>(&*gradient, "gather_last")?; + let one_hot = one_hot_matrix(&indices, width); + let input_gradient: RawTensor<2, f32> = one_hot.mul_(&gradient.reshape([shape[0], 1])); + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(input_gradient), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn embedding(&self, indices: &RawTensor<2, u32>) -> Tensor<3> { + let [rows, columns] = indices.shape(); + let width = self.shape()[1]; + let flat_indices = indices.clone().reshape([rows * columns]).to_concrete(); + self.index_select(0, &flat_indices) + .reshape([rows, columns, width]) + } +} + +impl Tensor<3> { + pub fn i(&self, index: (I1, I2, I3)) -> Tensor<2> + where + I1: Into, + I2: Into, + I3: Into, + { + self.index_ops([index.0.into(), index.1.into(), index.2.into()]) + } +} + +impl Tensor<4> { + pub fn i(&self, index: (I1, I2, I3, I4)) -> Tensor<3> + where + I1: Into, + I2: Into, + I3: Into, + I4: Into, + { + self.index_ops([ + index.0.into(), + index.1.into(), + index.2.into(), + index.3.into(), + ]) + } +} + +/// Build a `[indices.len(), size]` f32 matrix with 1.0 at `[row, indices[row]]` +/// so scatter-adds stay on-device as matmuls/products against it; duplicate +/// indices accumulate through the contraction. +fn one_hot_matrix(indices: &RawTensor<1, u32>, size: usize) -> RawTensor<2, f32> { + let device = indices.device(); + let rows = indices.shape()[0]; + let positions = (0..size) + .map(|position| position as f32) + .collect::>(); + let positions: RawTensor<2, f32> = RawTensor::from_slice(&device, [1, size], &positions); + let indices = indices.cast::().reshape([rows, 1]).to_concrete(); + indices.sub_(&positions).eq(0.0) +} diff --git a/fusor-ml/fusor/src/autograd/layers/conv.rs b/fusor-ml/fusor/src/autograd/layers/conv.rs new file mode 100644 index 000000000..21bc388c1 --- /dev/null +++ b/fusor-ml/fusor/src/autograd/layers/conv.rs @@ -0,0 +1,124 @@ +//! Trainable N-dimensional convolution layer. + +use super::super::{Graph, Tensor}; + +pub use crate::layers::ConvNdConfig; + +/// N-dimensional convolution layer with trainable parameters. +/// +/// Input / output tensors have rank `RANK = SPATIAL + 2`: +/// `(batch, channels, ...spatial)`. +/// Weight has shape `(out_channels, in_channels / groups, ...kernel)`. +pub struct ConvNd { + weight: Tensor, + bias: Option>, + config: ConvNdConfig, + in_channels: usize, + out_channels: usize, +} + +impl ConvNd { + /// Create a new convolution layer. + /// + /// `weight` shape: `(out_channels, in_channels / groups, ...kernel)`. + /// `bias` shape: `(out_channels,)`. + pub fn new( + weight: Tensor, + bias: Option>, + config: ConvNdConfig, + ) -> Self { + const { + assert!(RANK == SPATIAL + 2); + } + let shape = weight.shape(); + let out_channels = shape[0]; + let in_channels = shape[1] * config.groups; + + assert!(config.groups >= 1, "groups must be >= 1"); + assert_eq!( + out_channels % config.groups, + 0, + "out_channels ({out_channels}) must be divisible by groups ({})", + config.groups + ); + + if let Some(ref b) = bias { + assert_eq!( + b.shape()[0], + out_channels, + "Bias shape must match out_channels" + ); + } + + Self { + weight, + bias, + config, + in_channels, + out_channels, + } + } + + /// Import an inference layer's parameters as trainable leaves on `graph`. + pub fn from_inference( + graph: &Graph, + layer: &crate::layers::ConvNd, + ) -> Self { + Self::new( + graph.leaf(layer.weight().clone()), + layer.bias().map(|bias| graph.leaf(bias.clone())), + *layer.config(), + ) + } + + /// Get the configuration. + pub fn config(&self) -> &ConvNdConfig { + &self.config + } + + /// Number of input channels. + pub fn in_channels(&self) -> usize { + self.in_channels + } + + /// Number of output channels. + pub fn out_channels(&self) -> usize { + self.out_channels + } + + /// The weight leaf: `(out_channels, in_channels / groups, ...kernel)`. + pub fn weight(&self) -> &Tensor { + &self.weight + } + + /// The bias leaf: `(out_channels,)`. + pub fn bias(&self) -> Option<&Tensor<1>> { + self.bias.as_ref() + } + + /// Forward pass for any spatial rank. The free const generic `R2` equals + /// `RANK + SPATIAL` and is determined by the `LargerRank` bound, exactly + /// the same way the underlying `conv` operation infers it. + pub fn forward(&self, input: &Tensor) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank, + { + if self.config.groups == 1 { + input.conv( + &self.weight, + self.bias.as_ref(), + self.config.padding, + self.config.stride, + ) + } else { + input.grouped_conv( + &self.weight, + self.bias.as_ref(), + self.config.padding, + self.config.stride, + self.config.groups, + ) + } + } +} diff --git a/fusor-ml/fusor/src/autograd/layers/conv1d.rs b/fusor-ml/fusor/src/autograd/layers/conv1d.rs new file mode 100644 index 000000000..db486793b --- /dev/null +++ b/fusor-ml/fusor/src/autograd/layers/conv1d.rs @@ -0,0 +1,106 @@ +//! Trainable Conv1d layer implementation. + +use super::super::{Graph, Tensor}; + +pub use crate::layers::Conv1dConfig; + +/// 1D Convolution layer with trainable parameters. +/// +/// Applies a 1D convolution over an input signal. +/// Input shape: (batch, in_channels, length) +/// Output shape: (batch, out_channels, out_length) +/// where out_length = (length + 2*padding - kernel_size) / stride + 1 +pub struct Conv1d { + weight: Tensor<3>, + bias: Option>, + config: Conv1dConfig, + in_channels: usize, + out_channels: usize, + kernel_size: usize, +} + +impl Conv1d { + /// Create a new Conv1d layer with given weights and configuration. + /// + /// Weight shape: (out_channels, in_channels, kernel_size) + /// Bias shape: (out_channels,) + pub fn new(weight: Tensor<3>, bias: Option>, config: Conv1dConfig) -> Self { + let shape = weight.shape(); + let out_channels = shape[0]; + let in_channels = shape[1]; + let kernel_size = shape[2]; + + assert_eq!(config.groups, 1, "Only groups=1 is currently supported"); + assert_eq!(config.dilation, 1, "Only dilation=1 is currently supported"); + + if let Some(ref b) = bias { + assert_eq!( + b.shape()[0], + out_channels, + "Bias shape must match out_channels" + ); + } + + Self { + weight, + bias, + config, + in_channels, + out_channels, + kernel_size, + } + } + + /// Import an inference layer's parameters as trainable graph leaves. + pub fn from_inference(graph: &Graph, layer: &crate::layers::Conv1d) -> Self { + Self::new( + graph.leaf(layer.weight().clone()), + layer.bias().map(|bias| graph.leaf(bias.clone())), + *layer.config(), + ) + } + + /// Forward pass. + /// + /// Input shape: (batch, in_channels, length) + /// Output shape: (batch, out_channels, out_length) + pub fn forward(&self, input: &Tensor<3>) -> Tensor<3> { + input.conv( + &self.weight, + self.bias.as_ref(), + [self.config.padding], + [self.config.stride], + ) + } + + /// Get the weight tensor. + pub fn weight(&self) -> &Tensor<3> { + &self.weight + } + + /// Get the bias tensor. + pub fn bias(&self) -> Option<&Tensor<1>> { + self.bias.as_ref() + } + + /// Get the configuration. + pub fn config(&self) -> &Conv1dConfig { + &self.config + } + + /// Get the number of input channels. + pub fn in_channels(&self) -> usize { + self.in_channels + } + + /// Get the number of output channels. + pub fn out_channels(&self) -> usize { + self.out_channels + } + + /// Get the kernel size. + pub fn kernel_size(&self) -> usize { + self.kernel_size + } +} + diff --git a/fusor-ml/fusor/src/autograd/layers/embedding.rs b/fusor-ml/fusor/src/autograd/layers/embedding.rs new file mode 100644 index 000000000..19310ba4d --- /dev/null +++ b/fusor-ml/fusor/src/autograd/layers/embedding.rs @@ -0,0 +1,70 @@ +//! Trainable embedding layer implementation. + +use super::super::{Graph, RawTensor, Tensor}; + +/// Embedding layer for token/position embeddings. +/// +/// Maps integer indices to dense vectors. +/// Embedding table shape: (num_embeddings, embedding_dim) +#[derive(Clone)] +pub struct Embedding { + embeddings: Tensor<2>, + num_embeddings: usize, + embedding_dim: usize, +} + +impl Embedding { + /// Create a new embedding layer with the given embedding table. + pub fn new_from_tensor(embeddings: Tensor<2>) -> Self { + let shape = embeddings.shape(); + let num_embeddings = shape[0]; + let embedding_dim = shape[1]; + + Self { + embeddings, + num_embeddings, + embedding_dim, + } + } + + /// Import an inference embedding layer as a trainable layer whose + /// embedding table is a gradient leaf on `graph`. + pub fn from_inference(graph: &Graph, layer: &crate::layers::Embedding) -> Self { + let table = match layer.dense_embeddings() { + Some(dense) => dense.clone(), + None => layer.embeddings_quantized().dequantize().to_concrete(), + }; + Self::new_from_tensor(graph.leaf(table)) + } + + /// Forward pass: lookup embeddings for the given indices. + /// + /// Input: [batch, seq_len] with indices + /// Output: [batch, seq_len, embedding_dim] with embeddings + pub fn forward(&self, indices: &RawTensor<2, u32>) -> Tensor<3> { + self.embeddings.embedding(indices) + } + + /// Forward pass: lookup embeddings for flat indices. + /// + /// Input: [seq_len] with indices + /// Output: [seq_len, embedding_dim] with embeddings + pub fn forward_1d(&self, indices: &RawTensor<1, u32>) -> Tensor<2> { + self.embeddings.index_select(0, indices) + } + + /// Get the embedding table. + pub fn embeddings(&self) -> &Tensor<2> { + &self.embeddings + } + + /// Get the number of embeddings. + pub fn num_embeddings(&self) -> usize { + self.num_embeddings + } + + /// Get the embedding dimension. + pub fn embedding_dim(&self) -> usize { + self.embedding_dim + } +} diff --git a/fusor-ml/fusor/src/autograd/layers/layer_norm.rs b/fusor-ml/fusor/src/autograd/layers/layer_norm.rs new file mode 100644 index 000000000..95a1a5c01 --- /dev/null +++ b/fusor-ml/fusor/src/autograd/layers/layer_norm.rs @@ -0,0 +1,195 @@ +//! Trainable layer normalization. + +use super::super::{Graph, Tensor}; + +/// Layer Normalization. +/// +/// Normalizes the input over the last dimension. +/// Formula: output = (input - mean) / sqrt(variance + eps) * weight + bias +pub struct LayerNorm { + weight: Tensor, + bias: Option>, + eps: f32, +} + +impl LayerNorm { + /// Create a new LayerNorm layer. + /// + /// Weight and bias should have shape (normalized_dim,). + pub fn new(weight: Tensor, bias: Option>, eps: f32) -> Self { + Self { weight, bias, eps } + } + + /// Import an inference [`crate::layers::LayerNorm`]'s weights as trainable leaves. + pub fn from_inference(graph: &Graph, layer: &crate::layers::LayerNorm) -> Self { + Self::new( + graph.leaf(layer.weight().clone()), + layer.bias().map(|bias| graph.leaf(bias.clone())), + layer.eps(), + ) + } + + /// Get the weight tensor. + pub fn weight(&self) -> &Tensor { + &self.weight + } + + /// Get the bias tensor if present. + pub fn bias(&self) -> Option<&Tensor> { + self.bias.as_ref() + } + + /// Get the epsilon value. + pub fn eps(&self) -> f32 { + self.eps + } +} + +impl LayerNorm<1> { + /// Forward pass for 2D input (batch, features). + /// + /// Normalizes over the last dimension (features). + pub fn forward_2d(&self, input: &Tensor<2>) -> Tensor<2> { + let weight_broadcast = self.weight.broadcast_as(input.shape()); + let bias_broadcast = self + .bias + .as_ref() + .map(|bias| bias.broadcast_as(input.shape())); + input.layer_norm::<1>(&weight_broadcast, bias_broadcast.as_ref(), self.eps, true) + } + + /// Forward pass for 3D input (batch, seq_len, features). + /// + /// Normalizes over the last dimension (features). + pub fn forward(&self, input: &Tensor<3>) -> Tensor<3> { + let weight_broadcast = self.weight.broadcast_as(input.shape()); + let bias_broadcast = self + .bias + .as_ref() + .map(|bias| bias.broadcast_as(input.shape())); + input.layer_norm::<2>(&weight_broadcast, bias_broadcast.as_ref(), self.eps, true) + } + + /// Forward pass through the fused last-dim kernel (3D input). + pub fn forward_fused(&self, input: &Tensor<3>) -> Tensor<3> { + input.layer_norm_last_dim_fused::<2, 1>(&self.weight, self.bias.as_ref(), self.eps) + } +} + +/// Layer normalization with a selectable reduction axis. +/// +/// `axis == None` normalizes the last dimension. `axis == Some(a)` normalizes +/// dimension `a` by transposing that axis to the end, applying the common +/// last-dimension path, then transposing back. +pub struct LayerNormNd { + weight: Tensor<1>, + bias: Option>, + axis: Option, + eps: f32, +} + +impl LayerNormNd { + /// Create a LayerNorm that normalizes the last dimension. + pub fn new(weight: Tensor<1>, bias: Option>, eps: f32) -> Self { + Self { + weight, + bias, + axis: None, + eps, + } + } + + /// Create a LayerNorm that normalizes the given axis. + pub fn new_over_axis( + weight: Tensor<1>, + bias: Option>, + axis: usize, + eps: f32, + ) -> Self { + Self { + weight, + bias, + axis: Some(axis), + eps, + } + } + + /// Import an inference [`crate::layers::LayerNormNd`]'s weights as trainable + /// leaves, normalizing the last dimension. + pub fn from_inference(graph: &Graph, layer: &crate::layers::LayerNormNd) -> Self { + Self::new( + graph.leaf(layer.weight().clone()), + layer.bias().map(|bias| graph.leaf(bias.clone())), + layer.eps(), + ) + } + + /// Import an inference [`crate::layers::LayerNormNd`]'s weights as trainable + /// leaves, normalizing `axis`. + pub fn from_inference_over_axis( + graph: &Graph, + layer: &crate::layers::LayerNormNd, + axis: usize, + ) -> Self { + Self::new_over_axis( + graph.leaf(layer.weight().clone()), + layer.bias().map(|bias| graph.leaf(bias.clone())), + axis, + layer.eps(), + ) + } + + pub fn weight(&self) -> &Tensor<1> { + &self.weight + } + + pub fn bias(&self) -> Option<&Tensor<1>> { + self.bias.as_ref() + } + + pub fn eps(&self) -> f32 { + self.eps + } + + /// Forward pass for 2D input. + pub fn forward_2d(&self, input: &Tensor<2>) -> Tensor<2> { + self.forward::<2, 1>(input) + } + + /// Forward pass for any input rank. `OUT_RANK` equals `N - 1`. + pub fn forward(&self, input: &Tensor) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + let shape = input.shape(); + let axis = self.axis.unwrap_or(N - 1); + + if axis == N - 1 { + let weight_b = self.weight.broadcast_as(shape); + let bias_b = self.bias.as_ref().map(|bias| bias.broadcast_as(shape)); + return input.layer_norm::(&weight_b, bias_b.as_ref(), self.eps, true); + } + + let mut permuted_shape = shape; + permuted_shape.swap(axis, N - 1); + let permuted = input.transpose(axis, N - 1); + let weight_b = self.weight.broadcast_as(permuted_shape); + let bias_b = self + .bias + .as_ref() + .map(|bias| bias.broadcast_as(permuted_shape)); + let normed = permuted.layer_norm::(&weight_b, bias_b.as_ref(), self.eps, true); + normed.transpose(axis, N - 1) + } + + /// Fused fast path for normalizing the last dim of a rank-3 tensor. + pub fn forward_fused(&self, input: &Tensor<3>) -> Tensor<3> { + if matches!(self.axis, Some(axis) if axis != 2) { + return self.forward::<3, 2>(input); + } + + input.layer_norm_last_dim_fused::<2, 1>(&self.weight, self.bias.as_ref(), self.eps) + } +} diff --git a/fusor-ml/fusor/src/autograd/layers/linear.rs b/fusor-ml/fusor/src/autograd/layers/linear.rs new file mode 100644 index 000000000..e3f7e9b1d --- /dev/null +++ b/fusor-ml/fusor/src/autograd/layers/linear.rs @@ -0,0 +1,72 @@ +//! Trainable linear layer implementation. + +use super::super::{Graph, Tensor}; + +/// A trainable linear (fully connected) layer. +/// +/// Computes `output = input @ weight.T + bias`. +pub struct Linear { + weight: Tensor<2>, + bias: Option>, +} + +impl Linear { + /// Create a new Linear layer with the given weight and optional bias. + /// + /// Weight shape: (out_features, in_features) + /// Bias shape: (out_features,) + pub fn new(weight: Tensor<2>, bias: Option>) -> Self { + Self { weight, bias } + } + + /// Import an inference layer as trainable graph leaves, dequantizing the weight to f32. + pub fn from_inference(graph: &Graph, layer: &crate::layers::Linear) -> Self { + let weight = graph.leaf(layer.weight().dequantize::<2>().to_concrete()); + let bias = layer.bias().map(|bias| graph.leaf(bias.clone())); + Self { weight, bias } + } + + /// Get the weight tensor. + pub fn weight(&self) -> &Tensor<2> { + &self.weight + } + + /// Get the bias tensor if present. + pub fn bias(&self) -> Option<&Tensor<1>> { + self.bias.as_ref() + } + + /// Get the input features size. + pub fn in_features(&self) -> usize { + self.weight.shape()[1] + } + + /// Get the output features size. + pub fn out_features(&self) -> usize { + self.weight.shape()[0] + } + + /// Forward pass for 3D input (batch, seq_len, in_features) + /// + /// Input shape: (batch, seq_len, in_features) + /// Output shape: (batch, seq_len, out_features) + pub fn forward(&self, input: &Tensor<3>) -> Tensor<3> { + let [batch, seq_len, in_features] = input.shape(); + let input_2d = input.reshape([batch * seq_len, in_features]); + self.forward_2d(&input_2d) + .reshape([batch, seq_len, self.out_features()]) + } + + /// Forward pass for 2D input (batch, in_features) + /// + /// Input shape: (batch, in_features) + /// Output shape: (batch, out_features) + pub fn forward_2d(&self, input: &Tensor<2>) -> Tensor<2> { + let output = input.mat_mul(&self.weight.t()); + if let Some(bias) = &self.bias { + output.add_(bias) + } else { + output + } + } +} diff --git a/fusor-ml/fusor/src/autograd/layers/mod.rs b/fusor-ml/fusor/src/autograd/layers/mod.rs new file mode 100644 index 000000000..74676e288 --- /dev/null +++ b/fusor-ml/fusor/src/autograd/layers/mod.rs @@ -0,0 +1,13 @@ +mod conv; +mod conv1d; +mod embedding; +mod layer_norm; +mod linear; +mod rms_norm; + +pub use conv::{ConvNd, ConvNdConfig}; +pub use conv1d::{Conv1d, Conv1dConfig}; +pub use embedding::Embedding; +pub use layer_norm::{LayerNorm, LayerNormNd}; +pub use linear::Linear; +pub use rms_norm::RmsNorm; diff --git a/fusor-ml/fusor/src/autograd/layers/rms_norm.rs b/fusor-ml/fusor/src/autograd/layers/rms_norm.rs new file mode 100644 index 000000000..b70f65214 --- /dev/null +++ b/fusor-ml/fusor/src/autograd/layers/rms_norm.rs @@ -0,0 +1,69 @@ +//! Trainable RMS normalization implementation. + +use super::super::{Graph, Tensor}; + +/// Root Mean Square Normalization. +/// +/// Normalizes the input over the last dimension without centering. +/// Formula: output = input / sqrt(mean(x^2) + eps) * weight +pub struct RmsNorm { + weight: Tensor, + bias: Option>, + eps: f32, +} + +impl RmsNorm { + /// Create a new RmsNorm layer. + /// + /// Weight should have shape matching the normalized dimension. + pub fn new(weight: Tensor, bias: Option>, eps: f32) -> Self { + Self { weight, bias, eps } + } + + /// Import a raw inference layer's weights as trainable graph leaves. + pub fn from_inference(graph: &Graph, layer: &crate::layers::RmsNorm) -> Self { + Self { + weight: graph.leaf(layer.weight().clone()), + bias: layer.bias().map(|bias| graph.leaf(bias.clone())), + eps: layer.eps(), + } + } + + /// Get the weight tensor. + pub fn weight(&self) -> &Tensor { + &self.weight + } + + /// Get the bias tensor if present. + pub fn bias(&self) -> Option<&Tensor> { + self.bias.as_ref() + } + + /// Get the epsilon value. + pub fn eps(&self) -> f32 { + self.eps + } +} + +impl RmsNorm<1> { + /// Forward pass for 2D input (batch, features). + pub fn forward_2d(&self, input: &Tensor<2>) -> Tensor<2> { + input.rms_norm_fused::<1, 1>(&self.weight, self.bias.as_ref(), self.eps) + } + + /// Forward pass for 3D input (batch, seq_len, features). + pub fn forward(&self, input: &Tensor<3>) -> Tensor<3> { + input.rms_norm_fused::<1, 2>(&self.weight, self.bias.as_ref(), self.eps) + } + + /// Forward pass for 4D input (batch, heads, seq_len, features). + pub fn forward_4d(&self, input: &Tensor<4>) -> Tensor<4> { + input.rms_norm_fused::<1, 3>(&self.weight, self.bias.as_ref(), self.eps) + } + + /// Forward pass for `input + residual` followed by RMSNorm. + pub fn forward_residual(&self, input: &Tensor<3>, residual: &Tensor<3>) -> Tensor<3> { + input.rms_norm_residual_fused::<1, 2>(residual, &self.weight, self.bias.as_ref(), self.eps) + } +} + diff --git a/fusor-ml/fusor/src/autograd/matmul.rs b/fusor-ml/fusor/src/autograd/matmul.rs new file mode 100644 index 000000000..9bf0c9caa --- /dev/null +++ b/fusor-ml/fusor/src/autograd/matmul.rs @@ -0,0 +1,80 @@ +use super::*; + +impl Tensor { + pub(super) fn mat_mul_internal(&self, rhs: &Self) -> Self { + assert_same_graph(self, rhs); + let value = self.value.mat_mul(&rhs.value); + let lhs_id = self.handle.id; + let rhs_id = rhs.handle.id; + let lhs_value = self.value.clone(); + let rhs_value = rhs.value.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "mat_mul")?; + Ok(vec![ + BackwardTarget { + node: lhs_id, + gradient: Box::new( + gradient.clone().mat_mul(&rhs_value.transpose(R - 2, R - 1)), + ), + }, + BackwardTarget { + node: rhs_id, + gradient: Box::new(lhs_value.transpose(R - 2, R - 1).mat_mul(&gradient)), + }, + ]) + }); + self.emit_op( + value, + vec![self.handle.clone(), rhs.handle.clone()], + Some(backward), + ) + } + + pub fn mat_mul(&self, rhs: &Self) -> Self { + self.mat_mul_internal(rhs) + } + + pub fn matmul(&self, rhs: &Self) -> Self { + self.mat_mul_internal(rhs) + } + + pub fn t(&self) -> Self { + assert!(R >= 2, "t requires rank >= 2"); + self.transpose(R - 2, R - 1) + } + + pub fn q_mat_mul(&self, weights: &crate::QMatrix) -> Self { + if R == 1 { + let k = self.shape()[0]; + let n = weights.shape()[0]; + let out_shape: [usize; R] = std::array::from_fn(|_| n); + return self.reshape([1, k]).q_mat_mul(weights).reshape(out_shape); + } + let value = self.value.q_mat_mul(weights).to_concrete(); + if !self.requires_grad() { + return self.emit_op(value, vec![self.handle.clone()], None); + } + let weights = weights.clone(); + self.replay_unary("q_mat_mul", value, move |input| { + let n = weights.shape()[0]; + let k = weights.shape()[1]; + let batch_dims = R - 2; + let weight_shape: [usize; R] = std::array::from_fn(|i| { + if i < batch_dims { + 1 + } else if i == batch_dims { + k + } else { + n + } + }); + let dequantized: RawTensor<2, f32> = weights.dequantize(); + let weight = dequantized + .transpose(0, 1) + .reshape(weight_shape) + .to_concrete(); + let weight = Tensor::constant_from_raw(&input.graph(), weight); + input.mat_mul_internal(&weight) + }) + } +} diff --git a/fusor-ml/fusor/src/autograd/mod.rs b/fusor-ml/fusor/src/autograd/mod.rs new file mode 100644 index 000000000..9e35fc50c --- /dev/null +++ b/fusor-ml/fusor/src/autograd/mod.rs @@ -0,0 +1,823 @@ +use std::{ + any::Any, + collections::{HashMap, HashSet, VecDeque}, + sync::{Arc, Mutex}, +}; + +use crate::{Device, Error, Result, Tensor as RawTensor}; + +mod composite; +mod elementwise; +mod indexing; +pub mod layers; +mod matmul; +mod reduce; +mod view; + +#[cfg(test)] +mod tests; + +type NodeId = usize; +#[cfg(not(target_arch = "wasm32"))] +type BackwardRule = + Arc) -> Result> + Send + Sync>; +#[cfg(target_arch = "wasm32")] +type BackwardRule = Arc) -> Result>>; + +#[cfg(not(target_arch = "wasm32"))] +trait BackwardClosure: Send + Sync + 'static {} +#[cfg(not(target_arch = "wasm32"))] +impl BackwardClosure for T where T: Send + Sync + 'static {} + +#[cfg(target_arch = "wasm32")] +trait BackwardClosure: 'static {} +#[cfg(target_arch = "wasm32")] +impl BackwardClosure for T where T: 'static {} + +#[derive(Clone)] +pub struct Graph { + inner: Arc, +} + +#[derive(Clone)] +pub struct Tensor { + value: RawTensor, + handle: NodeHandle, +} + +pub struct Gradients { + gradients: HashMap>, +} + +pub struct BackwardTarget { + node: NodeId, + gradient: Box, +} + +#[derive(Clone)] +pub struct Parent { + handle: NodeHandle, +} + +#[derive(Clone)] +struct NodeHandle { + graph: Arc, + id: NodeId, +} + +#[derive(Clone)] +struct Node { + parents: Vec, + backward: Option, + requires_grad: bool, +} + +struct GraphInner { + state: Mutex, +} + +struct GraphState { + next_id: NodeId, + nodes: HashMap, +} + +#[cfg(not(target_arch = "wasm32"))] +trait AnyTensorValue: Send + Sync { + fn as_any(&self) -> &dyn Any; + fn clone_box(&self) -> Box; + fn into_detached(self: Box) -> Box; + fn add_box(&self, other: &dyn AnyTensorValue) -> Result>; +} + +#[cfg(target_arch = "wasm32")] +trait AnyTensorValue { + fn as_any(&self) -> &dyn Any; + fn clone_box(&self) -> Box; + fn into_detached(self: Box) -> Box; + fn add_box(&self, other: &dyn AnyTensorValue) -> Result>; +} + +impl Graph { + pub fn new() -> Self { + Self { + inner: Arc::new(GraphInner { + state: Mutex::new(GraphState { + next_id: 0, + nodes: HashMap::new(), + }), + }), + } + } + + pub fn leaf(&self, value: RawTensor) -> Tensor { + self.tensor_with_grad(value, true) + } + + pub fn constant(&self, value: RawTensor) -> Tensor { + self.tensor_with_grad(value, false) + } + + pub fn tensor(&self, device: &Device, data: T) -> Tensor + where + RawTensor: fusor_types::FromArray, + { + self.leaf(RawTensor::new(device, data)) + } + + pub fn constant_from_data(&self, device: &Device, data: T) -> Tensor + where + RawTensor: fusor_types::FromArray, + { + self.constant(RawTensor::new(device, data)) + } + + fn tensor_with_grad( + &self, + value: RawTensor, + requires_grad: bool, + ) -> Tensor { + let id = self.inner.add_node(Vec::new(), None, requires_grad); + Tensor { + value, + handle: NodeHandle { + graph: self.inner.clone(), + id, + }, + } + } +} + +impl Default for Graph { + fn default() -> Self { + Self::new() + } +} + +impl Tensor { + pub fn from_raw(graph: &Graph, value: RawTensor) -> Self { + graph.leaf(value) + } + + pub fn constant_from_raw(graph: &Graph, value: RawTensor) -> Self { + graph.constant(value) + } + + pub fn new(graph: &Graph, device: &Device, data: T) -> Self + where + RawTensor: fusor_types::FromArray, + { + graph.tensor(device, data) + } + + pub fn from_array(graph: &Graph, device: &Device, data: T) -> Self + where + RawTensor: fusor_types::FromArray, + { + Self::new(graph, device, data) + } + + pub fn from_slice(graph: &Graph, device: &Device, shape: [usize; R], data: &[f32]) -> Self { + graph.leaf(RawTensor::from_slice(device, shape, data)) + } + + pub fn zeros(graph: &Graph, device: &Device, shape: [usize; R]) -> Self { + graph.leaf(RawTensor::zeros(device, shape)) + } + + pub fn ones(graph: &Graph, device: &Device, shape: [usize; R]) -> Self { + Self::splat(graph, device, 1.0, shape) + } + + pub fn splat(graph: &Graph, device: &Device, value: f32, shape: [usize; R]) -> Self { + graph.leaf(RawTensor::splat(device, value, shape)) + } + + pub fn full(graph: &Graph, device: &Device, shape: [usize; R], value: f32) -> Self { + Self::splat(graph, device, value, shape) + } + + pub fn zeros_like(&self) -> Self { + Self::zeros(&self.graph(), &self.device(), self.shape()) + } + + pub fn ones_like(&self) -> Self { + Self::ones(&self.graph(), &self.device(), self.shape()) + } + + pub fn raw(&self) -> &RawTensor { + &self.value + } + + pub fn into_raw(self) -> RawTensor { + self.value + } + + pub fn shape(&self) -> [usize; R] { + self.value.shape() + } + + pub fn device(&self) -> Device { + self.value.device() + } + + pub fn graph(&self) -> Graph { + Graph { + inner: self.handle.graph.clone(), + } + } + + pub fn requires_grad(&self) -> bool { + self.handle.graph.requires_grad(self.handle.id) + } + + pub fn parent(&self) -> Parent { + Parent { + handle: self.handle.clone(), + } + } + + pub fn detach(&self) -> Self { + let requires_grad = self.requires_grad(); + let id = self.handle.graph.add_node(Vec::new(), None, requires_grad); + Self { + value: self.value.to_concrete(), + handle: NodeHandle { + graph: self.handle.graph.clone(), + id, + }, + } + } + + #[cfg(not(target_arch = "wasm32"))] + pub fn with_backwards(self, parents: I, backwards: F) -> Self + where + I: IntoIterator, + F: Fn(RawTensor) -> Result> + Send + Sync + 'static, + { + self.with_backwards_impl(parents, backwards) + } + + #[cfg(target_arch = "wasm32")] + pub fn with_backwards(self, parents: I, backwards: F) -> Self + where + I: IntoIterator, + F: Fn(RawTensor) -> Result> + 'static, + { + self.with_backwards_impl(parents, backwards) + } + + fn with_backwards_impl(self, parents: I, backwards: F) -> Self + where + I: IntoIterator, + F: Fn(RawTensor) -> Result> + BackwardClosure, + { + let parent_handles = parents + .into_iter() + .map(|parent| parent.handle) + .collect::>(); + let requires_grad = parent_handles + .iter() + .any(|parent| parent.graph.requires_grad(parent.id)); + let parent_ids = parent_handles + .iter() + .map(|parent| parent.id) + .collect::>(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = gradient + .as_any() + .downcast_ref::>() + .ok_or_else(|| Error::msg("gradient rank mismatch in custom backward"))? + .clone(); + let targets = backwards(gradient)?; + // The scheduler only unlocks a parent once every child sends it a + // gradient, so a missing target would silently starve that + // parent's whole subgraph. + for parent in &parent_handles { + if parent.graph.requires_grad(parent.id) + && !targets.iter().any(|target| target.node == parent.id) + { + return Err(Error::msg( + "custom backward omitted a gradient for a parent that requires grad", + )); + } + } + Ok(targets) + }); + self.handle.graph.replace_node( + self.handle.id, + Node { + parents: parent_ids, + backward: Some(backward), + requires_grad, + }, + ); + self + } + + pub fn backward(&self) -> Result { + let elements = self.shape().iter().product::(); + if elements != 1 { + return Err(Error::msg( + "backward() requires a single-element tensor; use backward_with() for non-scalars", + )); + } + let seed = RawTensor::splat(&self.device(), 1.0, self.shape()); + self.backward_with(seed) + } + + pub fn backward_with(&self, seed: RawTensor) -> Result { + self.handle.graph.backward(self.handle.id, Box::new(seed)) + } + + fn emit_op( + &self, + value: RawTensor, + parents: Vec, + backward: Option, + ) -> Tensor { + for parent in &parents { + assert!( + Arc::ptr_eq(&self.handle.graph, &parent.graph), + "cannot mix autograd tensors from different graphs" + ); + } + let requires_grad = parents + .iter() + .any(|parent| parent.graph.requires_grad(parent.id)); + let parent_ids = parents.into_iter().map(|parent| parent.id).collect(); + let id = self + .handle + .graph + .add_node(parent_ids, backward, requires_grad); + Tensor { + value, + handle: NodeHandle { + graph: self.handle.graph.clone(), + id, + }, + } + } + + fn unary_from_value( + &self, + value: RawTensor, + backward: impl Fn(RawTensor, RawTensor) -> RawTensor + BackwardClosure, + ) -> Self { + let input_id = self.handle.id; + let output = value.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "unary")?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(backward(gradient, output.clone()).to_concrete()), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + fn binary_op( + &self, + rhs: &Self, + value: RawTensor, + backward: impl Fn( + RawTensor, + RawTensor, + RawTensor, + ) -> Vec> + + BackwardClosure, + ) -> Self { + assert!( + Arc::ptr_eq(&self.handle.graph, &rhs.handle.graph), + "cannot mix autograd tensors from different graphs" + ); + let lhs_id = self.handle.id; + let rhs_id = rhs.handle.id; + let lhs_value = self.value.clone(); + let rhs_value = rhs.value.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "binary")?; + let gradients = backward(gradient, lhs_value.clone(), rhs_value.clone()); + Ok(vec![ + BackwardTarget { + node: lhs_id, + gradient: Box::new(gradients[0].clone().to_concrete()), + }, + BackwardTarget { + node: rhs_id, + gradient: Box::new(gradients[1].clone().to_concrete()), + }, + ]) + }); + self.emit_op( + value, + vec![self.handle.clone(), rhs.handle.clone()], + Some(backward), + ) + } + + fn replay_unary( + &self, + context: &'static str, + value: RawTensor, + replay: impl Fn(Tensor) -> Tensor + BackwardClosure, + ) -> Tensor { + let input_id = self.handle.id; + let input_value = self.value.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, context)?; + let graph = Graph::new(); + let replay_input = Tensor::from_raw(&graph, input_value.clone()); + let replay_output = replay(replay_input.clone()); + let gradients = replay_output.backward_with(gradient)?; + let input_gradient = gradients + .get(&replay_input) + .ok_or_else(|| Error::msg(format!("missing replay gradient in {context}")))?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(input_gradient), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + fn replay_binary( + &self, + rhs: &Tensor, + context: &'static str, + value: RawTensor, + replay: impl Fn(Tensor, Tensor) -> Tensor + BackwardClosure, + ) -> Tensor { + assert_same_graph(self, rhs); + let lhs_id = self.handle.id; + let rhs_id = rhs.handle.id; + let lhs_value = self.value.clone(); + let rhs_value = rhs.value.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, context)?; + let graph = Graph::new(); + let replay_lhs = Tensor::from_raw(&graph, lhs_value.clone()); + let replay_rhs = Tensor::from_raw(&graph, rhs_value.clone()); + let replay_output = replay(replay_lhs.clone(), replay_rhs.clone()); + let gradients = replay_output.backward_with(gradient)?; + let lhs_gradient = gradients + .get(&replay_lhs) + .ok_or_else(|| Error::msg(format!("missing lhs replay gradient in {context}")))?; + let rhs_gradient = gradients + .get(&replay_rhs) + .ok_or_else(|| Error::msg(format!("missing rhs replay gradient in {context}")))?; + Ok(vec![ + BackwardTarget { + node: lhs_id, + gradient: Box::new(lhs_gradient), + }, + BackwardTarget { + node: rhs_id, + gradient: Box::new(rhs_gradient), + }, + ]) + }); + self.emit_op( + value, + vec![self.handle.clone(), rhs.handle.clone()], + Some(backward), + ) + } + + fn replay_quaternary( + &self, + second: &Tensor, + third: &Tensor, + fourth: &Tensor, + context: &'static str, + value: RawTensor, + replay: impl Fn(Tensor, Tensor, Tensor, Tensor) -> Tensor + BackwardClosure, + ) -> Tensor { + assert_same_graph(self, second); + assert_same_graph(self, third); + assert_same_graph(self, fourth); + let ids = [ + self.handle.id, + second.handle.id, + third.handle.id, + fourth.handle.id, + ]; + let first_value = self.value.clone(); + let second_value = second.value.clone(); + let third_value = third.value.clone(); + let fourth_value = fourth.value.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, context)?; + let graph = Graph::new(); + let replay_first = Tensor::from_raw(&graph, first_value.clone()); + let replay_second = Tensor::from_raw(&graph, second_value.clone()); + let replay_third = Tensor::from_raw(&graph, third_value.clone()); + let replay_fourth = Tensor::from_raw(&graph, fourth_value.clone()); + let replay_output = replay( + replay_first.clone(), + replay_second.clone(), + replay_third.clone(), + replay_fourth.clone(), + ); + let gradients = replay_output.backward_with(gradient)?; + let missing = + || Error::msg(format!("missing replay gradient in {context}")); + Ok(vec![ + BackwardTarget { + node: ids[0], + gradient: Box::new(gradients.get(&replay_first).ok_or_else(missing)?), + }, + BackwardTarget { + node: ids[1], + gradient: Box::new(gradients.get(&replay_second).ok_or_else(missing)?), + }, + BackwardTarget { + node: ids[2], + gradient: Box::new(gradients.get(&replay_third).ok_or_else(missing)?), + }, + BackwardTarget { + node: ids[3], + gradient: Box::new(gradients.get(&replay_fourth).ok_or_else(missing)?), + }, + ]) + }); + self.emit_op( + value, + vec![ + self.handle.clone(), + second.handle.clone(), + third.handle.clone(), + fourth.handle.clone(), + ], + Some(backward), + ) + } + + fn replay_ternary( + &self, + second: &Tensor, + third: &Tensor, + context: &'static str, + value: RawTensor, + replay: impl Fn(Tensor, Tensor, Tensor) -> Tensor + BackwardClosure, + ) -> Tensor { + assert_same_graph(self, second); + assert_same_graph(self, third); + let first_id = self.handle.id; + let second_id = second.handle.id; + let third_id = third.handle.id; + let first_value = self.value.clone(); + let second_value = second.value.clone(); + let third_value = third.value.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, context)?; + let graph = Graph::new(); + let replay_first = Tensor::from_raw(&graph, first_value.clone()); + let replay_second = Tensor::from_raw(&graph, second_value.clone()); + let replay_third = Tensor::from_raw(&graph, third_value.clone()); + let replay_output = replay( + replay_first.clone(), + replay_second.clone(), + replay_third.clone(), + ); + let gradients = replay_output.backward_with(gradient)?; + let first_gradient = gradients + .get(&replay_first) + .ok_or_else(|| Error::msg(format!("missing first replay gradient in {context}")))?; + let second_gradient = gradients.get(&replay_second).ok_or_else(|| { + Error::msg(format!("missing second replay gradient in {context}")) + })?; + let third_gradient = gradients + .get(&replay_third) + .ok_or_else(|| Error::msg(format!("missing third replay gradient in {context}")))?; + Ok(vec![ + BackwardTarget { + node: first_id, + gradient: Box::new(first_gradient), + }, + BackwardTarget { + node: second_id, + gradient: Box::new(second_gradient), + }, + BackwardTarget { + node: third_id, + gradient: Box::new(third_gradient), + }, + ]) + }); + self.emit_op( + value, + vec![ + self.handle.clone(), + second.handle.clone(), + third.handle.clone(), + ], + Some(backward), + ) + } +} + +impl Tensor<1> { + pub fn arange(graph: &Graph, device: &Device, start: f32, end: f32) -> Tensor<1> { + graph.leaf(crate::arange(device, start, end)) + } + + pub fn arange_step( + graph: &Graph, + device: &Device, + start: f32, + end: f32, + step: f32, + ) -> Tensor<1> { + graph.leaf(crate::arange_step(device, start, end, step)) + } +} + +impl Gradients { + pub fn get(&self, tensor: &Tensor) -> Option> { + self.gradients + .get(&tensor.handle.id) + .and_then(|gradient| gradient.as_any().downcast_ref::>()) + .cloned() + } + + pub fn into_detached(self) -> Self { + Self { + gradients: self + .gradients + .into_iter() + .map(|(id, gradient)| (id, gradient.into_detached())) + .collect(), + } + } +} + +impl BackwardTarget { + pub fn wrt(tensor: &Tensor, gradient: RawTensor) -> Self { + Self { + node: tensor.handle.id, + gradient: Box::new(gradient), + } + } +} + +impl GraphInner { + fn add_node( + &self, + parents: Vec, + backward: Option, + requires_grad: bool, + ) -> NodeId { + let mut state = self.state.lock().unwrap(); + let id = state.next_id; + state.next_id += 1; + state.nodes.insert( + id, + Node { + parents, + backward, + requires_grad, + }, + ); + id + } + + fn replace_node(&self, id: NodeId, node: Node) { + self.state.lock().unwrap().nodes.insert(id, node); + } + + fn requires_grad(&self, id: NodeId) -> bool { + self.state + .lock() + .unwrap() + .nodes + .get(&id) + .map(|node| node.requires_grad) + .unwrap_or(false) + } + + fn backward(&self, root: NodeId, seed: Box) -> Result { + let nodes = self.reachable_nodes(root); + let mut pending_children = HashMap::::new(); + for (id, node) in &nodes { + pending_children.entry(*id).or_insert(0); + for parent in &node.parents { + *pending_children.entry(*parent).or_insert(0) += 1; + } + } + + let mut gradients = HashMap::>::new(); + gradients.insert(root, seed); + + let mut queue = VecDeque::new(); + queue.push_back(root); + + while let Some(node_id) = queue.pop_front() { + let Some(node) = nodes.get(&node_id) else { + continue; + }; + let Some(backward) = node.backward.as_ref() else { + continue; + }; + let gradient = gradients + .get(&node_id) + .ok_or_else(|| Error::msg(format!("missing gradient for node {node_id}")))? + .clone_box(); + + for target in backward(gradient)? { + let Some(parent_node) = nodes.get(&target.node) else { + continue; + }; + if !parent_node.requires_grad { + continue; + } + accumulate_gradient(&mut gradients, target.node, target.gradient)?; + let remaining = pending_children.get_mut(&target.node).ok_or_else(|| { + Error::msg(format!("missing child count for node {}", target.node)) + })?; + *remaining = remaining.saturating_sub(1); + if *remaining == 0 { + queue.push_back(target.node); + } + } + } + + Ok(Gradients { gradients }) + } + + fn reachable_nodes(&self, root: NodeId) -> HashMap { + let snapshot = self.state.lock().unwrap().nodes.clone(); + let mut reachable = HashMap::new(); + let mut stack = vec![root]; + let mut visited = HashSet::new(); + while let Some(node_id) = stack.pop() { + if !visited.insert(node_id) { + continue; + } + if let Some(node) = snapshot.get(&node_id) { + reachable.insert(node_id, node.clone()); + stack.extend(node.parents.iter().copied()); + } + } + reachable + } +} + +impl AnyTensorValue for RawTensor { + fn as_any(&self) -> &dyn Any { + self + } + + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } + + fn into_detached(self: Box) -> Box { + match *self { + RawTensor::Cpu(tensor) => Box::new(RawTensor::Cpu(tensor.to_concrete())), + RawTensor::Gpu(tensor) => Box::new(RawTensor::Gpu(tensor.detach())), + } + } + + fn add_box(&self, other: &dyn AnyTensorValue) -> Result> { + let other = other + .as_any() + .downcast_ref::>() + .ok_or_else(|| Error::msg("gradient rank mismatch while accumulating"))?; + Ok(Box::new((self.clone() + other.clone()).to_concrete())) + } +} + +fn accumulate_gradient( + gradients: &mut HashMap>, + node: NodeId, + gradient: Box, +) -> Result<()> { + match gradients.get(&node) { + Some(existing) => { + let accumulated = existing.add_box(&*gradient)?; + gradients.insert(node, accumulated); + } + None => { + gradients.insert(node, gradient); + } + } + Ok(()) +} + +fn downcast_tensor( + value: &dyn AnyTensorValue, + context: &str, +) -> Result> { + value + .as_any() + .downcast_ref::>() + .cloned() + .ok_or_else(|| Error::msg(format!("gradient rank mismatch in {context}"))) +} + +fn assert_same_graph(lhs: &Tensor, rhs: &Tensor) { + assert!( + Arc::ptr_eq(&lhs.handle.graph, &rhs.handle.graph), + "cannot mix autograd tensors from different graphs" + ); +} diff --git a/fusor-ml/fusor/src/autograd/reduce.rs b/fusor-ml/fusor/src/autograd/reduce.rs new file mode 100644 index 000000000..399160940 --- /dev/null +++ b/fusor-ml/fusor/src/autograd/reduce.rs @@ -0,0 +1,419 @@ +use fusor_types::{SlidingWindow, StrideSpec}; + +use super::*; + +impl Tensor { + pub(super) fn sum_keepdim_any(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + let input_shape = self.shape(); + let value = self.value.sum_keepdim::(axis).to_concrete(); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "sum_keepdim")?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(gradient.broadcast_as(input_shape).to_concrete()), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub(super) fn sum_any(&self, axis: usize) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + let input_shape = self.shape(); + let value = self.value.sum::(axis).to_concrete(); + let input_id = self.handle.id; + let mut keepdim_shape = input_shape; + keepdim_shape[axis] = 1; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "sum")?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new( + gradient + .reshape(keepdim_shape) + .to_concrete() + .broadcast_as(input_shape) + .to_concrete(), + ), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub(super) fn max_keepdim_any(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + let input = self.value.clone(); + let value = input.max_keepdim::(axis).to_concrete(); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "max_keepdim")?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(reduction_extrema_keepdim_grad::( + input.clone(), + axis, + gradient, + true, + )), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + fn min_keepdim_any(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + let input = self.value.clone(); + let value = input.min_keepdim::(axis).to_concrete(); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "min_keepdim")?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(reduction_extrema_keepdim_grad::( + input.clone(), + axis, + gradient, + false, + )), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub(super) fn mean_keepdim_any(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.sum_keepdim_any::(axis) + .div_scalar(self.shape()[axis] as f32) + } + + fn product_keepdim_any(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::ProdOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + { + let input = self.value.clone(); + let input_shape = self.shape(); + let value = input.product_keepdim::(axis).to_concrete(); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "product_keepdim")?; + let upstream = gradient.broadcast_as(input_shape).to_concrete(); + let zeros = RawTensor::zeros(&input.device(), input_shape); + let ones = RawTensor::splat(&input.device(), 1.0, input_shape); + let zero_mask = input.eq(0.0).to_concrete(); + let safe_input = zero_mask.where_cond(&ones, &input).to_concrete(); + let zero_count = zero_mask.sum_keepdim::(axis).to_concrete(); + let zero_count_broadcast = zero_count.broadcast_as(input_shape).to_concrete(); + let product_non_zero = safe_input + .product_keepdim::(axis) + .broadcast_as(input_shape) + .to_concrete(); + let no_zero_grad = (upstream.clone() + * (product_non_zero.clone() / safe_input).to_concrete()) + .to_concrete(); + let single_zero_grad = zero_mask + .where_cond(&(upstream * product_non_zero).to_concrete(), &zeros) + .to_concrete(); + let gradient = ((no_zero_grad * zero_count_broadcast.eq(0.0).to_concrete()) + .to_concrete() + + (single_zero_grad * zero_count_broadcast.eq(1.0).to_concrete()).to_concrete()) + .to_concrete(); + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(gradient), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + fn var_keepdim_any(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + let mean = self.mean_keepdim_any::(axis); + let centered = self.sub(&mean.broadcast_as(self.shape())); + centered.sqr().mean_keepdim_any::(axis) + } + + pub fn pool( + &self, + pools: [impl Into; DIFF], + with: impl Fn(&Tensor, usize) -> Self + Copy, + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank, + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank<1, R3, f32>, + crate::gpu::Tensor: crate::gpu::SmallerRank, + { + let pools: [crate::composite::pool::PoolSize; DIFF] = pools.map(|pool| pool.into()); + let axis_start = R - DIFF; + let windows: [SlidingWindow; DIFF] = std::array::from_fn(|i| { + let pool = pools[i]; + SlidingWindow::new(axis_start + i, pool.size, pool.stride) + }); + let shape = self.shape(); + let mut sorted_windows = windows; + sorted_windows.sort_by_key(|window| window.axis); + let specs: [StrideSpec; R2] = std::array::from_fn(|out_i| { + if out_i < R { + if let Some(window) = sorted_windows.iter().find(|window| window.axis == out_i) { + let positions = (shape[out_i] - window.window_size) / window.step + 1; + StrideSpec::dim_with(out_i, positions, window.step) + } else { + StrideSpec::dim(out_i, shape[out_i]) + } + } else { + let window = &sorted_windows[out_i - R]; + StrideSpec::dim(window.axis, window.window_size) + } + }); + + let tiled: Tensor = self.restride(specs); + let unsqueezed: Tensor = tiled.unsqueeze_dims::<1, R3>([R2]); + let flattened: Tensor = unsqueezed.flatten_last_n::(); + with(&flattened, O - 1) + } + + pub fn pool_max( + &self, + pools: [impl Into; DIFF], + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank, + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank<1, R3, f32>, + crate::gpu::Tensor: crate::gpu::SmallerRank, + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: + crate::gpu::LastRank + crate::gpu::SmallerRank<1, R, f32>, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.pool::(pools, |windowed, axis| windowed.max::(axis)) + } + + pub fn pool_min( + &self, + pools: [impl Into; DIFF], + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank, + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank<1, R3, f32>, + crate::gpu::Tensor: crate::gpu::SmallerRank, + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: + crate::gpu::LastRank + crate::gpu::SmallerRank<1, R, f32>, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.pool::(pools, |windowed, axis| windowed.min::(axis)) + } + + pub fn max_keepdim(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.max_keepdim_any::(axis) + } + + pub fn max(&self, axis: usize) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: + crate::gpu::LastRank + crate::gpu::SmallerRank<1, OUT_RANK, f32>, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.max_keepdim_any::(axis) + .squeeze_dims::<1, OUT_RANK>([axis]) + } + + pub fn min_keepdim(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.min_keepdim_any::(axis) + } + + pub fn min(&self, axis: usize) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: + crate::gpu::LastRank + crate::gpu::SmallerRank<1, OUT_RANK, f32>, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.min_keepdim_any::(axis) + .squeeze_dims::<1, OUT_RANK>([axis]) + } + + pub fn mean_keepdim(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.mean_keepdim_any::(axis) + } + + pub fn mean(&self, axis: usize) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: + crate::gpu::LastRank + crate::gpu::SmallerRank<1, OUT_RANK, f32>, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.mean_keepdim_any::(axis) + .squeeze_dims::<1, OUT_RANK>([axis]) + } + + pub fn product(&self, axis: usize) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: + crate::gpu::LastRank + crate::gpu::SmallerRank<1, OUT_RANK, f32>, + crate::cpu::ProdOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + { + self.product_keepdim_any::(axis) + .squeeze_dims::<1, OUT_RANK>([axis]) + } + + pub fn product_keepdim(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::ProdOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + { + self.product_keepdim_any::(axis) + } + + pub fn var(&self, axis: usize) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: + crate::gpu::LastRank + crate::gpu::SmallerRank<1, OUT_RANK, f32>, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.var_keepdim_any::(axis) + .squeeze_dims::<1, OUT_RANK>([axis]) + } + + pub fn var_keepdim(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.var_keepdim_any::(axis) + } +} + +impl Tensor<1> { + pub fn sum(&self) -> Tensor<0> { + let input_shape = self.shape(); + let value = self.value.sum::<0>(0); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::<0>(&*gradient, "sum")?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(gradient.broadcast_as(input_shape).to_concrete()), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn sum_keepdim(&self, axis: usize) -> Tensor<1> { + self.sum_keepdim_any::<0>(axis) + } +} + +macro_rules! sum_wrappers { + ($($rank:literal => $out:literal),* $(,)?) => {$( + impl Tensor<$rank> { + pub fn sum(&self, axis: usize) -> Tensor<$out> { + self.sum_any::<$out>(axis) + } + + pub fn sum_keepdim(&self, axis: usize) -> Tensor<$rank> { + self.sum_keepdim_any::<$out>(axis) + } + } + )*}; +} + +sum_wrappers!(2 => 1, 3 => 2, 4 => 3, 5 => 4, 6 => 5, 7 => 6, 8 => 7, 9 => 8, 10 => 9); + +fn reduction_extrema_keepdim_grad( + input: RawTensor, + axis: usize, + gradient: RawTensor, + is_max: bool, +) -> RawTensor +where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, +{ + let input_shape = input.shape(); + let extrema = if is_max { + input.max_keepdim::(axis) + } else { + input.min_keepdim::(axis) + } + .to_concrete(); + let extrema_broadcast = extrema.broadcast_as(input_shape).to_concrete(); + let mask = (input - extrema_broadcast) + .to_concrete() + .eq(0.0) + .to_concrete(); + let tie_count = mask + .sum_keepdim::(axis) + .broadcast_as(input_shape) + .to_concrete(); + ((mask * gradient.broadcast_as(input_shape)).to_concrete() / tie_count).to_concrete() +} diff --git a/fusor-ml/fusor/src/autograd/tests.rs b/fusor-ml/fusor/src/autograd/tests.rs new file mode 100644 index 000000000..6716897c2 --- /dev/null +++ b/fusor-ml/fusor/src/autograd/tests.rs @@ -0,0 +1,6713 @@ +use super::*; +use crate::{Layout, ToVec1, ToVec2}; +use fusor_types::StrideSpec; + +fn assert_close(left: f32, right: f32) { + assert!((left - right).abs() < 1e-3, "expected {right}, got {left}"); +} + +fn assert_slice_close(left: &[f32], right: &[f32]) { + assert_eq!(left.len(), right.len(), "slice lengths differ"); + for (index, (left, right)) in left.iter().zip(right.iter()).enumerate() { + assert!( + (*left - *right).abs() < 1e-3, + "mismatch at index {index}: expected {right}, got {left}", + ); + } +} + +async fn test_devices() -> Vec { + let mut devices = vec![Device::cpu()]; + match Device::gpu().await { + Ok(gpu) => devices.push(gpu), + Err(_) => eprintln!("skipping GPU coverage: GPU unavailable"), + } + devices +} + +async fn flatten(tensor: RawTensor) -> Vec { + let elements = tensor.shape().into_iter().product(); + tensor + .reshape([elements]) + .as_slice() + .await + .unwrap() + .to_vec1() +} + +async fn finite_difference_gradient( + device: &Device, + shape: [usize; R], + data: &[f32], + loss: &F, +) -> Vec +where + F: Fn(&Graph, Tensor) -> Tensor<0>, +{ + let epsilon = 1e-2f32; + let mut numeric = Vec::with_capacity(data.len()); + for index in 0..data.len() { + let mut perturbed = data.to_vec(); + perturbed[index] = data[index] + epsilon; + let graph = Graph::new(); + let plus = loss( + &graph, + Tensor::from_slice(&graph, device, shape, &perturbed), + ); + let plus = plus.raw().to_scalar().await.unwrap(); + perturbed[index] = data[index] - epsilon; + let graph = Graph::new(); + let minus = loss( + &graph, + Tensor::from_slice(&graph, device, shape, &perturbed), + ); + let minus = minus.raw().to_scalar().await.unwrap(); + numeric.push((plus - minus) / (2.0 * epsilon)); + } + numeric +} + +async fn assert_gradient_matches_finite_difference( + device: &Device, + shape: [usize; R], + data: &[f32], + loss: F, +) where + F: Fn(&Graph, Tensor) -> Tensor<0>, +{ + let graph = Graph::new(); + let input = Tensor::from_slice(&graph, device, shape, data); + let gradients = loss(&graph, input.clone()).backward().unwrap(); + let analytic = flatten(gradients.get(&input).unwrap()).await; + let numeric = finite_difference_gradient(device, shape, data, &loss).await; + assert_eq!(analytic.len(), numeric.len(), "gradient lengths differ"); + for (index, (analytic, numeric)) in analytic.iter().zip(numeric.iter()).enumerate() { + assert!( + (analytic - numeric).abs() < 1e-2 + 1e-2 * numeric.abs(), + "gradient mismatch at index {index}: analytic {analytic}, finite difference {numeric}", + ); + } +} + +#[tokio::test] +async fn test_backward_squared_sum() { + for device in test_devices().await { + let graph = Graph::new(); + + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + let loss = x.sqr().sum(); + let gradients = loss.backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(dx[0], 2.0); + assert_close(dx[1], 4.0); + assert_close(dx[2], 6.0); + } +} + +#[tokio::test] +async fn test_autograd_silu() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, -2.0, 0.5]); + + let output = x.silu(); + let values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + + let expected = [1.0f32, -2.0, 0.5].map(|v| v / (1.0 + (-v).exp())); + for (value, expected) in values.iter().zip(expected) { + assert_close(*value, expected); + } + + let gradients = output.sum().backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + let expected_grads = [1.0f32, -2.0, 0.5].map(|v| { + let sigmoid = 1.0 / (1.0 + (-v).exp()); + sigmoid * (1.0 + v * (1.0 - sigmoid)) + }); + for (value, expected) in dx.iter().zip(expected_grads) { + assert_close(*value, expected); + } + + assert_gradient_matches_finite_difference(&device, [3], &[1.0, -2.0, 0.5], |_, x| { + x.silu().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_gelu() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, -2.0, 0.5]); + + let output = x.gelu(); + let values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + + let expected = [1.0f32, -2.0, 0.5].map(|v| { + 0.5 * v + * (1.0 + ((2.0 / std::f32::consts::PI).sqrt() * (v + 0.044_715 * v.powi(3))).tanh()) + }); + for (value, expected) in values.iter().zip(expected) { + assert_close(*value, expected); + } + + let gradients = output.sum().backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + let expected_grads = [1.0f32, -2.0, 0.5].map(|v| { + let scale = (2.0 / std::f32::consts::PI).sqrt(); + let inner = scale * (v + 0.044_715 * v.powi(3)); + let tanh = inner.tanh(); + let dinner = scale * (1.0 + 3.0 * 0.044_715 * v * v); + 0.5 * (1.0 + tanh) + 0.5 * v * (1.0 - tanh * tanh) * dinner + }); + for (value, expected) in dx.iter().zip(expected_grads) { + assert_close(*value, expected); + } + + assert_gradient_matches_finite_difference(&device, [3], &[1.0, -2.0, 0.5], |_, x| { + x.gelu().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_where_cond() { + for device in test_devices().await { + let graph = Graph::new(); + let condition: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 0.0, -2.0]); + let on_true: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32, 3.0, 4.0]); + let on_false: Tensor<1> = Tensor::new(&graph, &device, &[10.0f32, 20.0, 30.0]); + + let output = condition.where_cond(&on_true, &on_false); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.flatten_all().sum().backward().unwrap(); + + let dcondition = gradients + .get(&condition) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + let dtrue = gradients + .get(&on_true) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + let dfalse = gradients + .get(&on_false) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![2.0, 20.0, 4.0]); + assert_eq!(dcondition, vec![0.0, 0.0, 0.0]); + assert_eq!(dtrue, vec![1.0, 0.0, 1.0]); + assert_eq!(dfalse, vec![0.0, 1.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_index_select() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]); + let indices = RawTensor::from_slice(&device, [3], &[2u32, 0, 2]); + + let output = input.index_select(1, &indices); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_eq!( + output_values, + vec![vec![3.0, 1.0, 3.0], vec![6.0, 4.0, 6.0]] + ); + assert_eq!(dinput, vec![vec![1.0, 0.0, 2.0], vec![1.0, 0.0, 2.0]]); + } +} + +#[tokio::test] +async fn test_backward_slice_assign() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new( + &graph, + &device, + &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], + ); + let value: Tensor<2> = Tensor::new(&graph, &device, &[[10.0f32, 11.0], [12.0, 13.0]]); + + let output = input.slice_assign([0..2, 1..3], &value); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + let gradients = output.sum(1).sum().backward().unwrap(); + + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + let dvalue = gradients + .get(&value) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_eq!( + output_values, + vec![ + vec![1.0, 10.0, 11.0], + vec![4.0, 12.0, 13.0], + vec![7.0, 8.0, 9.0] + ] + ); + assert_eq!( + dinput, + vec![ + vec![1.0, 0.0, 0.0], + vec![1.0, 0.0, 0.0], + vec![1.0, 1.0, 1.0] + ] + ); + assert_eq!(dvalue, vec![vec![1.0, 1.0], vec![1.0, 1.0]]); + } +} + +#[tokio::test] +async fn test_backward_expand() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[2.0f32, 3.0, 4.0]]); + + let output = input.expand([2, 3]); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_eq!( + output_values, + vec![vec![2.0, 3.0, 4.0], vec![2.0, 3.0, 4.0]] + ); + assert_eq!(dinput, vec![vec![2.0, 2.0, 2.0]]); + } +} + +#[tokio::test] +async fn test_backward_flatten_all() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0], [3.0, 4.0]]); + + let output = input.flatten_all(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_eq!(output_values, vec![1.0, 2.0, 3.0, 4.0]); + assert_eq!(dinput, vec![vec![1.0, 1.0], vec![1.0, 1.0]]); + } +} + +#[tokio::test] +async fn test_backward_flatten_last_n() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<3> = Tensor::new( + &graph, + &device, + &[ + [[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]], + [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]], + ], + ); + + let output = input.flatten_last_n::<1, 2>(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .reshape([2, 6]) + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_eq!( + output_values, + vec![ + vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], + vec![7.0, 8.0, 9.0, 10.0, 11.0, 12.0] + ] + ); + assert_eq!( + dinput, + vec![ + vec![1.0, 1.0, 1.0, 1.0, 1.0, 1.0], + vec![1.0, 1.0, 1.0, 1.0, 1.0, 1.0] + ] + ); + } +} + +#[tokio::test] +async fn test_backward_flatten_first_n() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<3> = Tensor::new( + &graph, + &device, + &[ + [[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]], + [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]], + ], + ); + + let output = input.flatten_first_n::<1, 2>(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .reshape([4, 3]) + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_eq!( + output_values, + vec![ + vec![1.0, 2.0, 3.0], + vec![4.0, 5.0, 6.0], + vec![7.0, 8.0, 9.0], + vec![10.0, 11.0, 12.0] + ] + ); + assert_eq!( + dinput, + vec![ + vec![1.0, 1.0, 1.0], + vec![1.0, 1.0, 1.0], + vec![1.0, 1.0, 1.0], + vec![1.0, 1.0, 1.0] + ] + ); + } +} + +#[tokio::test] +async fn test_backward_narrow() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new( + &graph, + &device, + &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], + ); + + let output = input.narrow(1usize, 1, 2); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_eq!( + output_values, + vec![vec![2.0, 3.0], vec![5.0, 6.0], vec![8.0, 9.0]] + ); + assert_eq!( + dinput, + vec![ + vec![0.0, 1.0, 1.0], + vec![0.0, 1.0, 1.0], + vec![0.0, 1.0, 1.0] + ] + ); + } +} + +#[tokio::test] +async fn test_backward_repeat() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0], [3.0, 4.0]]); + + let output = input.repeat([2, 3]); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_eq!( + output_values, + vec![ + vec![1.0, 2.0, 1.0, 2.0, 1.0, 2.0], + vec![3.0, 4.0, 3.0, 4.0, 3.0, 4.0], + vec![1.0, 2.0, 1.0, 2.0, 1.0, 2.0], + vec![3.0, 4.0, 3.0, 4.0, 3.0, 4.0] + ] + ); + assert_eq!(dinput, vec![vec![6.0, 6.0], vec![6.0, 6.0]]); + } +} + +#[tokio::test] +async fn test_backward_resize() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new( + &graph, + &device, + &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], + ); + + let output = input.resize([2, 2]); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_eq!(output_values, vec![vec![1.0, 2.0], vec![4.0, 5.0]]); + assert_eq!( + dinput, + vec![ + vec![1.0, 1.0, 0.0], + vec![1.0, 1.0, 0.0], + vec![0.0, 0.0, 0.0] + ] + ); + } +} + +#[tokio::test] +async fn test_backward_restride() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0, 4.0]); + + let output = input.restride([StrideSpec::dim(0, 2), StrideSpec::dim(0, 3)]); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!( + output_values, + vec![vec![1.0, 2.0, 3.0], vec![2.0, 3.0, 4.0]] + ); + assert_eq!(dinput, vec![1.0, 2.0, 2.0, 1.0]); + } +} + +#[tokio::test] +async fn test_backward_restride_strided_overlap() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new( + &graph, + &device, + &[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], + ); + + let output = input.restride([StrideSpec::dim_with(0, 3, 2), StrideSpec::dim(0, 3)]); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!( + output_values, + vec![ + vec![1.0, 2.0, 3.0], + vec![3.0, 4.0, 5.0], + vec![5.0, 6.0, 7.0] + ] + ); + assert_eq!(dinput, vec![1.0, 1.0, 2.0, 1.0, 2.0, 1.0, 1.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_restride_layout() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0, 4.0, 5.0]); + let layout = Layout::contiguous(&[5]) + .restride(&[StrideSpec::dim(0, 2).with_offset(1), StrideSpec::dim(0, 2)]); + + let output = input.restride_layout(layout); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![vec![2.0, 3.0], vec![3.0, 4.0]]); + assert_eq!(dinput, vec![0.0, 1.0, 2.0, 1.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_squeeze_dims() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<4> = Tensor::new( + &graph, + &device, + &[[[[1.0f32], [2.0], [3.0]]], [[[4.0], [5.0], [6.0]]]], + ); + + let output = input.squeeze_dims::<2, 2>([1, 3]); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .reshape([2, 3]) + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_eq!( + output_values, + vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]] + ); + assert_eq!(dinput, vec![vec![1.0, 1.0, 1.0], vec![1.0, 1.0, 1.0]]); + } +} + +#[tokio::test] +async fn test_backward_unsqueeze_dims() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]); + + let output = input.unsqueeze_dims::<2, 4>([0, 2]); + let output_values = output + .raw() + .clone() + .reshape([2, 3]) + .as_slice() + .await + .unwrap() + .to_vec2(); + let gradients = output.sum(3).sum(2).sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_eq!(output.shape(), [1, 2, 1, 3]); + assert_eq!( + output_values, + vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]] + ); + assert_eq!(dinput, vec![vec![1.0, 1.0, 1.0], vec![1.0, 1.0, 1.0]]); + } +} + +#[tokio::test] +async fn test_backward_max() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 5.0, 5.0], [4.0, 2.0, 0.0]]); + + let output = input.max::<1>(1); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_eq!(output_values, vec![5.0, 4.0]); + assert_eq!(dinput, vec![vec![0.0, 0.5, 0.5], vec![1.0, 0.0, 0.0]]); + } +} + +#[tokio::test] +async fn test_backward_min() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 1.0, 5.0], [4.0, 2.0, 0.0]]); + + let output = input.min::<1>(1); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_eq!(output_values, vec![1.0, 0.0]); + assert_eq!(dinput, vec![vec![0.5, 0.5, 0.0], vec![0.0, 0.0, 1.0]]); + } +} + +#[tokio::test] +async fn test_backward_mean() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]); + + let output = input.mean::<1>(1); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_eq!(output_values, vec![2.0, 5.0]); + assert_eq!( + dinput, + vec![ + vec![1.0 / 3.0, 1.0 / 3.0, 1.0 / 3.0], + vec![1.0 / 3.0, 1.0 / 3.0, 1.0 / 3.0] + ] + ); + } +} + +#[tokio::test] +async fn test_backward_product() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new( + &graph, + &device, + &[[2.0f32, 3.0, 4.0], [5.0, 0.0, 7.0], [0.0, 0.0, 9.0]], + ); + + let output = input.product::<1>(1); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_eq!(output_values, vec![24.0, 0.0, 0.0]); + assert_eq!( + dinput, + vec![ + vec![12.0, 8.0, 6.0], + vec![0.0, 35.0, 0.0], + vec![0.0, 0.0, 0.0] + ] + ); + } +} + +#[tokio::test] +async fn test_backward_product_keepdim() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[2.0f32, 3.0, 4.0], [5.0, 0.0, 7.0]]); + + let output = input.product_keepdim::<1>(1); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_eq!(output_values, vec![vec![24.0], vec![0.0]]); + assert_eq!(dinput, vec![vec![12.0, 8.0, 6.0], vec![0.0, 35.0, 0.0]]); + } +} + +#[tokio::test] +async fn test_backward_var() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]); + + let output = input.var::<1>(1); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_eq!(output_values, vec![2.0 / 3.0, 2.0 / 3.0]); + assert_eq!( + dinput, + vec![ + vec![-2.0 / 3.0, 0.0, 2.0 / 3.0], + vec![-2.0 / 3.0, 0.0, 2.0 / 3.0] + ] + ); + } +} + +#[tokio::test] +async fn test_backward_var_keepdim() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]); + + let output = input.var_keepdim::<1>(1); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_eq!(output_values, vec![vec![2.0 / 3.0], vec![2.0 / 3.0]]); + assert_eq!( + dinput, + vec![ + vec![-2.0 / 3.0, 0.0, 2.0 / 3.0], + vec![-2.0 / 3.0, 0.0, 2.0 / 3.0] + ] + ); + } +} + +#[tokio::test] +async fn test_backward_clamp() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[-1.0f32, 0.0, 2.0, 5.0]); + + let output = input.clamp(0.0, 3.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![0.0, 0.0, 2.0, 3.0]); + assert_eq!(dinput, vec![0.0, 0.0, 1.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_eq() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 1.0]); + + let output = input.eq(1.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![1.0, 0.0, 1.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_eq_scalar() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[3.0f32, 2.0, 3.0]); + + let output = input.eq_scalar(3.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![1.0, 0.0, 1.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_eq_tensor() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + let rhs: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 0.0, 3.0]); + + let output = lhs.eq_tensor(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![1.0, 0.0, 1.0]); + assert_eq!(dlhs, vec![0.0, 0.0, 0.0]); + assert_eq!(drhs, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_gt_scalar() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + + let output = input.gt_scalar(2.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![0.0, 0.0, 1.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_gt_tensor() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 4.0, 3.0]); + let rhs: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32, 1.0, 3.0]); + + let output = lhs.gt_tensor(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![0.0, 1.0, 0.0]); + assert_eq!(dlhs, vec![0.0, 0.0, 0.0]); + assert_eq!(drhs, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_gte_scalar() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + + let output = input.gte_scalar(2.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![0.0, 1.0, 1.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_gte_tensor() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 4.0, 3.0]); + let rhs: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32, 4.0, 2.0]); + + let output = lhs.gte_tensor(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![0.0, 1.0, 1.0]); + assert_eq!(dlhs, vec![0.0, 0.0, 0.0]); + assert_eq!(drhs, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_lt() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + + let output = input.lt(2.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![1.0, 0.0, 0.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_lt_scalar() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + + let output = input.lt_scalar(3.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![1.0, 1.0, 0.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_lt_tensor() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + let rhs: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32, 1.0, 3.0]); + + let output = lhs.lt_tensor(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![1.0, 0.0, 0.0]); + assert_eq!(dlhs, vec![0.0, 0.0, 0.0]); + assert_eq!(drhs, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_lte() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + + let output = input.lte(2.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![1.0, 1.0, 0.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_lte_scalar() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + + let output = input.lte_scalar(1.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![1.0, 0.0, 0.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_lte_tensor() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + let rhs: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32, 2.0, 1.0]); + + let output = lhs.lte_tensor(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![1.0, 1.0, 0.0]); + assert_eq!(dlhs, vec![0.0, 0.0, 0.0]); + assert_eq!(drhs, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_max_elementwise() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[-1.0f32, 0.0, 2.0]); + + let output = input.max_elementwise(0.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![0.0, 0.0, 2.0]); + assert_eq!(dinput, vec![0.0, 0.0, 1.0]); + } +} + +#[tokio::test] +async fn test_backward_max_scalar() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 4.0, 2.0]); + + let output = input.max_scalar(3.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![3.0, 4.0, 3.0]); + assert_eq!(dinput, vec![0.0, 1.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_min_elementwise() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 4.0, 2.0]); + + let output = input.min_elementwise(3.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![1.0, 3.0, 2.0]); + assert_eq!(dinput, vec![1.0, 0.0, 1.0]); + } +} + +#[tokio::test] +async fn test_backward_min_scalar() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 4.0, 2.0]); + + let output = input.min_scalar(2.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![1.0, 2.0, 2.0]); + assert_eq!(dinput, vec![1.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_mt() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 4.0, 2.0]); + + let output = input.mt(2.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![0.0, 1.0, 0.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_mte() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 4.0, 2.0]); + + let output = input.mte(2.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![0.0, 1.0, 1.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_ne() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 4.0, 2.0]); + + let output = input.ne(2.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![1.0, 1.0, 0.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_ne_scalar() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 4.0, 2.0]); + + let output = input.ne_scalar(4.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![1.0, 0.0, 1.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_ne_tensor() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 4.0, 2.0]); + let rhs: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 0.0, 3.0]); + + let output = lhs.ne_tensor(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![0.0, 1.0, 1.0]); + assert_eq!(dlhs, vec![0.0, 0.0, 0.0]); + assert_eq!(drhs, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_abs() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[-2.0f32, 0.0, 3.0]); + + let output = input.abs(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![2.0, 0.0, 3.0]); + assert_eq!(dinput, vec![-1.0, 0.0, 1.0]); + } +} + +#[tokio::test] +async fn test_backward_acos() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32]); + + let output = input.acos(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0], 0.5f32.acos()); + assert_close(dinput[0], -1.0f32 / (1.0f32 - 0.25f32).sqrt()); + } +} + +#[tokio::test] +async fn test_backward_acosh() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32]); + + let output = input.acosh(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0], 2.0f32.acosh()); + assert_close( + dinput[0], + 1.0f32 / ((2.0f32 - 1.0f32).sqrt() * (2.0f32 + 1.0f32).sqrt()), + ); + } +} + +#[tokio::test] +async fn test_backward_approximate_exp() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32]); + + let output = input.approximate_exp(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0], 1.0f32.exp()); + assert_close(dinput[0], 1.0f32.exp()); + } +} + +#[tokio::test] +async fn test_backward_asin() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32]); + + let output = input.asin(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0], 0.5f32.asin()); + assert_close(dinput[0], 1.0f32 / (1.0f32 - 0.25f32).sqrt()); + } +} + +#[tokio::test] +async fn test_backward_asinh() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.5f32]); + + let output = input.asinh(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0], 1.5f32.asinh()); + assert_close(dinput[0], 1.0f32 / (1.5f32 * 1.5f32 + 1.0f32).sqrt()); + } +} + +#[tokio::test] +async fn test_backward_atan() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32]); + + let output = input.atan(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0], 0.5f32.atan()); + assert_close(dinput[0], 1.0f32 / (1.0f32 + 0.25f32)); + } +} + +#[tokio::test] +async fn test_backward_atanh() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32]); + + let output = input.atanh(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0], 0.5f32.atanh()); + assert_close(dinput[0], 1.0f32 / (1.0f32 - 0.25f32)); + } +} + +#[tokio::test] +async fn test_backward_cos() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32]); + + let output = input.cos(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0], 0.5f32.cos()); + assert_close(dinput[0], -0.5f32.sin()); + } +} + +#[tokio::test] +async fn test_backward_cosh() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32]); + + let output = input.cosh(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0], 0.5f32.cosh()); + assert_close(dinput[0], 0.5f32.sinh()); + } +} + +#[tokio::test] +async fn test_backward_exp2() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32]); + + let output = input.exp2(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0], 2.0f32.exp2()); + assert_close(dinput[0], std::f32::consts::LN_2 * 2.0f32.exp2()); + } +} + +#[tokio::test] +async fn test_backward_less_approximate_exp() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32]); + + let output = input.less_approximate_exp(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0], 1.0f32.exp()); + assert_close(dinput[0], 1.0f32.exp()); + } +} + +#[tokio::test] +async fn test_backward_log2() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[4.0f32]); + + let output = input.log2(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0], 4.0f32.log2()); + assert_close(dinput[0], 1.0f32 / (4.0f32 * std::f32::consts::LN_2)); + } +} + +#[tokio::test] +async fn test_backward_sin() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32]); + + let output = input.sin(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0], 0.5f32.sin()); + assert_close(dinput[0], 0.5f32.cos()); + } +} + +#[tokio::test] +async fn test_backward_sinh() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32]); + + let output = input.sinh(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0], 0.5f32.sinh()); + assert_close(dinput[0], 0.5f32.cosh()); + } +} + +#[tokio::test] +async fn test_backward_tan() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32]); + + let output = input.tan(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0], 0.5f32.tan()); + assert_close(dinput[0], 1.0f32 / (0.5f32.cos() * 0.5f32.cos())); + } +} + +#[tokio::test] +async fn test_backward_tanh_exact() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32]); + + let output = input.tanh_exact(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0], 0.5f32.tanh()); + assert_close(dinput[0], 1.0f32 - 0.5f32.tanh().powi(2)); + } +} + +#[tokio::test] +async fn test_cast() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + + let output = input.cast::(); + let output_values = output.as_slice().await.unwrap().to_vec1(); + + assert_close(f32::from(output_values[0]), 1.0); + assert_close(f32::from(output_values[1]), 2.0); + assert_close(f32::from(output_values[2]), 3.0); + } +} + +#[tokio::test] +async fn test_arange() { + for device in test_devices().await { + let graph = Graph::new(); + + let output = Tensor::<1>::arange(&graph, &device, 1.0, 5.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + + assert_eq!(output_values, vec![1.0, 2.0, 3.0, 4.0]); + } +} + +#[tokio::test] +async fn test_arange_step() { + for device in test_devices().await { + let graph = Graph::new(); + + let output = Tensor::<1>::arange_step(&graph, &device, 1.0, 6.0, 2.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + + assert_eq!(output_values, vec![1.0, 3.0, 5.0]); + } +} + +#[tokio::test] +async fn test_full() { + for device in test_devices().await { + let graph = Graph::new(); + + let output: Tensor<2> = Tensor::full(&graph, &device, [2, 3], 1.5); + let output_values = output.raw().clone().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[2, 3]); + for row in 0..2 { + for col in 0..3 { + assert_close(output_values[[row, col]], 1.5); + } + } + } +} + +#[tokio::test] +async fn test_zeros_like() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0], [3.0, 4.0]]); + + let output = input.zeros_like(); + let output_values = output.raw().clone().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[2, 2]); + assert_close(output_values[[0, 0]], 0.0); + assert_close(output_values[[0, 1]], 0.0); + assert_close(output_values[[1, 0]], 0.0); + assert_close(output_values[[1, 1]], 0.0); + } +} + +#[tokio::test] +async fn test_from_array() { + for device in test_devices().await { + let graph = Graph::new(); + + let output: Tensor<2> = Tensor::from_array(&graph, &device, &[[1.0f32, 2.0], [3.0, 4.0]]); + let output_values = output.raw().clone().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[2, 2]); + assert_close(output_values[[0, 0]], 1.0); + assert_close(output_values[[0, 1]], 2.0); + assert_close(output_values[[1, 0]], 3.0); + assert_close(output_values[[1, 1]], 4.0); + } +} + +#[tokio::test] +async fn test_backward_add_broadcast_api() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0]); + let rhs: Tensor<2> = Tensor::new(&graph, &device, &[[10.0f32], [20.0]]); + + let output: Tensor<2> = lhs.add_::<2, 2>(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_close(output_values[0][0], 11.0); + assert_close(output_values[0][1], 12.0); + assert_close(output_values[1][0], 21.0); + assert_close(output_values[1][1], 22.0); + assert_close(dlhs[0], 2.0); + assert_close(dlhs[1], 2.0); + assert_close(drhs[0][0], 2.0); + assert_close(drhs[1][0], 2.0); + } +} + +#[tokio::test] +async fn test_backward_sub_broadcast_api() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<2> = Tensor::new(&graph, &device, &[[3.0f32], [4.0]]); + let rhs: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0]); + + let output: Tensor<2> = lhs.sub_::<1, 2>(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0][0], 2.0); + assert_close(output_values[0][1], 1.0); + assert_close(output_values[1][0], 3.0); + assert_close(output_values[1][1], 2.0); + assert_close(dlhs[0][0], 2.0); + assert_close(dlhs[1][0], 2.0); + assert_close(drhs[0], -2.0); + assert_close(drhs[1], -2.0); + } +} + +#[tokio::test] +async fn test_backward_mul_broadcast_api() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32, 3.0]); + let rhs: Tensor<2> = Tensor::new(&graph, &device, &[[10.0f32], [20.0]]); + + let output: Tensor<2> = lhs.mul_::<2, 2>(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_close(output_values[0][0], 20.0); + assert_close(output_values[0][1], 30.0); + assert_close(output_values[1][0], 40.0); + assert_close(output_values[1][1], 60.0); + assert_close(dlhs[0], 30.0); + assert_close(dlhs[1], 30.0); + assert_close(drhs[0][0], 5.0); + assert_close(drhs[1][0], 5.0); + } +} + +#[tokio::test] +async fn test_backward_div_broadcast_api() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<2> = Tensor::new(&graph, &device, &[[10.0f32], [20.0]]); + let rhs: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32, 4.0]); + + let output: Tensor<2> = lhs.div_::<1, 2>(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0][0], 5.0); + assert_close(output_values[0][1], 2.5); + assert_close(output_values[1][0], 10.0); + assert_close(output_values[1][1], 5.0); + assert_close(dlhs[0][0], 0.75); + assert_close(dlhs[1][0], 0.75); + assert_close(drhs[0], -7.5); + assert_close(drhs[1], -1.875); + } +} + +#[tokio::test] +async fn test_backward_pow_broadcast_api() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32, 3.0]); + let rhs: Tensor<2> = Tensor::new(&graph, &device, &[[2.0f32], [1.0]]); + + let output: Tensor<2> = lhs.pow_::<2, 2>(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_close(output_values[0][0], 4.0); + assert_close(output_values[0][1], 9.0); + assert_close(output_values[1][0], 2.0); + assert_close(output_values[1][1], 3.0); + assert_close(dlhs[0], 5.0); + assert_close(dlhs[1], 7.0); + assert_close(drhs[0][0], 12.660099); + assert_close(drhs[1][0], 4.6821313); + } +} + +#[tokio::test] +async fn test_backward_chunk() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new( + &graph, + &device, + &[[1.0f32, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], + ); + + let chunks = input.chunk(2, 1); + assert_eq!(chunks.len(), 2); + let first = chunks[0].raw().clone().as_slice().await.unwrap().to_vec2(); + let second = chunks[1].raw().clone().as_slice().await.unwrap().to_vec2(); + let loss = chunks[0] + .flatten_all() + .sum() + .add(&chunks[1].flatten_all().sum().mul_scalar(2.0)); + let gradients = loss.backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_close(first[0][0], 1.0); + assert_close(first[0][1], 2.0); + assert_close(first[1][0], 5.0); + assert_close(first[1][1], 6.0); + assert_close(second[0][0], 3.0); + assert_close(second[0][1], 4.0); + assert_close(second[1][0], 7.0); + assert_close(second[1][1], 8.0); + assert_close(dinput[0][0], 1.0); + assert_close(dinput[0][1], 1.0); + assert_close(dinput[0][2], 2.0); + assert_close(dinput[0][3], 2.0); + assert_close(dinput[1][0], 1.0); + assert_close(dinput[1][1], 1.0); + assert_close(dinput[1][2], 2.0); + assert_close(dinput[1][3], 2.0); + } +} + +#[tokio::test] +async fn test_backward_matmul() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0], [3.0, 4.0]]); + let rhs: Tensor<2> = Tensor::new(&graph, &device, &[[5.0f32, 6.0], [7.0, 8.0]]); + + let output = lhs.matmul(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dlhs = gradients.get(&lhs).unwrap().as_slice().await.unwrap(); + let drhs = gradients.get(&rhs).unwrap().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[2, 2]); + assert_close(output_values[[0, 0]], 19.0); + assert_close(output_values[[0, 1]], 22.0); + assert_close(output_values[[1, 0]], 43.0); + assert_close(output_values[[1, 1]], 50.0); + + assert_close(dlhs[[0, 0]], 11.0); + assert_close(dlhs[[0, 1]], 15.0); + assert_close(dlhs[[1, 0]], 11.0); + assert_close(dlhs[[1, 1]], 15.0); + + assert_close(drhs[[0, 0]], 4.0); + assert_close(drhs[[0, 1]], 4.0); + assert_close(drhs[[1, 0]], 6.0); + assert_close(drhs[[1, 1]], 6.0); + } +} + +#[tokio::test] +async fn test_backward_t() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0], [3.0, 4.0]]); + + let output = input.t(); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dinput = gradients.get(&input).unwrap().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[2, 2]); + assert_close(output_values[[0, 0]], 1.0); + assert_close(output_values[[0, 1]], 3.0); + assert_close(output_values[[1, 0]], 2.0); + assert_close(output_values[[1, 1]], 4.0); + + assert_close(dinput[[0, 0]], 1.0); + assert_close(dinput[[0, 1]], 1.0); + assert_close(dinput[[1, 0]], 1.0); + assert_close(dinput[[1, 1]], 1.0); + } +} + +#[tokio::test] +async fn test_backward_pool() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<3> = Tensor::new(&graph, &device, &[[[1.0f32, 2.0, 3.0, 4.0]]]); + + let output = input.pool::<1, 4, 5, 4>([(2, 1)], |windowed, axis| windowed.mean::<3>(axis)); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dinput = gradients.get(&input).unwrap().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[1, 1, 3]); + assert_close(output_values[[0, 0, 0]], 1.5); + assert_close(output_values[[0, 0, 1]], 2.5); + assert_close(output_values[[0, 0, 2]], 3.5); + + assert_close(dinput[[0, 0, 0]], 0.5); + assert_close(dinput[[0, 0, 1]], 1.0); + assert_close(dinput[[0, 0, 2]], 1.0); + assert_close(dinput[[0, 0, 3]], 0.5); + } +} + +#[tokio::test] +async fn test_backward_pool_max() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<3> = Tensor::new(&graph, &device, &[[[1.0f32, 4.0, 2.0, 3.0]]]); + + let output = input.pool_max::<1, 4, 5, 4>([(2, 1)]); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dinput = gradients.get(&input).unwrap().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[1, 1, 3]); + assert_close(output_values[[0, 0, 0]], 4.0); + assert_close(output_values[[0, 0, 1]], 4.0); + assert_close(output_values[[0, 0, 2]], 3.0); + + assert_close(dinput[[0, 0, 0]], 0.0); + assert_close(dinput[[0, 0, 1]], 2.0); + assert_close(dinput[[0, 0, 2]], 0.0); + assert_close(dinput[[0, 0, 3]], 1.0); + } +} + +#[tokio::test] +async fn test_backward_pool_min() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<3> = Tensor::new(&graph, &device, &[[[1.0f32, 4.0, 2.0, 3.0]]]); + + let output = input.pool_min::<1, 4, 5, 4>([(2, 1)]); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dinput = gradients.get(&input).unwrap().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[1, 1, 3]); + assert_close(output_values[[0, 0, 0]], 1.0); + assert_close(output_values[[0, 0, 1]], 2.0); + assert_close(output_values[[0, 0, 2]], 2.0); + + assert_close(dinput[[0, 0, 0]], 1.0); + assert_close(dinput[[0, 0, 1]], 0.0); + assert_close(dinput[[0, 0, 2]], 2.0); + assert_close(dinput[[0, 0, 3]], 0.0); + } +} + +#[tokio::test] +async fn test_backward_q_mat_mul() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 1.0, 1.0, 1.0]]); + let weight_bytes: Vec = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0] + .into_iter() + .flat_map(|value| value.to_le_bytes()) + .collect(); + let weights = crate::QMatrix::from_raw_bytes( + &device, + [2, 4], + &weight_bytes, + fusor_gguf::GgmlType::F32, + ) + .unwrap(); + + let output = input.q_mat_mul(&weights); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dinput = gradients.get(&input).unwrap().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[1, 2]); + assert_close(output_values[[0, 0]], 10.0); + assert_close(output_values[[0, 1]], 26.0); + + assert_close(dinput[[0, 0]], 6.0); + assert_close(dinput[[0, 1]], 8.0); + assert_close(dinput[[0, 2]], 10.0); + assert_close(dinput[[0, 3]], 12.0); + } +} + +#[tokio::test] +async fn test_backward_stack() { + for device in test_devices().await { + let graph = Graph::new(); + let first: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0]); + let second: Tensor<1> = Tensor::new(&graph, &device, &[3.0f32, 4.0]); + + let output = Tensor::stack::<2>(vec![first.clone(), second.clone()], 0); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dfirst = gradients + .get(&first) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + let dsecond = gradients + .get(&second) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values.shape(), &[2, 2]); + assert_close(output_values[[0, 0]], 1.0); + assert_close(output_values[[0, 1]], 2.0); + assert_close(output_values[[1, 0]], 3.0); + assert_close(output_values[[1, 1]], 4.0); + + assert_close(dfirst[0], 1.0); + assert_close(dfirst[1], 1.0); + assert_close(dsecond[0], 1.0); + assert_close(dsecond[1], 1.0); + } +} + +#[tokio::test] +async fn test_backward_rope() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<4> = Tensor::new( + &graph, + &device, + &[[[[1.0f32, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]]], + ); + let cos: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 1.0], [1.0, 1.0]]); + let sin: Tensor<2> = Tensor::new(&graph, &device, &[[0.0f32, 0.0], [0.0, 0.0]]); + + let output = input.rope(&cos, &sin); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dinput = gradients.get(&input).unwrap().as_slice().await.unwrap(); + let dcos = gradients.get(&cos).unwrap().as_slice().await.unwrap(); + let dsin = gradients.get(&sin).unwrap().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[1, 1, 2, 4]); + assert_close(output_values[[0, 0, 0, 0]], 1.0); + assert_close(output_values[[0, 0, 0, 1]], 2.0); + assert_close(output_values[[0, 0, 0, 2]], 3.0); + assert_close(output_values[[0, 0, 0, 3]], 4.0); + assert_close(output_values[[0, 0, 1, 0]], 5.0); + assert_close(output_values[[0, 0, 1, 1]], 6.0); + assert_close(output_values[[0, 0, 1, 2]], 7.0); + assert_close(output_values[[0, 0, 1, 3]], 8.0); + + for index in [ + [0, 0, 0, 0], + [0, 0, 0, 1], + [0, 0, 0, 2], + [0, 0, 0, 3], + [0, 0, 1, 0], + [0, 0, 1, 1], + [0, 0, 1, 2], + [0, 0, 1, 3], + ] { + assert_close(dinput[index], 1.0); + } + + assert_close(dcos[[0, 0]], 4.0); + assert_close(dcos[[0, 1]], 6.0); + assert_close(dcos[[1, 0]], 12.0); + assert_close(dcos[[1, 1]], 14.0); + + assert_close(dsin[[0, 0]], -2.0); + assert_close(dsin[[0, 1]], -2.0); + assert_close(dsin[[1, 0]], -2.0); + assert_close(dsin[[1, 1]], -2.0); + } +} + +#[tokio::test] +async fn test_backward_rope_fused() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<4> = Tensor::new( + &graph, + &device, + &[[[[1.0f32, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]]], + ); + let cos: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 1.0], [1.0, 1.0]]); + let sin: Tensor<2> = Tensor::new(&graph, &device, &[[0.0f32, 0.0], [0.0, 0.0]]); + + let output = input.rope_fused(&cos, &sin); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dinput = gradients.get(&input).unwrap().as_slice().await.unwrap(); + let dcos = gradients.get(&cos).unwrap().as_slice().await.unwrap(); + let dsin = gradients.get(&sin).unwrap().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[1, 1, 2, 4]); + assert_close(output_values[[0, 0, 0, 0]], 1.0); + assert_close(output_values[[0, 0, 0, 1]], 2.0); + assert_close(output_values[[0, 0, 0, 2]], 3.0); + assert_close(output_values[[0, 0, 0, 3]], 4.0); + assert_close(output_values[[0, 0, 1, 0]], 5.0); + assert_close(output_values[[0, 0, 1, 1]], 6.0); + assert_close(output_values[[0, 0, 1, 2]], 7.0); + assert_close(output_values[[0, 0, 1, 3]], 8.0); + + for index in [ + [0, 0, 0, 0], + [0, 0, 0, 1], + [0, 0, 0, 2], + [0, 0, 0, 3], + [0, 0, 1, 0], + [0, 0, 1, 1], + [0, 0, 1, 2], + [0, 0, 1, 3], + ] { + assert_close(dinput[index], 1.0); + } + + assert_close(dcos[[0, 0]], 3.0); + assert_close(dcos[[0, 1]], 7.0); + assert_close(dcos[[1, 0]], 11.0); + assert_close(dcos[[1, 1]], 15.0); + + assert_close(dsin[[0, 0]], -1.0); + assert_close(dsin[[0, 1]], -1.0); + assert_close(dsin[[1, 0]], -1.0); + assert_close(dsin[[1, 1]], -1.0); + } +} + +#[tokio::test] +async fn test_backward_rope_interleaved() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<4> = Tensor::new( + &graph, + &device, + &[[[[1.0f32, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]]], + ); + let cos: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 1.0], [1.0, 1.0]]); + let sin: Tensor<2> = Tensor::new(&graph, &device, &[[0.0f32, 0.0], [0.0, 0.0]]); + + let output = input.rope_interleaved(&cos, &sin); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dinput = gradients.get(&input).unwrap().as_slice().await.unwrap(); + let dcos = gradients.get(&cos).unwrap().as_slice().await.unwrap(); + let dsin = gradients.get(&sin).unwrap().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[1, 1, 2, 4]); + assert_close(output_values[[0, 0, 0, 0]], 1.0); + assert_close(output_values[[0, 0, 0, 1]], 2.0); + assert_close(output_values[[0, 0, 0, 2]], 3.0); + assert_close(output_values[[0, 0, 0, 3]], 4.0); + assert_close(output_values[[0, 0, 1, 0]], 5.0); + assert_close(output_values[[0, 0, 1, 1]], 6.0); + assert_close(output_values[[0, 0, 1, 2]], 7.0); + assert_close(output_values[[0, 0, 1, 3]], 8.0); + + for index in [ + [0, 0, 0, 0], + [0, 0, 0, 1], + [0, 0, 0, 2], + [0, 0, 0, 3], + [0, 0, 1, 0], + [0, 0, 1, 1], + [0, 0, 1, 2], + [0, 0, 1, 3], + ] { + assert_close(dinput[index], 1.0); + } + + assert_close(dcos[[0, 0]], 3.0); + assert_close(dcos[[0, 1]], 7.0); + assert_close(dcos[[1, 0]], 11.0); + assert_close(dcos[[1, 1]], 15.0); + + assert_close(dsin[[0, 0]], -1.0); + assert_close(dsin[[0, 1]], -1.0); + assert_close(dsin[[1, 0]], -1.0); + assert_close(dsin[[1, 1]], -1.0); + } +} + +#[tokio::test] +async fn test_backward_rope_normal_fused() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<4> = Tensor::new( + &graph, + &device, + &[[[[1.0f32, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]]], + ); + let cos: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 1.0], [1.0, 1.0]]); + let sin: Tensor<2> = Tensor::new(&graph, &device, &[[0.0f32, 0.0], [0.0, 0.0]]); + + let output = input.rope_normal_fused(&cos, &sin); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dinput = gradients.get(&input).unwrap().as_slice().await.unwrap(); + let dcos = gradients.get(&cos).unwrap().as_slice().await.unwrap(); + let dsin = gradients.get(&sin).unwrap().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[1, 1, 2, 4]); + assert_close(output_values[[0, 0, 0, 0]], 1.0); + assert_close(output_values[[0, 0, 0, 1]], 2.0); + assert_close(output_values[[0, 0, 0, 2]], 3.0); + assert_close(output_values[[0, 0, 0, 3]], 4.0); + assert_close(output_values[[0, 0, 1, 0]], 5.0); + assert_close(output_values[[0, 0, 1, 1]], 6.0); + assert_close(output_values[[0, 0, 1, 2]], 7.0); + assert_close(output_values[[0, 0, 1, 3]], 8.0); + + for index in [ + [0, 0, 0, 0], + [0, 0, 0, 1], + [0, 0, 0, 2], + [0, 0, 0, 3], + [0, 0, 1, 0], + [0, 0, 1, 1], + [0, 0, 1, 2], + [0, 0, 1, 3], + ] { + assert_close(dinput[index], 1.0); + } + + assert_close(dcos[[0, 0]], 4.0); + assert_close(dcos[[0, 1]], 6.0); + assert_close(dcos[[1, 0]], 12.0); + assert_close(dcos[[1, 1]], 14.0); + + assert_close(dsin[[0, 0]], -2.0); + assert_close(dsin[[0, 1]], -2.0); + assert_close(dsin[[1, 0]], -2.0); + assert_close(dsin[[1, 1]], -2.0); + } +} + +#[tokio::test] +async fn test_backward_pow() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32]); + let rhs: Tensor<1> = Tensor::new(&graph, &device, &[3.0f32]); + + let output = lhs.pow(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0], 8.0); + assert_close(dlhs[0], 12.0); + assert_close(drhs[0], 8.0 * 2.0f32.ln()); + } +} + +#[tokio::test] +async fn test_backward_pow_elementwise() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[3.0f32]); + + let output = input.pow_elementwise(2.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0], 9.0); + assert_close(dinput[0], 6.0); + } +} + +#[tokio::test] +async fn test_backward_pow_scalar() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[4.0f32]); + + let output = input.pow_scalar(0.5); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(output_values[0], 2.0); + assert_close(dinput[0], 0.25); + } +} + +#[tokio::test] +async fn test_autograd_rms_norm() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]); + let weight: Tensor<2> = Tensor::constant_from_raw( + &graph, + RawTensor::from_slice(&device, [1, 3], &[1.0f32, 1.0, 1.0]), + ); + + let output = input.rms_norm(&weight, 1e-5); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + + let expected = [[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]].map(|row| { + let mean_sq = row.iter().map(|value| value * value).sum::() / row.len() as f32; + let scale = 1.0 / (mean_sq + 1e-5).sqrt(); + row.map(|value| value * scale) + }); + + for (actual_row, expected_row) in output_values.iter().zip(expected.iter()) { + for (actual, expected) in actual_row.iter().zip(expected_row.iter()) { + assert_close(*actual, *expected); + } + } + + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + // d/dx_j sum_k x_k * (mean(x^2) + eps)^-1/2 + // = 1/rms - x_j * sum(x) / (n * rms^3) + let expected_grads = [[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]].map(|row| { + let n = row.len() as f32; + let mean_sq = row.iter().map(|value| value * value).sum::() / n; + let rms = (mean_sq + 1e-5).sqrt(); + let sum = row.iter().sum::(); + row.map(|value| 1.0 / rms - value * sum / (n * rms.powi(3))) + }); + for (actual_row, expected_row) in dinput.iter().zip(expected_grads.iter()) { + for (actual, expected) in actual_row.iter().zip(expected_row.iter()) { + assert_close(*actual, *expected); + } + } + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 3], + &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], + move |graph, x| { + let weight = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &[1.0f32, 1.0, 1.0]), + ); + x.rms_norm(&weight, 1e-5).sum(1).sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_backward_matmul_with_broadcast_bias() { + for device in test_devices().await { + let graph = Graph::new(); + + let x: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]); + let w: Tensor<2> = Tensor::new(&graph, &device, &[[0.5f32], [1.0], [1.5]]); + let b: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32]); + + let y = x.mat_mul(&w).add(&b.broadcast_as([2, 1])); + let loss = y.sum(1).sum(); + + let gradients = loss.backward().unwrap(); + let dw = gradients + .get(&w) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + let db = gradients + .get(&b) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_close(dw[0][0], 5.0); + assert_close(dw[1][0], 7.0); + assert_close(dw[2][0], 9.0); + assert_close(db[0], 2.0); + } +} + +#[tokio::test] +async fn test_backward_embedding() { + for device in test_devices().await { + let graph = Graph::new(); + + let table: Tensor<2> = + Tensor::new(&graph, &device, &[[1.0f32, 2.0], [3.0, 4.0], [5.0, 6.0]]); + let indices: RawTensor<2, u32> = RawTensor::new(&device, &[[0u32, 2u32]]); + let embedded = table.embedding(&indices); + let loss = embedded.sum(2).sum(1).sum(); + + let gradients = loss.backward().unwrap(); + let dtable = gradients + .get(&table) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_close(dtable[0][0], 1.0); + assert_close(dtable[0][1], 1.0); + assert_close(dtable[1][0], 0.0); + assert_close(dtable[1][1], 0.0); + assert_close(dtable[2][0], 1.0); + assert_close(dtable[2][1], 1.0); + } +} + +#[tokio::test] +async fn test_backward_gather_last() { + for device in test_devices().await { + let graph = Graph::new(); + + let values: Tensor<2> = + Tensor::new(&graph, &device, &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]); + let indices: RawTensor<1, u32> = RawTensor::new(&device, &[2u32, 0u32]); + let gathered = values.gather_last(&indices); + let loss = gathered.sum(); + + let gradients = loss.backward().unwrap(); + let dvalues = gradients + .get(&values) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + assert_close(dvalues[0][0], 0.0); + assert_close(dvalues[0][1], 0.0); + assert_close(dvalues[0][2], 1.0); + assert_close(dvalues[1][0], 1.0); + assert_close(dvalues[1][1], 0.0); + assert_close(dvalues[1][2], 0.0); + } +} + +#[tokio::test] +async fn test_backward_softmax_last_dim_fused_matches_composite() { + for device in test_devices().await { + let input_data = &[ + [[0.2f32, -0.4, 1.1], [0.5, 0.3, -0.7]], + [[-1.0, 0.8, 0.6], [0.9, -0.2, 0.1]], + ]; + + let fused_graph = Graph::new(); + let fused_input: Tensor<3> = Tensor::new(&fused_graph, &device, input_data); + let fused_output = fused_input.softmax_last_dim_fused::<2>(); + let fused_loss = fused_output.sqr().reshape([12]).sum(); + let fused_gradients = fused_loss.backward().unwrap(); + + let composite_graph = Graph::new(); + let composite_input: Tensor<3> = Tensor::new(&composite_graph, &device, input_data); + let composite_output = composite_input.softmax_last_dim::<2>(); + let composite_loss = composite_output.sqr().reshape([12]).sum(); + let composite_gradients = composite_loss.backward().unwrap(); + + let fused_output = flatten(fused_output.raw().clone()).await; + let composite_output = flatten(composite_output.raw().clone()).await; + let fused_dx = flatten(fused_gradients.get(&fused_input).unwrap()).await; + let composite_dx = flatten(composite_gradients.get(&composite_input).unwrap()).await; + + assert_slice_close(&fused_output, &composite_output); + assert_slice_close(&fused_dx, &composite_dx); + } +} + +#[tokio::test] +async fn test_backward_rms_norm_fused_matches_composite() { + for device in test_devices().await { + let input_data = &[ + [[0.3f32, -1.2, 0.7], [1.5, 0.1, -0.8]], + [[-0.4, 0.9, 1.3], [0.2, -0.6, 0.5]], + ]; + let weight_data = &[1.0f32, 0.75, 1.25]; + let eps = 1e-5; + + let fused_graph = Graph::new(); + let fused_input: Tensor<3> = Tensor::new(&fused_graph, &device, input_data); + let fused_weight: Tensor<1> = Tensor::new(&fused_graph, &device, weight_data); + let fused_output = fused_input.rms_norm_fused_no_bias::<1, 2>(&fused_weight, eps); + let fused_loss = fused_output.sqr().reshape([12]).sum(); + let fused_gradients = fused_loss.backward().unwrap(); + + let composite_graph = Graph::new(); + let composite_input: Tensor<3> = Tensor::new(&composite_graph, &device, input_data); + let composite_weight: Tensor<3> = Tensor::from_slice( + &composite_graph, + &device, + [1, 1, 3], + weight_data, + ); + let composite_output = composite_input.rms_norm(&composite_weight, eps); + let composite_loss = composite_output.sqr().reshape([12]).sum(); + let composite_gradients = composite_loss.backward().unwrap(); + + let fused_output = flatten(fused_output.raw().clone()).await; + let composite_output = flatten(composite_output.raw().clone()).await; + let fused_dx = flatten(fused_gradients.get(&fused_input).unwrap()).await; + let composite_dx = flatten(composite_gradients.get(&composite_input).unwrap()).await; + let fused_dw = flatten(fused_gradients.get(&fused_weight).unwrap()).await; + let composite_dw = flatten(composite_gradients.get(&composite_weight).unwrap()).await; + + assert_slice_close(&fused_output, &composite_output); + assert_slice_close(&fused_dx, &composite_dx); + assert_slice_close(&fused_dw, &composite_dw); + } +} + +#[tokio::test] +async fn test_backward_layer_norm_last_dim_fused_matches_composite() { + for device in test_devices().await { + let input_data = &[ + [[0.25f32, -0.5, 1.0], [1.25, -1.5, 0.75]], + [[-0.8, 0.4, 1.2], [0.6, -0.1, -0.9]], + ]; + let weight_data = &[1.0f32, 0.9, 1.1]; + let bias_data = &[0.1f32, -0.2, 0.05]; + let eps = 1e-5; + + let fused_graph = Graph::new(); + let fused_input: Tensor<3> = Tensor::new(&fused_graph, &device, input_data); + let fused_weight: Tensor<1> = Tensor::new(&fused_graph, &device, weight_data); + let fused_bias: Tensor<1> = Tensor::new(&fused_graph, &device, bias_data); + let fused_output = + fused_input.layer_norm_last_dim_fused::<2, 1>(&fused_weight, Some(&fused_bias), eps); + let fused_loss = fused_output.sqr().reshape([12]).sum(); + let fused_gradients = fused_loss.backward().unwrap(); + + let composite_graph = Graph::new(); + let composite_input: Tensor<3> = Tensor::new(&composite_graph, &device, input_data); + let composite_weight: Tensor<3> = Tensor::from_slice( + &composite_graph, + &device, + [1, 1, 3], + weight_data, + ); + let composite_bias: Tensor<3> = Tensor::from_slice( + &composite_graph, + &device, + [1, 1, 3], + bias_data, + ); + let composite_output = + composite_input.layer_norm(&composite_weight, Some(&composite_bias), eps, true); + let composite_loss = composite_output.sqr().reshape([12]).sum(); + let composite_gradients = composite_loss.backward().unwrap(); + + let fused_output = flatten(fused_output.raw().clone()).await; + let composite_output = flatten(composite_output.raw().clone()).await; + let fused_dx = flatten(fused_gradients.get(&fused_input).unwrap()).await; + let composite_dx = flatten(composite_gradients.get(&composite_input).unwrap()).await; + let fused_dw = flatten(fused_gradients.get(&fused_weight).unwrap()).await; + let composite_dw = flatten(composite_gradients.get(&composite_weight).unwrap()).await; + let fused_db = flatten(fused_gradients.get(&fused_bias).unwrap()).await; + let composite_db = flatten(composite_gradients.get(&composite_bias).unwrap()).await; + + assert_slice_close(&fused_output, &composite_output); + assert_slice_close(&fused_dx, &composite_dx); + assert_slice_close(&fused_dw, &composite_dw); + assert_slice_close(&fused_db, &composite_db); + } +} + +#[tokio::test] +async fn test_backward_flash_attention_matches_composite() { + for device in test_devices().await { + let q_data = &[[[[0.2f32, 0.6], [1.0, -0.3]]]]; + let k_data = &[[[[0.4f32, -0.7], [0.9, 0.1]]]]; + let v_data = &[[[[1.1f32, -0.5], [0.3, 0.8]]]]; + let scale = (2.0f32).sqrt(); + + let fused_graph = Graph::new(); + let fused_q: Tensor<4> = Tensor::new(&fused_graph, &device, q_data); + let fused_k: Tensor<4> = Tensor::new(&fused_graph, &device, k_data); + let fused_v: Tensor<4> = Tensor::new(&fused_graph, &device, v_data); + let fused_output = fused_q.flash_attention(&fused_k, &fused_v, scale, None); + let fused_loss = fused_output.sqr().reshape([4]).sum(); + let fused_gradients = fused_loss.backward().unwrap(); + + let composite_graph = Graph::new(); + let composite_q: Tensor<4> = Tensor::new(&composite_graph, &device, q_data); + let composite_k: Tensor<4> = Tensor::new(&composite_graph, &device, k_data); + let composite_v: Tensor<4> = Tensor::new(&composite_graph, &device, v_data); + let composite_output = + composite_q.flash_attention_composite(&composite_k, &composite_v, scale, None); + let composite_loss = composite_output.sqr().reshape([4]).sum(); + let composite_gradients = composite_loss.backward().unwrap(); + + let fused_output = flatten(fused_output.raw().clone()).await; + let composite_output = flatten(composite_output.raw().clone()).await; + let fused_dq = flatten(fused_gradients.get(&fused_q).unwrap()).await; + let composite_dq = flatten(composite_gradients.get(&composite_q).unwrap()).await; + let fused_dk = flatten(fused_gradients.get(&fused_k).unwrap()).await; + let composite_dk = flatten(composite_gradients.get(&composite_k).unwrap()).await; + let fused_dv = flatten(fused_gradients.get(&fused_v).unwrap()).await; + let composite_dv = flatten(composite_gradients.get(&composite_v).unwrap()).await; + + assert_slice_close(&fused_output, &composite_output); + assert_slice_close(&fused_dq, &composite_dq); + assert_slice_close(&fused_dk, &composite_dk); + assert_slice_close(&fused_dv, &composite_dv); + } +} + +#[tokio::test] +async fn test_backward_mat_mul_rank3() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs_data = (1..=24).map(|n| n as f32).collect::>(); + let rhs_data = (1..=40).map(|n| n as f32).collect::>(); + let lhs: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 3, 4], &lhs_data); + let rhs: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 4, 5], &rhs_data); + + let output = lhs.mat_mul(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dlhs = flatten(gradients.get(&lhs).unwrap()).await; + let drhs = flatten(gradients.get(&rhs).unwrap()).await; + + assert_eq!(output_values.shape(), &[2, 3, 5]); + assert_close(output_values[[0, 0, 0]], 110.0); + assert_close(output_values[[1, 2, 4]], 2950.0); + + // with an all-ones seed, dlhs[b, i, k] = sum_j rhs[b, k, j] and + // drhs[b, k, j] = sum_i lhs[b, i, k] + for b in 0..2 { + for i in 0..3 { + for k in 0..4 { + let expected = (0..5).map(|j| rhs_data[b * 20 + k * 5 + j]).sum::(); + assert_close(dlhs[b * 12 + i * 4 + k], expected); + } + } + for k in 0..4 { + for j in 0..5 { + let expected = (0..3).map(|i| lhs_data[b * 12 + i * 4 + k]).sum::(); + assert_close(drhs[b * 20 + k * 5 + j], expected); + } + } + } + + let lhs_small = lhs_data + .iter() + .map(|value| value * 0.05) + .collect::>(); + let rhs_small = rhs_data + .iter() + .map(|value| value * 0.03) + .collect::>(); + let fd_device = device.clone(); + let fd_rhs = rhs_small.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 3, 4], + &lhs_small, + move |graph, lhs| { + let rhs = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 4, 5], &fd_rhs), + ); + lhs.mat_mul(&rhs).sqr().flatten_all().sum() + }, + ) + .await; + let fd_device = device.clone(); + let fd_lhs = lhs_small.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 4, 5], + &rhs_small, + move |graph, rhs| { + let lhs = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3, 4], &fd_lhs), + ); + lhs.mat_mul(&rhs).sqr().flatten_all().sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_backward_cat_dim0() { + for device in test_devices().await { + let graph = Graph::new(); + let first_data = (1..=6).map(|n| n as f32).collect::>(); + let second_data = (7..=18).map(|n| n as f32).collect::>(); + let first: Tensor<3> = Tensor::from_slice(&graph, &device, [1, 2, 3], &first_data); + let second: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 2, 3], &second_data); + + let output = Tensor::cat(vec![first.clone(), second.clone()], 0); + let output_values = flatten(output.raw().clone()).await; + let seed_data = (0..18).map(|n| n as f32 + 10.0).collect::>(); + let seed = RawTensor::from_slice(&device, [3, 2, 3], &seed_data); + let gradients = output.backward_with(seed).unwrap(); + let dfirst = flatten(gradients.get(&first).unwrap()).await; + let dsecond = flatten(gradients.get(&second).unwrap()).await; + + assert_eq!(output.shape(), [3, 2, 3]); + assert_eq!( + output_values, + (1..=18).map(|n| n as f32).collect::>() + ); + assert_eq!(dfirst, seed_data[..6].to_vec()); + assert_eq!(dsecond, seed_data[6..].to_vec()); + } +} + +#[tokio::test] +async fn test_backward_cat_dim1() { + for device in test_devices().await { + let graph = Graph::new(); + let first_data = (1..=6).map(|n| n as f32).collect::>(); + let second_data = (10..=21).map(|n| n as f32).collect::>(); + let first: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 1, 3], &first_data); + let second: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 2, 3], &second_data); + + let output = Tensor::cat(vec![first.clone(), second.clone()], 1); + let output_values = flatten(output.raw().clone()).await; + let seed_data = (0..18).map(|n| n as f32 + 10.0).collect::>(); + let seed = RawTensor::from_slice(&device, [2, 3, 3], &seed_data); + let gradients = output.backward_with(seed).unwrap(); + let dfirst = flatten(gradients.get(&first).unwrap()).await; + let dsecond = flatten(gradients.get(&second).unwrap()).await; + + assert_eq!(output.shape(), [2, 3, 3]); + assert_eq!( + output_values, + vec![ + 1.0, 2.0, 3.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 4.0, 5.0, 6.0, 16.0, 17.0, 18.0, + 19.0, 20.0, 21.0 + ] + ); + + let mut expected_dfirst = Vec::new(); + let mut expected_dsecond = Vec::new(); + for i in 0..2 { + for j in 0..3 { + for k in 0..3 { + let value = seed_data[i * 9 + j * 3 + k]; + if j < 1 { + expected_dfirst.push(value); + } else { + expected_dsecond.push(value); + } + } + } + } + assert_eq!(dfirst, expected_dfirst); + assert_eq!(dsecond, expected_dsecond); + + let first_small = first_data + .iter() + .map(|value| value * 0.1) + .collect::>(); + let second_small = second_data + .iter() + .map(|value| value * 0.1) + .collect::>(); + let fd_device = device.clone(); + let fd_second = second_small.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 1, 3], + &first_small, + move |graph, first| { + let second = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 3], &fd_second), + ); + Tensor::cat(vec![first, second], 1) + .sqr() + .flatten_all() + .sum() + }, + ) + .await; + let fd_device = device.clone(); + let fd_first = first_small.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 2, 3], + &second_small, + move |graph, second| { + let first = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 1, 3], &fd_first), + ); + Tensor::cat(vec![first, second], 1) + .sqr() + .flatten_all() + .sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_backward_cat_dim2() { + for device in test_devices().await { + let graph = Graph::new(); + let first_data = (1..=4).map(|n| n as f32).collect::>(); + let second_data = (5..=12).map(|n| n as f32).collect::>(); + let first: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 2, 1], &first_data); + let second: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 2, 2], &second_data); + + let output = Tensor::cat(vec![first.clone(), second.clone()], 2); + let output_values = flatten(output.raw().clone()).await; + let seed_data = (0..12).map(|n| n as f32 + 10.0).collect::>(); + let seed = RawTensor::from_slice(&device, [2, 2, 3], &seed_data); + let gradients = output.backward_with(seed).unwrap(); + let dfirst = flatten(gradients.get(&first).unwrap()).await; + let dsecond = flatten(gradients.get(&second).unwrap()).await; + + assert_eq!(output.shape(), [2, 2, 3]); + assert_eq!( + output_values, + vec![ + 1.0, 5.0, 6.0, 2.0, 7.0, 8.0, 3.0, 9.0, 10.0, 4.0, 11.0, 12.0 + ] + ); + + let mut expected_dfirst = Vec::new(); + let mut expected_dsecond = Vec::new(); + for i in 0..2 { + for j in 0..2 { + for k in 0..3 { + let value = seed_data[i * 6 + j * 3 + k]; + if k < 1 { + expected_dfirst.push(value); + } else { + expected_dsecond.push(value); + } + } + } + } + assert_eq!(dfirst, expected_dfirst); + assert_eq!(dsecond, expected_dsecond); + + let first_small = first_data + .iter() + .map(|value| value * 0.1) + .collect::>(); + let second_small = second_data + .iter() + .map(|value| value * 0.1) + .collect::>(); + let fd_device = device.clone(); + let fd_second = second_small.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 2, 1], + &first_small, + move |graph, first| { + let second = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 2], &fd_second), + ); + Tensor::cat(vec![first, second], 2) + .sqr() + .flatten_all() + .sum() + }, + ) + .await; + let fd_device = device.clone(); + let fd_first = first_small.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 2, 2], + &second_small, + move |graph, second| { + let first = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 1], &fd_first), + ); + Tensor::cat(vec![first, second], 2) + .sqr() + .flatten_all() + .sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_backward_log() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32, 1.5, 2.5]); + + let output = input.log(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + for (value, input) in output_values.iter().zip([0.5f32, 1.5, 2.5]) { + assert_close(*value, input.ln()); + } + for (value, input) in dinput.iter().zip([0.5f32, 1.5, 2.5]) { + assert_close(*value, 1.0 / input); + } + + assert_gradient_matches_finite_difference(&device, [3], &[0.5, 1.5, 2.5], |_, x| { + x.log().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_neg() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.5f32, -2.0, 0.5]); + + let output = input.neg(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![-1.5, 2.0, -0.5]); + assert_eq!(dinput, vec![-1.0, -1.0, -1.0]); + + assert_gradient_matches_finite_difference(&device, [3], &[1.5, -2.0, 0.5], |_, x| { + x.neg().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_exp() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.0f32, 0.5, -1.0]); + + let output = input.exp(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + for (value, input) in output_values.iter().zip([0.0f32, 0.5, -1.0]) { + assert_close(*value, input.exp()); + } + for (value, input) in dinput.iter().zip([0.0f32, 0.5, -1.0]) { + assert_close(*value, input.exp()); + } + + assert_gradient_matches_finite_difference(&device, [3], &[0.0, 0.5, -1.0], |_, x| { + x.exp().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_log_sum_exp() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = + Tensor::new(&graph, &device, &[[0.0f32, 0.5, 1.0], [1.0, -1.0, 0.0]]); + + let output = input.exp().sum_keepdim(1).log(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + let gradients = output.reshape([2]).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + + // d/dx_j log(sum_k exp(x_k)) = softmax(x)_j + let rows = [[0.0f32, 0.5, 1.0], [1.0, -1.0, 0.0]]; + for (row_index, row) in rows.iter().enumerate() { + let sum_exp = row.iter().map(|value| value.exp()).sum::(); + assert_close(output_values[row_index][0], sum_exp.ln()); + for (column, value) in row.iter().enumerate() { + assert_close(dinput[row_index][column], value.exp() / sum_exp); + } + } + + assert_gradient_matches_finite_difference( + &device, + [2, 3], + &[0.0, 0.5, 1.0, 1.0, -1.0, 0.0], + |_, x| x.exp().sum_keepdim(1).log().reshape([2]).sum(), + ) + .await; + } +} + +#[tokio::test] +async fn test_backward_with_backwards() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + let y: Tensor<1> = Tensor::new(&graph, &device, &[4.0f32, 5.0, 6.0]); + + let x_target = x.clone(); + let y_target = y.clone(); + let output = x + .add(&y) + .with_backwards([x.parent(), y.parent()], move |grad| { + Ok(vec![ + BackwardTarget::wrt(&x_target, grad.clone().mul_scalar(2.0).to_concrete()), + BackwardTarget::wrt(&y_target, grad.mul_scalar(-3.0).to_concrete()), + ]) + }); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + let dy = gradients + .get(&y) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_eq!(output_values, vec![5.0, 7.0, 9.0]); + // the custom rule replaces add's backward, so the gradients are the + // custom 2x/-3x rather than add's 1/1 + assert_eq!(dx, vec![2.0, 2.0, 2.0]); + assert_eq!(dy, vec![-3.0, -3.0, -3.0]); + } +} + +#[tokio::test] +async fn test_backward_with_backwards_missing_parent_errors() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + let y: Tensor<1> = Tensor::new(&graph, &device, &[4.0f32, 5.0, 6.0]); + + let x_target = x.clone(); + let output = x + .add(&y) + .with_backwards([x.parent(), y.parent()], move |grad| { + Ok(vec![BackwardTarget::wrt(&x_target, grad)]) + }); + + // The scheduler waits on a gradient from every child edge, so a custom + // rule that skips a live parent must fail loudly instead of silently + // dropping every gradient upstream of it. + let Err(error) = output.sum().backward() else { + panic!("backward succeeded despite an omitted parent gradient"); + }; + assert!( + error.to_string().contains("omitted a gradient"), + "expected missing-parent error, got: {error}", + ); + } +} + +#[tokio::test] +async fn test_graph_drops_after_backward() { + for device in test_devices().await { + let graph = Graph::new(); + let weak = Arc::downgrade(&graph.inner); + + let x: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0], [3.0, 4.0]]); + let w: Tensor<2> = Tensor::new(&graph, &device, &[[0.5f32, -1.0], [1.5, 2.0]]); + let loss = x.mat_mul(&w).sum(1).sum(); + let gradients = loss.backward().unwrap(); + assert!(gradients.get(&x).is_some()); + assert!(gradients.get(&w).is_some()); + + drop(gradients); + drop(loss); + drop(x); + drop(w); + drop(graph); + + assert!( + weak.upgrade().is_none(), + "autograd graph stayed alive after all tensors were dropped", + ); + } +} + +#[test] +fn test_gpu_gradients_can_detach() { + let Ok(device) = Device::gpu_blocking() else { + eprintln!("skipping GPU gradient detach regression test: GPU unavailable"); + return; + }; + + let graph = Graph::new(); + let x: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0], [3.0, 4.0]]); + let w: Tensor<2> = Tensor::new(&graph, &device, &[[0.5f32, -1.0], [1.5, 2.0]]); + let gradients = x + .mat_mul(&w) + .sum(1) + .sum() + .backward() + .unwrap() + .into_detached(); + let dx = gradients.get(&x).expect("missing x gradient"); + let dw = gradients.get(&w).expect("missing w gradient"); + + assert_eq!( + dx.as_gpu() + .expect("expected GPU x gradient") + .count_kernels_to_resolve(), + 0, + "detached x gradient should not retain backward compute graph", + ); + assert_eq!( + dw.as_gpu() + .expect("expected GPU w gradient") + .count_kernels_to_resolve(), + 0, + "detached w gradient should not retain backward compute graph", + ); +} + +/// 8x8 XOR grid plus deterministic LCG weight init for a 2-`hidden`-2 MLP, +/// shared by the XOR training tests. +fn xor_training_data(hidden: usize) -> (Vec, Vec, Vec, Vec) { + let mut features = Vec::with_capacity(128); + let mut labels = Vec::with_capacity(64); + for row in 0..8 { + for column in 0..8 { + let x = -0.875 + 0.25 * row as f32; + let y = -0.875 + 0.25 * column as f32; + features.extend([x, y]); + labels.push(u32::from((x > 0.0) != (y > 0.0))); + } + } + + let mut state = 42u64; + let mut next_uniform = move || { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + (state >> 33) as f32 / (1u64 << 32) as f32 - 0.5 + }; + let w1_init: Vec = (0..2 * hidden).map(|_| next_uniform()).collect(); + let w2_init: Vec = (0..hidden * 2).map(|_| next_uniform() * 0.5).collect(); + (features, labels, w1_init, w2_init) +} + +/// End-to-end training: a 2-16-2 MLP learns XOR over an 8x8 grid of 2D +/// points with softmax cross-entropy and full-batch SGD. Exercises the +/// whole tape (matmul, broadcast bias, relu, softmax, gather, log, +/// reduce) plus the detach/update loop across many resolves per device. +#[tokio::test] +async fn test_train_xor_classifier() { + const SAMPLES: usize = 64; + const HIDDEN: usize = 16; + const STEPS: usize = 500; + const LEARNING_RATE: f32 = 1.0; + + let (features, labels, w1_init, w2_init) = xor_training_data(HIDDEN); + + for (device, name) in test_devices().await.into_iter().zip(["cpu", "gpu"]) { + let inputs = RawTensor::from_slice(&device, [SAMPLES, 2], &features); + let targets = RawTensor::from_slice(&device, [SAMPLES], &labels); + + let mut w1 = RawTensor::from_slice(&device, [2, HIDDEN], &w1_init); + let mut b1 = RawTensor::zeros(&device, [HIDDEN]); + let mut w2 = RawTensor::from_slice(&device, [HIDDEN, 2], &w2_init); + let mut b2 = RawTensor::zeros(&device, [2]); + + let mut final_loss = f32::INFINITY; + for step in 0..STEPS { + let graph = Graph::new(); + let x = Tensor::constant_from_raw(&graph, inputs.clone()); + let w1_t = Tensor::from_raw(&graph, w1.clone()); + let b1_t = Tensor::from_raw(&graph, b1.clone()); + let w2_t = Tensor::from_raw(&graph, w2.clone()); + let b2_t = Tensor::from_raw(&graph, b2.clone()); + + let hidden = b1_t.add_::<2, 2>(&x.mat_mul(&w1_t)).relu(); + let logits = b2_t.add_::<2, 2>(&hidden.mat_mul(&w2_t)); + // Numerically stable cross-entropy: log softmax via log-sum-exp + // so a saturated class cannot underflow to log(0). + let shifted = logits.sub_::<2, 2>(&logits.max_keepdim::<1>(1)); + let log_sum_exp = shifted.exp().sum_keepdim(1).log(); + let label_log_probs = shifted.sub_::<2, 2>(&log_sum_exp).gather_last(&targets); + let loss: Tensor<0> = label_log_probs.sum().mul_scalar(-1.0 / SAMPLES as f32); + + let loss_value = flatten(loss.raw().clone()).await[0]; + let gradients = loss.backward().unwrap().into_detached(); + let dw1 = gradients.get(&w1_t).unwrap(); + let db1 = gradients.get(&b1_t).unwrap(); + let dw2 = gradients.get(&w2_t).unwrap(); + let db2 = gradients.get(&b2_t).unwrap(); + + w1 = (w1 - dw1 * LEARNING_RATE).to_concrete(); + b1 = (b1 - db1 * LEARNING_RATE).to_concrete(); + w2 = (w2 - dw2 * LEARNING_RATE).to_concrete(); + b2 = (b2 - db2 * LEARNING_RATE).to_concrete(); + + final_loss = loss_value; + if step % 100 == 0 { + eprintln!("[{name}] step {step}: loss {loss_value:.4}"); + } + } + eprintln!("[{name}] final loss {final_loss:.4}"); + + let graph = Graph::new(); + let x = Tensor::constant_from_raw(&graph, inputs.clone()); + let w1_t = Tensor::constant_from_raw(&graph, w1.clone()); + let b1_t = Tensor::constant_from_raw(&graph, b1.clone()); + let w2_t = Tensor::constant_from_raw(&graph, w2.clone()); + let b2_t = Tensor::constant_from_raw(&graph, b2.clone()); + let hidden = b1_t.add_::<2, 2>(&x.mat_mul(&w1_t)).relu(); + let logits = b2_t.add_::<2, 2>(&hidden.mat_mul(&w2_t)); + let logits = logits.raw().clone().as_slice().await.unwrap().to_vec2(); + let correct = logits + .iter() + .zip(&labels) + .filter(|(row, label)| u32::from(row[1] > row[0]) == **label) + .count(); + eprintln!("[{name}] accuracy {correct}/{SAMPLES}"); + + assert!( + final_loss < 0.1, + "training did not converge: final loss {final_loss}", + ); + assert_eq!(correct, SAMPLES, "classifier misclassified training points"); + } +} + +#[tokio::test] +async fn test_autograd_sigmoid() { + for device in test_devices().await { + let graph = Graph::new(); + let inputs = [-2.0f32, -0.5, 0.0, 1.0, 3.0]; + let x: Tensor<1> = Tensor::new(&graph, &device, &inputs); + + let output = x.sigmoid(); + let values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + + let expected = inputs.map(|v| 1.0 / (1.0 + (-v).exp())); + for (value, expected) in values.iter().zip(expected) { + assert_close(*value, expected); + } + + let gradients = output.sum().backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + let expected_grads = inputs.map(|v| { + let sigmoid = 1.0 / (1.0 + (-v).exp()); + sigmoid * (1.0 - sigmoid) + }); + for (value, expected) in dx.iter().zip(expected_grads) { + assert_close(*value, expected); + } + + assert_gradient_matches_finite_difference(&device, [5], &inputs, |_, x| x.sigmoid().sum()) + .await; + } +} + +#[tokio::test] +async fn test_autograd_to_concrete() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, -2.0, 3.0]); + + let output = x.mul_scalar(2.0).to_concrete(); + let values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + assert_eq!(values, vec![2.0, -4.0, 6.0]); + + let gradients = output.sqr().sum().backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + assert_slice_close(&dx, &[8.0, -16.0, 24.0]); + + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, -2.0, 3.0]); + let gradients = x.to_concrete().sqr().sum().backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + assert_slice_close(&dx, &[2.0, -4.0, 6.0]); + + assert_gradient_matches_finite_difference(&device, [3], &[1.0f32, -2.0, 3.0], |_, x| { + x.to_concrete().sqr().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_index_select_rank_generic() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0, 4.0]); + let indices = RawTensor::from_slice(&device, [3], &[2u32, 0, 2]); + let selected = x.index_select(0, &indices); + assert_eq!( + selected.raw().clone().as_slice().await.unwrap().to_vec1(), + vec![3.0, 1.0, 3.0] + ); + let gradients = selected.sum().backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_eq!(dx, vec![1.0, 0.0, 2.0, 0.0]); + + let graph = Graph::new(); + let data: Vec = (0..12).map(|v| v as f32).collect(); + let x: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 3, 2], &data); + let indices = RawTensor::from_slice(&device, [4], &[2u32, 0, 1, 0]); + let selected = x.index_select(1, &indices); + assert_eq!(selected.shape(), [2, 4, 2]); + assert_eq!( + flatten(selected.raw().clone()).await, + vec![4.0, 5.0, 0.0, 1.0, 2.0, 3.0, 0.0, 1.0, 10.0, 11.0, 6.0, 7.0, 8.0, 9.0, 6.0, 7.0] + ); + let gradients = selected.flatten_all().sum().backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_eq!( + dx, + vec![2.0, 2.0, 1.0, 1.0, 1.0, 1.0, 2.0, 2.0, 1.0, 1.0, 1.0, 1.0] + ); + + assert_gradient_matches_finite_difference( + &device, + [2, 3], + &[0.5f32, -1.0, 2.0, 3.0, -0.5, 1.5], + |graph, x| { + let indices = RawTensor::from_slice(&x.device(), [2], &[2u32, 0]); + let weights = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&x.device(), [2, 2], &[1.0f32, 2.0, 3.0, 4.0]), + ); + x.index_select(1, &indices) + .mul(&weights) + .flatten_all() + .sum() + }, + ) + .await; + + assert_gradient_matches_finite_difference( + &device, + [2, 3], + &[0.5f32, -1.0, 2.0, 3.0, -0.5, 1.5], + |graph, x| { + let indices = RawTensor::from_slice(&x.device(), [3], &[1u32, 1, 0]); + let weights = Tensor::constant_from_raw( + graph, + RawTensor::from_slice( + &x.device(), + [3, 3], + &[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0], + ), + ); + x.index_select(0, &indices) + .mul(&weights) + .flatten_all() + .sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_autograd_i_indexing() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<2> = + Tensor::from_slice(&graph, &device, [2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + let row = x.i((1, ..)); + assert_eq!( + row.raw().clone().as_slice().await.unwrap().to_vec1(), + vec![4.0, 5.0, 6.0] + ); + let gradients = row.sum().backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_eq!(dx, vec![0.0, 0.0, 0.0, 1.0, 1.0, 1.0]); + + let graph = Graph::new(); + let x: Tensor<2> = + Tensor::from_slice(&graph, &device, [2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + let column = x.i((0..2, 1)); + assert_eq!( + column.raw().clone().as_slice().await.unwrap().to_vec1(), + vec![2.0, 5.0] + ); + let gradients = column.sum().backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_eq!(dx, vec![0.0, 1.0, 0.0, 0.0, 1.0, 0.0]); + + let graph = Graph::new(); + let data: Vec = (0..12).map(|v| v as f32).collect(); + let x: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 2, 3], &data); + let plane = x.i((.., 1, ..)); + assert_eq!(plane.shape(), [2, 3]); + assert_eq!( + flatten(plane.raw().clone()).await, + vec![3.0, 4.0, 5.0, 9.0, 10.0, 11.0] + ); + let gradients = plane.flatten_all().sum().backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_eq!( + dx, + vec![0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0] + ); + + let graph = Graph::new(); + let data: Vec = (0..16).map(|v| v as f32).collect(); + let x: Tensor<4> = Tensor::from_slice(&graph, &device, [2, 2, 2, 2], &data); + let cube = x.i((1, .., .., ..)); + assert_eq!(cube.shape(), [2, 2, 2]); + assert_eq!( + flatten(cube.raw().clone()).await, + (8..16).map(|v| v as f32).collect::>() + ); + let gradients = cube.flatten_all().sum().backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + let expected: Vec = (0..16).map(|v| if v < 8 { 0.0 } else { 1.0 }).collect(); + assert_eq!(dx, expected); + + assert_gradient_matches_finite_difference( + &device, + [2, 3], + &[0.5f32, -1.0, 2.0, 3.0, -0.5, 1.5], + |graph, x| { + let weights = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&x.device(), [3], &[1.0f32, 2.0, 3.0]), + ); + x.i((1, ..)).mul(&weights).sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_backward_squeeze() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<3> = + Tensor::from_slice(&graph, &device, [2, 1, 2], &[1.0, 2.0, 3.0, 4.0]); + let output: Tensor<2> = input.squeeze::<2>(1); + assert_eq!(output.shape(), [2, 2]); + let values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + assert_eq!(values, vec![vec![1.0, 2.0], vec![3.0, 4.0]]); + + let gradients = output.sqr().flatten_all().sum().backward().unwrap(); + let dx = flatten(gradients.get(&input).unwrap()).await; + assert_slice_close(&dx, &[2.0, 4.0, 6.0, 8.0]); + + assert_gradient_matches_finite_difference( + &device, + [2, 1, 2], + &[1.0, 2.0, 3.0, 4.0], + |_, x| x.squeeze::<2>(1).sqr().flatten_all().sum(), + ) + .await; + } +} + +#[tokio::test] +async fn test_backward_unsqueeze() { + for device in test_devices().await { + let graph = Graph::new(); + let data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; + let input: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 2, 2], &data); + let output: Tensor<4> = input.unsqueeze::<4>(1); + assert_eq!(output.shape(), [2, 1, 2, 2]); + assert_eq!(flatten(output.raw().clone()).await, data.to_vec()); + + let gradients = output.sqr().flatten_all().sum().backward().unwrap(); + let dx = flatten(gradients.get(&input).unwrap()).await; + assert_slice_close(&dx, &data.map(|value| 2.0 * value)); + + assert_gradient_matches_finite_difference(&device, [2, 2, 2], &data, |_, x| { + x.unsqueeze::<4>(3).sqr().flatten_all().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_cat_rank1() { + for device in test_devices().await { + let graph = Graph::new(); + let first: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0]); + let second: Tensor<1> = Tensor::new(&graph, &device, &[3.0f32, 4.0, 5.0]); + let output = Tensor::cat(vec![first.clone(), second.clone()], 0); + assert_eq!(output.shape(), [5]); + let values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + assert_eq!(values, vec![1.0, 2.0, 3.0, 4.0, 5.0]); + + let gradients = output.sqr().sum().backward().unwrap(); + assert_slice_close(&flatten(gradients.get(&first).unwrap()).await, &[2.0, 4.0]); + assert_slice_close( + &flatten(gradients.get(&second).unwrap()).await, + &[6.0, 8.0, 10.0], + ); + + assert_gradient_matches_finite_difference(&device, [2], &[1.0, 2.0], |graph, x| { + let other = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&x.device(), [3], &[3.0, 4.0, 5.0]), + ); + Tensor::cat(vec![x, other], 0).sqr().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_cat_rank2() { + for device in test_devices().await { + let graph = Graph::new(); + let first: Tensor<2> = + Tensor::from_slice(&graph, &device, [2, 2], &[1.0, 2.0, 3.0, 4.0]); + let second: Tensor<2> = Tensor::from_slice(&graph, &device, [2, 1], &[5.0, 6.0]); + let output = Tensor::cat(vec![first.clone(), second.clone()], 1); + assert_eq!(output.shape(), [2, 3]); + let values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + assert_eq!(values, vec![vec![1.0, 2.0, 5.0], vec![3.0, 4.0, 6.0]]); + + let gradients = output.sqr().flatten_all().sum().backward().unwrap(); + assert_slice_close( + &flatten(gradients.get(&first).unwrap()).await, + &[2.0, 4.0, 6.0, 8.0], + ); + assert_slice_close(&flatten(gradients.get(&second).unwrap()).await, &[10.0, 12.0]); + + assert_gradient_matches_finite_difference(&device, [2, 1], &[5.0, 6.0], |graph, x| { + let other = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&x.device(), [2, 2], &[1.0, 2.0, 3.0, 4.0]), + ); + Tensor::cat(vec![other, x], 1).sqr().flatten_all().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_pad_with_zeros() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + let output = input.pad_with_zeros(0, 1, 2); + assert_eq!(output.shape(), [6]); + let values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + assert_eq!(values, vec![0.0, 1.0, 2.0, 3.0, 0.0, 0.0]); + + let gradients = output.sqr().sum().backward().unwrap(); + assert_slice_close( + &flatten(gradients.get(&input).unwrap()).await, + &[2.0, 4.0, 6.0], + ); + + let passthrough = input.pad_with_zeros(0, 0, 0); + let gradients = passthrough.sqr().sum().backward().unwrap(); + assert_slice_close( + &flatten(gradients.get(&input).unwrap()).await, + &[2.0, 4.0, 6.0], + ); + + assert_gradient_matches_finite_difference(&device, [3], &[1.0, 2.0, 3.0], |_, x| { + x.pad_with_zeros(0, 1, 2).sqr().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_pad_axis() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = + Tensor::from_slice(&graph, &device, [2, 2], &[1.0, 2.0, 3.0, 4.0]); + let output = input.pad_axis(1, 1); + assert_eq!(output.shape(), [2, 4]); + let values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + assert_eq!( + values, + vec![vec![0.0, 1.0, 2.0, 0.0], vec![0.0, 3.0, 4.0, 0.0]] + ); + + let gradients = output.sqr().flatten_all().sum().backward().unwrap(); + assert_slice_close( + &flatten(gradients.get(&input).unwrap()).await, + &[2.0, 4.0, 6.0, 8.0], + ); + + assert_gradient_matches_finite_difference( + &device, + [2, 2], + &[1.0, 2.0, 3.0, 4.0], + |_, x| x.pad_axis(0, 2).sqr().flatten_all().sum(), + ) + .await; + } +} + +#[tokio::test] +async fn test_backward_sliding_window_view() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0, 4.0, 5.0]); + let output: Tensor<2> = + input.sliding_window_view::<1, 2>([fusor_types::SlidingWindow::new(0, 3, 1)]); + assert_eq!(output.shape(), [3, 3]); + let values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + assert_eq!( + values, + vec![ + vec![1.0, 2.0, 3.0], + vec![2.0, 3.0, 4.0], + vec![3.0, 4.0, 5.0] + ] + ); + + let gradients = output.flatten_all().sum().backward().unwrap(); + assert_slice_close( + &flatten(gradients.get(&input).unwrap()).await, + &[1.0, 2.0, 3.0, 2.0, 1.0], + ); + + assert_gradient_matches_finite_difference( + &device, + [5], + &[1.0, 2.0, 3.0, 4.0, 5.0], + |_, x| { + x.sliding_window_view::<1, 2>([fusor_types::SlidingWindow::new(0, 3, 1)]) + .sqr() + .flatten_all() + .sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_backward_sliding_window_view_strided() { + for device in test_devices().await { + let graph = Graph::new(); + let data: Vec = (1..=10).map(|value| value as f32).collect(); + let input: Tensor<2> = Tensor::from_slice(&graph, &device, [2, 5], &data); + let output: Tensor<3> = + input.sliding_window_view::<1, 3>([fusor_types::SlidingWindow::new(1, 3, 2)]); + assert_eq!(output.shape(), [2, 2, 3]); + assert_eq!( + flatten(output.raw().clone()).await, + vec![1.0, 2.0, 3.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 8.0, 9.0, 10.0] + ); + + let gradients = output.flatten_all().sum().backward().unwrap(); + assert_slice_close( + &flatten(gradients.get(&input).unwrap()).await, + &[1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 2.0, 1.0, 1.0], + ); + + assert_gradient_matches_finite_difference(&device, [2, 5], &data, |_, x| { + x.sliding_window_view::<1, 3>([fusor_types::SlidingWindow::new(1, 3, 2)]) + .sqr() + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_sum_axis() { + for device in test_devices().await { + let graph = Graph::new(); + + let x: Tensor<2> = + Tensor::from_slice(&graph, &device, [2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + let summed = x.sum(1); + let forward = summed.raw().clone().as_slice().await.unwrap().to_vec1(); + assert_slice_close(&forward, &[6.0, 15.0]); + + let weight = + Tensor::constant_from_raw(&graph, RawTensor::from_slice(&device, [2], &[1.0, 2.0])); + let loss = summed.mul(&weight).sum(); + let gradients = loss.backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_slice_close(&dx, &[1.0, 1.0, 1.0, 2.0, 2.0, 2.0]); + + let graph = Graph::new(); + let x: Tensor<2> = + Tensor::from_slice(&graph, &device, [2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + let summed = x.sum(0); + let forward = summed.raw().clone().as_slice().await.unwrap().to_vec1(); + assert_slice_close(&forward, &[5.0, 7.0, 9.0]); + + let weight = Tensor::constant_from_raw( + &graph, + RawTensor::from_slice(&device, [3], &[1.0, 2.0, 3.0]), + ); + let loss = summed.mul(&weight).sum(); + let gradients = loss.backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_slice_close(&dx, &[1.0, 2.0, 3.0, 1.0, 2.0, 3.0]); + + let data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; + assert_gradient_matches_finite_difference(&device, [2, 3], &data, |_, x| { + x.sum(1).sqr().sum() + }) + .await; + assert_gradient_matches_finite_difference(&device, [2, 3], &data, |_, x| { + x.sum(0).sqr().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_sum_high_rank() { + for device in test_devices().await { + let graph = Graph::new(); + + let data: Vec = (1..=8).map(|v| v as f32).collect(); + let x: Tensor<5> = Tensor::from_slice(&graph, &device, [2, 1, 2, 1, 2], &data); + let summed = x.sum(2); + assert_eq!(summed.shape(), [2, 1, 1, 2]); + let forward = flatten(summed.raw().clone()).await; + assert_slice_close(&forward, &[4.0, 6.0, 12.0, 14.0]); + + let loss = summed.sqr().sum(3).sum(2).sum(1).sum(); + let gradients = loss.backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_slice_close(&dx, &[8.0, 12.0, 8.0, 12.0, 24.0, 28.0, 24.0, 28.0]); + + assert_gradient_matches_finite_difference(&device, [2, 1, 2, 1, 2], &data, |_, x| { + x.sum(2).sqr().sum(3).sum(2).sum(1).sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_sum_keepdim() { + for device in test_devices().await { + let graph = Graph::new(); + + let x: Tensor<2> = + Tensor::from_slice(&graph, &device, [2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + let summed = x.sum_keepdim(1); + assert_eq!(summed.shape(), [2, 1]); + let forward = flatten(summed.raw().clone()).await; + assert_slice_close(&forward, &[6.0, 15.0]); + + let weight = Tensor::constant_from_raw( + &graph, + RawTensor::from_slice(&device, [2, 1], &[1.0, 2.0]), + ); + let loss = summed.mul(&weight).flatten_all().sum(); + let gradients = loss.backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_slice_close(&dx, &[1.0, 1.0, 1.0, 2.0, 2.0, 2.0]); + + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + let summed = x.sum_keepdim(0); + assert_eq!(summed.shape(), [1]); + let forward = summed.raw().clone().as_slice().await.unwrap().to_vec1(); + assert_slice_close(&forward, &[6.0]); + let gradients = summed.sqr().sum().backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + assert_slice_close(&dx, &[12.0, 12.0, 12.0]); + + let graph = Graph::new(); + let data: Vec = (1..=8).map(|v| v as f32).collect(); + let x: Tensor<5> = Tensor::from_slice(&graph, &device, [2, 1, 2, 1, 2], &data); + let summed = x.sum_keepdim(2); + assert_eq!(summed.shape(), [2, 1, 1, 1, 2]); + let forward = flatten(summed.raw().clone()).await; + assert_slice_close(&forward, &[4.0, 6.0, 12.0, 14.0]); + let gradients = summed.flatten_all().sum().backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_slice_close(&dx, &[1.0; 8]); + + let data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; + assert_gradient_matches_finite_difference(&device, [2, 3], &data, |_, x| { + x.sum_keepdim(0).sqr().flatten_all().sum() + }) + .await; + assert_gradient_matches_finite_difference(&device, [3], &data[..3], |_, x| { + x.sum_keepdim(0).sqr().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_q_mat_mul_rank1() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0, 4.0]); + let weight_bytes: Vec = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0] + .into_iter() + .flat_map(|value| value.to_le_bytes()) + .collect(); + let weights = crate::QMatrix::from_raw_bytes( + &device, + [2, 4], + &weight_bytes, + fusor_gguf::GgmlType::F32, + ) + .unwrap(); + + let output = input.q_mat_mul(&weights); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec1(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + + assert_slice_close(&output_values, &[30.0, 70.0]); + assert_slice_close(&dinput, &[6.0, 8.0, 10.0, 12.0]); + + assert_gradient_matches_finite_difference(&device, [4], &[0.5, -1.0, 2.0, 0.25], |_, x| { + x.q_mat_mul(&weights).sum() + }) + .await; + } +} + +fn composite_ramp(graph: &Graph, device: &Device, shape: [usize; R]) -> Tensor { + let elements: usize = shape.iter().product(); + Tensor::constant_from_raw( + graph, + crate::arange(device, 1.0, elements as f32 + 1.0) + .reshape(shape) + .to_concrete(), + ) +} + +#[tokio::test] +async fn test_autograd_conv1d() { + for device in test_devices().await { + let x_data: Vec = (0..8).map(|i| (i as f32 * 0.7).sin()).collect(); + let w_data: Vec = (0..8).map(|i| (i as f32 * 0.3).cos()).collect(); + let b_data = [0.5f32, -1.0]; + + let graph = Graph::new(); + let x: Tensor<3> = Tensor::from_slice(&graph, &device, [1, 2, 4], &x_data); + let w: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 2, 2], &w_data); + let b: Tensor<1> = Tensor::from_slice(&graph, &device, [2], &b_data); + let output = x.conv(&w, Some(&b), [1], [1]); + + let raw_x = RawTensor::from_slice(&device, [1, 2, 4], &x_data); + let raw_w = RawTensor::from_slice(&device, [2, 2, 2], &w_data); + let raw_b = RawTensor::from_slice(&device, [2], &b_data); + let expected = raw_x.conv(&raw_w, Some(&raw_b), [1], [1]); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + + let fd_device = device.clone(); + let w_fd = w_data.clone(); + assert_gradient_matches_finite_difference( + &device, + [1, 2, 4], + &x_data, + move |graph, x| { + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 2], &w_fd), + ); + let b = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2], &[0.5f32, -1.0]), + ); + let out = x.conv(&w, Some(&b), [1], [1]); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }, + ) + .await; + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 2, 2], + &w_data, + move |graph, w| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 2, 4], &x_fd), + ); + let out = x.conv(&w, None, [1], [1]); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }, + ) + .await; + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + let w_fd = w_data.clone(); + assert_gradient_matches_finite_difference(&device, [2], &b_data, move |graph, b| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 2, 4], &x_fd), + ); + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 2], &w_fd), + ); + let out = x.conv(&w, Some(&b), [1], [1]); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_conv2d_strided() { + for device in test_devices().await { + let x_data: Vec = (0..18).map(|i| (i as f32 * 0.41).sin()).collect(); + let w_data: Vec = (0..16).map(|i| (i as f32 * 0.23).cos()).collect(); + + let graph = Graph::new(); + let x: Tensor<4> = Tensor::from_slice(&graph, &device, [1, 2, 3, 3], &x_data); + let w: Tensor<4> = Tensor::from_slice(&graph, &device, [2, 2, 2, 2], &w_data); + let output = x.conv(&w, None, [1, 1], [2, 2]); + + let raw_x = RawTensor::from_slice(&device, [1, 2, 3, 3], &x_data); + let raw_w = RawTensor::from_slice(&device, [2, 2, 2, 2], &w_data); + let expected = raw_x.conv(&raw_w, None, [1, 1], [2, 2]); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + + let fd_device = device.clone(); + let w_fd = w_data.clone(); + assert_gradient_matches_finite_difference( + &device, + [1, 2, 3, 3], + &x_data, + move |graph, x| { + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 2, 2], &w_fd), + ); + let out = x.conv(&w, None, [1, 1], [2, 2]); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }, + ) + .await; + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 2, 2, 2], + &w_data, + move |graph, w| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 2, 3, 3], &x_fd), + ); + let out = x.conv(&w, None, [1, 1], [2, 2]); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_autograd_grouped_conv() { + for device in test_devices().await { + let x_data: Vec = (0..16).map(|i| (i as f32 * 0.57).sin()).collect(); + let w_data: Vec = (0..16).map(|i| (i as f32 * 0.31).cos()).collect(); + let b_data = [0.2f32, -0.4, 0.6, -0.8]; + + let graph = Graph::new(); + let x: Tensor<3> = Tensor::from_slice(&graph, &device, [1, 4, 4], &x_data); + let w: Tensor<3> = Tensor::from_slice(&graph, &device, [4, 2, 2], &w_data); + let b: Tensor<1> = Tensor::from_slice(&graph, &device, [4], &b_data); + let output = x.grouped_conv(&w, Some(&b), [1], [2], 2); + + let raw_x = RawTensor::from_slice(&device, [1, 4, 4], &x_data); + let raw_w = RawTensor::from_slice(&device, [4, 2, 2], &w_data); + let raw_b = RawTensor::from_slice(&device, [4], &b_data); + let expected = raw_x.grouped_conv(&raw_w, Some(&raw_b), [1], [2], 2); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + + let fd_device = device.clone(); + let w_fd = w_data.clone(); + assert_gradient_matches_finite_difference( + &device, + [1, 4, 4], + &x_data, + move |graph, x| { + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [4, 2, 2], &w_fd), + ); + let b = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [4], &[0.2f32, -0.4, 0.6, -0.8]), + ); + let out = x.grouped_conv(&w, Some(&b), [1], [2], 2); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }, + ) + .await; + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + assert_gradient_matches_finite_difference( + &device, + [4, 2, 2], + &w_data, + move |graph, w| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 4, 4], &x_fd), + ); + let out = x.grouped_conv(&w, None, [1], [2], 2); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_autograd_upsample_nearest2d() { + for device in test_devices().await { + let data: Vec = (0..12).map(|i| i as f32 * 0.5).collect(); + let graph = Graph::new(); + let x: Tensor<4> = Tensor::from_slice(&graph, &device, [1, 2, 2, 3], &data); + let output = x.upsample_nearest2d(2, 3); + assert_eq!(output.shape(), [1, 2, 4, 9]); + + let raw_x = RawTensor::from_slice(&device, [1, 2, 2, 3], &data); + let expected = raw_x.upsample_nearest2d(2, 3); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + + let gradients = output.flatten_all().sum().backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + for value in dx { + assert_close(value, 6.0); + } + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference( + &device, + [1, 2, 2, 3], + &data, + move |graph, x| { + let out = x.upsample_nearest2d(2, 3); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_autograd_softmax_slow() { + for device in test_devices().await { + let data = [1.0f32, 2.0, 3.0, -1.0, 0.5, 0.0]; + let graph = Graph::new(); + let x: Tensor<2> = Tensor::from_slice(&graph, &device, [2, 3], &data); + let output = x.softmax_slow(1); + + let raw_x = RawTensor::from_slice(&device, [2, 3], &data); + let expected = raw_x.softmax_slow::<1>(1); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + + let last = x.softmax_slow_last_dim(); + assert_slice_close( + &flatten(last.raw().clone()).await, + &flatten(output.raw().clone()).await, + ); + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [2, 3], &data, move |graph, x| { + x.softmax_slow(1) + .mul(&composite_ramp(graph, &fd_device, [2, 3])) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_norm_rms_norm_rank4() { + for device in test_devices().await { + let data: Vec = (0..12).map(|i| (i as f32 * 0.83).sin() * 2.0).collect(); + let w_data = [0.5f32, 1.5, 2.0]; + let b_data = [0.1f32, -0.2, 0.3]; + let eps = 1e-5; + + let graph = Graph::new(); + let x: Tensor<4> = Tensor::from_slice(&graph, &device, [1, 2, 2, 3], &data); + let w: Tensor<4> = Tensor::from_slice(&graph, &device, [1, 1, 1, 3], &w_data); + let b: Tensor<4> = Tensor::from_slice(&graph, &device, [1, 1, 1, 3], &b_data); + let output = x.layer_norm(&w, Some(&b), eps, true); + + let raw_x = RawTensor::from_slice(&device, [1, 2, 2, 3], &data); + let raw_w = RawTensor::from_slice(&device, [3], &w_data); + let raw_b = RawTensor::from_slice(&device, [3], &b_data); + let expected = raw_x.layer_norm_last_dim_fused::<3, 1, _, _>(&raw_w, Some(&raw_b), eps); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference( + &device, + [1, 2, 2, 3], + &data, + move |graph, x| { + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 1, 1, 3], &[0.5f32, 1.5, 2.0]), + ); + let b = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 1, 1, 3], &[0.1f32, -0.2, 0.3]), + ); + let out = x.layer_norm(&w, Some(&b), eps, true); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }, + ) + .await; + + let rms = x.rms_norm(&w, eps); + let expected_rms = RawTensor::from_slice(&device, [1, 2, 2, 3], &data) + .rms_norm_fused::<1, 3>(&raw_w, None, eps); + assert_slice_close( + &flatten(rms.raw().clone()).await, + &flatten(expected_rms).await, + ); + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference( + &device, + [1, 2, 2, 3], + &data, + move |graph, x| { + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 1, 1, 3], &[0.5f32, 1.5, 2.0]), + ); + let out = x.rms_norm(&w, eps); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_autograd_rms_norm_residual_fused() { + for device in test_devices().await { + let x_data: Vec = (0..12).map(|i| (i as f32 * 0.29).sin()).collect(); + let r_data: Vec = (0..12).map(|i| (i as f32 * 0.61).cos()).collect(); + let w_data = [0.5f32, 1.5, 2.0]; + let b_data = [0.1f32, -0.2, 0.3]; + let eps = 1e-5; + + let graph = Graph::new(); + let x: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 2, 3], &x_data); + let r: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 2, 3], &r_data); + let w: Tensor<1> = Tensor::from_slice(&graph, &device, [3], &w_data); + let b: Tensor<1> = Tensor::from_slice(&graph, &device, [3], &b_data); + + let raw_x = RawTensor::from_slice(&device, [2, 2, 3], &x_data); + let raw_r = RawTensor::from_slice(&device, [2, 2, 3], &r_data); + let raw_w = RawTensor::from_slice(&device, [3], &w_data); + let raw_b = RawTensor::from_slice(&device, [3], &b_data); + + let output = x.rms_norm_residual_fused(&r, &w, Some(&b), eps); + let expected = raw_x.rms_norm_residual_fused::<1, 2, _>(&raw_r, &raw_w, Some(&raw_b), eps); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + + let no_bias = x.rms_norm_residual_fused(&r, &w, None, eps); + let expected_no_bias = RawTensor::from_slice(&device, [2, 2, 3], &x_data) + .rms_norm_residual_fused::<1, 2, _>(&raw_r, &raw_w, None, eps); + assert_slice_close( + &flatten(no_bias.raw().clone()).await, + &flatten(expected_no_bias).await, + ); + + let gradients = output.flatten_all().sum().backward().unwrap(); + assert!(gradients.get(&x).is_some()); + assert!(gradients.get(&r).is_some()); + assert!(gradients.get(&w).is_some()); + assert!(gradients.get(&b).is_some()); + + let fd_device = device.clone(); + let r_fd = r_data.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 2, 3], + &x_data, + move |graph, x| { + let r = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 3], &r_fd), + ); + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3], &[0.5f32, 1.5, 2.0]), + ); + let b = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3], &[0.1f32, -0.2, 0.3]), + ); + let out = x.rms_norm_residual_fused(&r, &w, Some(&b), eps); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }, + ) + .await; + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 2, 3], + &r_data, + move |graph, r| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 3], &x_fd), + ); + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3], &[0.5f32, 1.5, 2.0]), + ); + let out = x.rms_norm_residual_fused(&r, &w, None, eps); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }, + ) + .await; + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + let r_fd = r_data.clone(); + assert_gradient_matches_finite_difference(&device, [3], &w_data, move |graph, w| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 3], &x_fd), + ); + let r = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 3], &r_fd), + ); + let b = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3], &[0.1f32, -0.2, 0.3]), + ); + let out = x.rms_norm_residual_fused(&r, &w, Some(&b), eps); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + let r_fd = r_data.clone(); + assert_gradient_matches_finite_difference(&device, [3], &b_data, move |graph, b| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 3], &x_fd), + ); + let r = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 3], &r_fd), + ); + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3], &[0.5f32, 1.5, 2.0]), + ); + let out = x.rms_norm_residual_fused(&r, &w, Some(&b), eps); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_rope_pair_fused() { + for device in test_devices().await { + let q_data: Vec = (0..24).map(|i| (i as f32 * 0.37).sin()).collect(); + let k_data: Vec = (0..12).map(|i| (i as f32 * 0.53).cos()).collect(); + let cos_data: Vec = (0..6).map(|i| (i as f32 * 0.7).cos()).collect(); + let sin_data: Vec = (0..6).map(|i| (i as f32 * 0.7).sin()).collect(); + + let graph = Graph::new(); + let q: Tensor<4> = Tensor::from_slice(&graph, &device, [1, 2, 3, 4], &q_data); + let k: Tensor<4> = Tensor::from_slice(&graph, &device, [1, 1, 3, 4], &k_data); + let cos: Tensor<2> = + Tensor::constant_from_raw(&graph, RawTensor::from_slice(&device, [3, 2], &cos_data)); + let sin: Tensor<2> = + Tensor::constant_from_raw(&graph, RawTensor::from_slice(&device, [3, 2], &sin_data)); + let (q_out, k_out) = q.rope_pair_fused(&k, &cos, &sin); + + let raw_q = RawTensor::from_slice(&device, [1, 2, 3, 4], &q_data); + let raw_k = RawTensor::from_slice(&device, [1, 1, 3, 4], &k_data); + let raw_cos = RawTensor::from_slice(&device, [3, 2], &cos_data); + let raw_sin = RawTensor::from_slice(&device, [3, 2], &sin_data); + let (expected_q, expected_k) = raw_q.rope_pair_fused(&raw_k, &raw_cos, &raw_sin); + assert_slice_close( + &flatten(q_out.raw().clone()).await, + &flatten(expected_q).await, + ); + assert_slice_close( + &flatten(k_out.raw().clone()).await, + &flatten(expected_k).await, + ); + + let (normal_q, normal_k) = q.rope_normal_pair_fused(&k, &cos, &sin); + let (expected_nq, expected_nk) = RawTensor::from_slice(&device, [1, 2, 3, 4], &q_data) + .rope_normal_pair_fused(&raw_k, &raw_cos, &raw_sin); + assert_slice_close( + &flatten(normal_q.raw().clone()).await, + &flatten(expected_nq).await, + ); + assert_slice_close( + &flatten(normal_k.raw().clone()).await, + &flatten(expected_nk).await, + ); + + let fd_device = device.clone(); + let k_fd = k_data.clone(); + let cos_fd = cos_data.clone(); + let sin_fd = sin_data.clone(); + assert_gradient_matches_finite_difference( + &device, + [1, 2, 3, 4], + &q_data, + move |graph, q| { + let k = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 1, 3, 4], &k_fd), + ); + let cos = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3, 2], &cos_fd), + ); + let sin = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3, 2], &sin_fd), + ); + let (q_out, k_out) = q.rope_pair_fused(&k, &cos, &sin); + let q_loss = q_out + .mul(&composite_ramp(graph, &fd_device, q_out.shape())) + .flatten_all() + .sum(); + q_loss.add(&k_out.flatten_all().sum()) + }, + ) + .await; + + let fd_device = device.clone(); + let q_fd = q_data.clone(); + let cos_fd = cos_data.clone(); + let sin_fd = sin_data.clone(); + assert_gradient_matches_finite_difference( + &device, + [1, 1, 3, 4], + &k_data, + move |graph, k| { + let q = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 2, 3, 4], &q_fd), + ); + let cos = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3, 2], &cos_fd), + ); + let sin = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3, 2], &sin_fd), + ); + let (_, k_out) = q.rope_normal_pair_fused(&k, &cos, &sin); + k_out + .mul(&composite_ramp(graph, &fd_device, k_out.shape())) + .flatten_all() + .sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_autograd_rope_cache_forward() { + for device in test_devices().await { + let q_data: Vec = (0..24).map(|i| (i as f32 * 0.37).sin()).collect(); + let k_data: Vec = (0..24).map(|i| (i as f32 * 0.53).cos()).collect(); + let cache = crate::RopeCache::new(4, 8, 10000.0, &device).unwrap(); + + let graph = Graph::new(); + let q: Tensor<4> = Tensor::from_slice(&graph, &device, [1, 2, 3, 4], &q_data); + let k: Tensor<4> = Tensor::from_slice(&graph, &device, [1, 2, 3, 4], &k_data); + + let raw_q = RawTensor::from_slice(&device, [1, 2, 3, 4], &q_data); + let raw_k = RawTensor::from_slice(&device, [1, 2, 3, 4], &k_data); + + let (q_out, k_out) = q.rope_cache_forward(&k, &cache, 2); + let (expected_q, expected_k) = cache.forward(&raw_q, &raw_k, 2); + assert_slice_close( + &flatten(q_out.raw().clone()).await, + &flatten(expected_q).await, + ); + assert_slice_close( + &flatten(k_out.raw().clone()).await, + &flatten(expected_k).await, + ); + + let (qi_out, ki_out) = q.rope_cache_forward_interleaved(&k, &cache, 2); + let (expected_qi, expected_ki) = cache.forward_interleaved( + &RawTensor::from_slice(&device, [1, 2, 3, 4], &q_data), + &RawTensor::from_slice(&device, [1, 2, 3, 4], &k_data), + 2, + ); + assert_slice_close( + &flatten(qi_out.raw().clone()).await, + &flatten(expected_qi).await, + ); + assert_slice_close( + &flatten(ki_out.raw().clone()).await, + &flatten(expected_ki).await, + ); + + let fd_device = device.clone(); + let k_fd = k_data.clone(); + let fd_cache = cache.clone(); + assert_gradient_matches_finite_difference( + &device, + [1, 2, 3, 4], + &q_data, + move |graph, q| { + let k = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 2, 3, 4], &k_fd), + ); + let (q_out, _) = q.rope_cache_forward(&k, &fd_cache, 2); + q_out + .mul(&composite_ramp(graph, &fd_device, q_out.shape())) + .flatten_all() + .sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_autograd_ones() { + for device in test_devices().await { + let graph = Graph::new(); + + let ones: Tensor<1> = Tensor::ones(&graph, &device, [3]); + let forward = ones.raw().clone().as_slice().await.unwrap().to_vec1(); + assert_eq!(forward, vec![1.0, 1.0, 1.0]); + + let x: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32, -3.0, 4.0]); + let loss = x.mul(&ones).sum(); + assert_close(loss.raw().to_scalar().await.unwrap(), 3.0); + let gradients = loss.backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + assert_slice_close(&dx, &[1.0, 1.0, 1.0]); + let dones = gradients + .get(&ones) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + assert_slice_close(&dones, &[2.0, -3.0, 4.0]); + } +} + +#[tokio::test] +async fn test_autograd_ones_like() { + for device in test_devices().await { + let graph = Graph::new(); + + let x: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0], [3.0, 4.0]]); + let ones = x.ones_like(); + assert_eq!(ones.shape(), [2, 2]); + assert_eq!(flatten(ones.raw().clone()).await, vec![1.0, 1.0, 1.0, 1.0]); + + let loss = x.mul(&ones).flatten_all().sum(); + assert_close(loss.raw().to_scalar().await.unwrap(), 10.0); + let gradients = loss.backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_slice_close(&dx, &[1.0, 1.0, 1.0, 1.0]); + let dones = flatten(gradients.get(&ones).unwrap()).await; + assert_slice_close(&dones, &[1.0, 2.0, 3.0, 4.0]); + } +} + +#[tokio::test] +async fn test_backward_mat_mul_rank4() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs_data = (1..=24).map(|n| n as f32).collect::>(); + let rhs_data = (1..=24).map(|n| n as f32).collect::>(); + let lhs: Tensor<4> = Tensor::from_slice(&graph, &device, [2, 2, 2, 3], &lhs_data); + let rhs: Tensor<4> = Tensor::from_slice(&graph, &device, [2, 2, 3, 2], &rhs_data); + + let output = lhs.mat_mul(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dlhs = flatten(gradients.get(&lhs).unwrap()).await; + let drhs = flatten(gradients.get(&rhs).unwrap()).await; + + assert_eq!(output_values.shape(), &[2, 2, 2, 2]); + assert_close(output_values[[0, 0, 0, 0]], 22.0); + assert_close(output_values[[1, 1, 1, 1]], 1522.0); + + // with an all-ones seed, dlhs[b, i, k] = sum_j rhs[b, k, j] and + // drhs[b, k, j] = sum_i lhs[b, i, k] + for batch in 0..4 { + for i in 0..2 { + for k in 0..3 { + let expected = (0..2) + .map(|j| rhs_data[batch * 6 + k * 2 + j]) + .sum::(); + assert_close(dlhs[batch * 6 + i * 3 + k], expected); + } + } + for k in 0..3 { + for j in 0..2 { + let expected = (0..2) + .map(|i| lhs_data[batch * 6 + i * 3 + k]) + .sum::(); + assert_close(drhs[batch * 6 + k * 2 + j], expected); + } + } + } + + let lhs_small = lhs_data + .iter() + .map(|value| value * 0.05) + .collect::>(); + let rhs_small = rhs_data + .iter() + .map(|value| value * 0.03) + .collect::>(); + let fd_device = device.clone(); + let fd_rhs = rhs_small.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 2, 2, 3], + &lhs_small, + move |graph, lhs| { + let rhs = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 3, 2], &fd_rhs), + ); + lhs.mat_mul(&rhs).sqr().flatten_all().sum() + }, + ) + .await; + let fd_device = device.clone(); + let fd_lhs = lhs_small.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 2, 3, 2], + &rhs_small, + move |graph, rhs| { + let lhs = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 2, 3], &fd_lhs), + ); + lhs.mat_mul(&rhs).sqr().flatten_all().sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_autograd_std_ops_add_sub() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + let y: Tensor<1> = Tensor::new(&graph, &device, &[4.0f32, -5.0, 8.0]); + + for add in [ + &x + &y, + x.clone() + y.clone(), + &x + y.clone(), + x.clone() + &y, + ] { + let values = add.raw().clone().as_slice().await.unwrap().to_vec1(); + assert_slice_close(&values, &[5.0, -3.0, 11.0]); + } + for sub in [ + &x - &y, + x.clone() - y.clone(), + &x - y.clone(), + x.clone() - &y, + ] { + let values = sub.raw().clone().as_slice().await.unwrap().to_vec1(); + assert_slice_close(&values, &[-3.0, 7.0, -5.0]); + } + + let loss = ((&x + &y) * (&x - &y)).sum(); + let gradients = loss.backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + let dy = gradients + .get(&y) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + assert_slice_close(&dx, &[2.0, 4.0, 6.0]); + assert_slice_close(&dy, &[-8.0, 10.0, -16.0]); + + assert_gradient_matches_finite_difference(&device, [3], &[1.0, 2.0, 3.0], |graph, x| { + let y = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&x.device(), [3], &[4.0, -5.0, 8.0]), + ); + ((&x + &y) * (&x - &y)).sum() + }) + .await; + assert_gradient_matches_finite_difference(&device, [3], &[4.0, -5.0, 8.0], |graph, y| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&y.device(), [3], &[1.0, 2.0, 3.0]), + ); + ((&x + &y) * (&x - &y)).sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_std_ops_mul_div() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + let y: Tensor<1> = Tensor::new(&graph, &device, &[4.0f32, -5.0, 8.0]); + + for mul in [ + &x * &y, + x.clone() * y.clone(), + &x * y.clone(), + x.clone() * &y, + ] { + let values = mul.raw().clone().as_slice().await.unwrap().to_vec1(); + assert_slice_close(&values, &[4.0, -10.0, 24.0]); + } + for div in [ + &x / &y, + x.clone() / y.clone(), + &x / y.clone(), + x.clone() / &y, + ] { + let values = div.raw().clone().as_slice().await.unwrap().to_vec1(); + assert_slice_close(&values, &[0.25, -0.4, 0.375]); + } + + let loss = ((&x * &y) + (&x / &y)).sum(); + let gradients = loss.backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + let dy = gradients + .get(&y) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + assert_slice_close(&dx, &[4.25, -5.2, 8.125]); + assert_slice_close(&dy, &[0.9375, 1.92, 2.953125]); + + assert_gradient_matches_finite_difference(&device, [3], &[1.0, 2.0, 3.0], |graph, x| { + let y = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&x.device(), [3], &[4.0, -5.0, 8.0]), + ); + ((&x * &y) + (&x / &y)).sum() + }) + .await; + assert_gradient_matches_finite_difference(&device, [3], &[4.0, -5.0, 8.0], |graph, y| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&y.device(), [3], &[1.0, 2.0, 3.0]), + ); + ((&x * &y) + (&x / &y)).sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_std_ops_neg() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, -2.0, 3.0]); + + for neg in [-&x, -x.clone()] { + let values = neg.raw().clone().as_slice().await.unwrap().to_vec1(); + assert_slice_close(&values, &[-1.0, 2.0, -3.0]); + } + + let loss = ((-&x) * &x).sum(); + let gradients = loss.backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + assert_slice_close(&dx, &[-2.0, 4.0, -6.0]); + + assert_gradient_matches_finite_difference(&device, [3], &[1.0, -2.0, 3.0], |_, x| { + ((-&x) * &x).sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_std_ops_scalar() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + + for mul in [&x * 2.5, x.clone() * 2.5] { + let values = mul.raw().clone().as_slice().await.unwrap().to_vec1(); + assert_slice_close(&values, &[2.5, 5.0, 7.5]); + } + for add in [&x + 1.5, x.clone() + 1.5] { + let values = add.raw().clone().as_slice().await.unwrap().to_vec1(); + assert_slice_close(&values, &[2.5, 3.5, 4.5]); + } + for sub in [&x - 0.5, x.clone() - 0.5] { + let values = sub.raw().clone().as_slice().await.unwrap().to_vec1(); + assert_slice_close(&values, &[0.5, 1.5, 2.5]); + } + for div in [&x / 2.0, x.clone() / 2.0] { + let values = div.raw().clone().as_slice().await.unwrap().to_vec1(); + assert_slice_close(&values, &[0.5, 1.0, 1.5]); + } + + let loss = ((&x * 3.0) + 2.0).sum(); + let gradients = loss.backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec1(); + assert_slice_close(&dx, &[3.0, 3.0, 3.0]); + + assert_gradient_matches_finite_difference(&device, [3], &[1.0, 2.0, 3.0], |_, x| { + ((&x * 3.0) + 2.0).sum() + }) + .await; + + assert_gradient_matches_finite_difference(&device, [3], &[1.0, 2.0, 3.0], |_, x| { + ((&x - 1.5) / 4.0).sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_norm_rank_generic() { + for device in test_devices().await { + let eps = 1e-5f32; + let x_rows = [[1.0f32, 2.0, 4.0], [-1.0, 0.5, 3.0]]; + let x_data = [1.0f32, 2.0, 4.0, -1.0, 0.5, 3.0]; + let w_data = [0.5f32, 1.0, 1.5]; + let b_data = [0.1f32, -0.2, 0.3]; + for remove_mean in [true, false] { + let graph = Graph::new(); + let x: Tensor<2> = Tensor::from_slice(&graph, &device, [2, 3], &x_data); + let w: Tensor<2> = Tensor::from_slice(&graph, &device, [1, 3], &w_data); + let b: Tensor<2> = Tensor::from_slice(&graph, &device, [1, 3], &b_data); + + let output = x.layer_norm(&w, Some(&b), eps, remove_mean); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + for (row_index, row) in x_rows.iter().enumerate() { + let mean = if remove_mean { + row.iter().sum::() / 3.0 + } else { + 0.0 + }; + let var = row.iter().map(|v| (v - mean) * (v - mean)).sum::() / 3.0; + let std = (var + eps).sqrt(); + for column in 0..3 { + let expected = (row[column] - mean) / std * w_data[column] + b_data[column]; + assert_close(output_values[row_index][column], expected); + } + } + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 3], + &x_data, + move |graph, x| { + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &w_data), + ); + let b = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &b_data), + ); + x.layer_norm(&w, Some(&b), eps, remove_mean) + .flatten_all() + .sum() + }, + ) + .await; + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference( + &device, + [1, 3], + &w_data, + move |graph, w| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &x_data), + ); + let b = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &b_data), + ); + x.layer_norm(&w, Some(&b), eps, remove_mean) + .flatten_all() + .sum() + }, + ) + .await; + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference( + &device, + [1, 3], + &b_data, + move |graph, b| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &x_data), + ); + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &w_data), + ); + x.layer_norm(&w, Some(&b), eps, remove_mean) + .flatten_all() + .sum() + }, + ) + .await; + } + } +} + +#[tokio::test] +async fn test_autograd_rms_norm_rank_generic_weight() { + for device in test_devices().await { + let eps = 1e-5f32; + let x_rows = [[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]; + let x_data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; + let w_data = [0.5f32, 1.0, 1.5]; + let graph = Graph::new(); + let x: Tensor<2> = Tensor::from_slice(&graph, &device, [2, 3], &x_data); + let w: Tensor<2> = Tensor::from_slice(&graph, &device, [1, 3], &w_data); + + let output = x.rms_norm(&w, eps); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec2(); + for (row_index, row) in x_rows.iter().enumerate() { + let mean_sq = row.iter().map(|v| v * v).sum::() / 3.0; + let rms = (mean_sq + eps).sqrt(); + for column in 0..3 { + assert_close( + output_values[row_index][column], + row[column] / rms * w_data[column], + ); + } + } + + let gradients = output.flatten_all().sum().backward().unwrap(); + let dw = gradients + .get(&w) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec2(); + for column in 0..3 { + let mut expected = 0.0f32; + for row in x_rows.iter() { + let mean_sq = row.iter().map(|v| v * v).sum::() / 3.0; + expected += row[column] / (mean_sq + eps).sqrt(); + } + assert_close(dw[0][column], expected); + } + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [2, 3], &x_data, move |graph, x| { + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &w_data), + ); + x.rms_norm(&w, eps).flatten_all().sum() + }) + .await; + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [1, 3], &w_data, move |graph, w| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &x_data), + ); + x.rms_norm(&w, eps).flatten_all().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_rms_norm_fused_weight_rank() { + for device in test_devices().await { + let eps = 1e-5f32; + let x_rows = [[1.0f32, -2.0, 3.0], [0.5, 4.0, -1.5]]; + let x_data = [1.0f32, -2.0, 3.0, 0.5, 4.0, -1.5]; + let w_data = [0.5f32, 1.0, 1.5]; + let b_data = [0.1f32, -0.2, 0.3]; + let graph = Graph::new(); + let x: Tensor<2> = Tensor::from_slice(&graph, &device, [2, 3], &x_data); + let w1: Tensor<1> = Tensor::from_slice(&graph, &device, [3], &w_data); + let b1: Tensor<1> = Tensor::from_slice(&graph, &device, [3], &b_data); + let w2: Tensor<2> = Tensor::from_slice(&graph, &device, [1, 3], &w_data); + + let biased = x.rms_norm_fused::<1, 1>(&w1, Some(&b1), eps); + let biased_values = biased.raw().clone().as_slice().await.unwrap().to_vec2(); + let no_bias = x.rms_norm_fused_no_bias::<2, 1>(&w2, eps); + let no_bias_values = no_bias.raw().clone().as_slice().await.unwrap().to_vec2(); + for (row_index, row) in x_rows.iter().enumerate() { + let mean_sq = row.iter().map(|v| v * v).sum::() / 3.0; + let rms = (mean_sq + eps).sqrt(); + for column in 0..3 { + let scaled = row[column] / rms * w_data[column]; + assert_close(biased_values[row_index][column], scaled + b_data[column]); + assert_close(no_bias_values[row_index][column], scaled); + } + } + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [2, 3], &x_data, move |graph, x| { + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3], &w_data), + ); + let b = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3], &b_data), + ); + x.rms_norm_fused::<1, 1>(&w, Some(&b), eps) + .flatten_all() + .sum() + }) + .await; + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [1, 3], &w_data, move |graph, w| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &x_data), + ); + x.rms_norm_fused_no_bias::<2, 1>(&w, eps).flatten_all().sum() + }) + .await; + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [3], &b_data, move |graph, b| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &x_data), + ); + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3], &w_data), + ); + x.rms_norm_fused::<1, 1>(&w, Some(&b), eps) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_rms_norm_residual_fused_weight_rank() { + for device in test_devices().await { + let eps = 1e-5f32; + let x_data = [1.0f32, -2.0, 3.0, 0.5, 4.0, -1.5]; + let r_data = [0.5f32, 1.5, -1.0, 2.0, -0.5, 1.0]; + let w_data = [0.5f32, 1.0, 1.5]; + let b_data = [0.1f32, -0.2, 0.3]; + let graph = Graph::new(); + let x: Tensor<2> = Tensor::from_slice(&graph, &device, [2, 3], &x_data); + let r: Tensor<2> = Tensor::from_slice(&graph, &device, [2, 3], &r_data); + let w: Tensor<2> = Tensor::from_slice(&graph, &device, [1, 3], &w_data); + let b: Tensor<2> = Tensor::from_slice(&graph, &device, [1, 3], &b_data); + + let biased = x.rms_norm_residual_fused::<2, 1>(&r, &w, Some(&b), eps); + let biased_values = biased.raw().clone().as_slice().await.unwrap().to_vec2(); + let no_bias = x.rms_norm_residual_fused::<2, 1>(&r, &w, None, eps); + let no_bias_values = no_bias.raw().clone().as_slice().await.unwrap().to_vec2(); + for row_index in 0..2 { + let combined: Vec = (0..3) + .map(|column| x_data[row_index * 3 + column] + r_data[row_index * 3 + column]) + .collect(); + let mean_sq = combined.iter().map(|v| v * v).sum::() / 3.0; + let rms = (mean_sq + eps).sqrt(); + for column in 0..3 { + let scaled = combined[column] / rms * w_data[column]; + assert_close(biased_values[row_index][column], scaled + b_data[column]); + assert_close(no_bias_values[row_index][column], scaled); + } + } + + // Input/residual gradients are covered by + // test_autograd_rms_norm_residual_fused; only the rank-2 weight/bias + // path is new here. + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [1, 3], &w_data, move |graph, w| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &x_data), + ); + let r = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &r_data), + ); + let b = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &b_data), + ); + x.rms_norm_residual_fused::<2, 1>(&r, &w, Some(&b), eps) + .flatten_all() + .sum() + }) + .await; + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [1, 3], &b_data, move |graph, b| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &x_data), + ); + let r = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &r_data), + ); + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &w_data), + ); + x.rms_norm_residual_fused::<2, 1>(&r, &w, Some(&b), eps) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_norm_last_dim_fused_weight_rank() { + for device in test_devices().await { + let eps = 1e-5f32; + let x_rows = [[1.0f32, 2.0, 4.0], [-1.0, 0.5, 3.0]]; + let x_data = [1.0f32, 2.0, 4.0, -1.0, 0.5, 3.0]; + let w_data = [0.5f32, 1.0, 1.5]; + let b_data = [0.1f32, -0.2, 0.3]; + let graph = Graph::new(); + let x: Tensor<2> = Tensor::from_slice(&graph, &device, [2, 3], &x_data); + let w1: Tensor<1> = Tensor::from_slice(&graph, &device, [3], &w_data); + let b1: Tensor<1> = Tensor::from_slice(&graph, &device, [3], &b_data); + let w2: Tensor<2> = Tensor::from_slice(&graph, &device, [1, 3], &w_data); + let b2: Tensor<2> = Tensor::from_slice(&graph, &device, [1, 3], &b_data); + + let rank1 = x.layer_norm_last_dim_fused::<1, 1>(&w1, Some(&b1), eps); + let rank1_values = rank1.raw().clone().as_slice().await.unwrap().to_vec2(); + let rank2 = x.layer_norm_last_dim_fused::<1, 2>(&w2, Some(&b2), eps); + let rank2_values = rank2.raw().clone().as_slice().await.unwrap().to_vec2(); + for (row_index, row) in x_rows.iter().enumerate() { + let mean = row.iter().sum::() / 3.0; + let var = row.iter().map(|v| (v - mean) * (v - mean)).sum::() / 3.0; + let std = (var + eps).sqrt(); + for column in 0..3 { + let expected = (row[column] - mean) / std * w_data[column] + b_data[column]; + assert_close(rank1_values[row_index][column], expected); + assert_close(rank2_values[row_index][column], expected); + } + } + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [2, 3], &x_data, move |graph, x| { + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &w_data), + ); + let b = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &b_data), + ); + x.layer_norm_last_dim_fused::<1, 2>(&w, Some(&b), eps) + .flatten_all() + .sum() + }) + .await; + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [1, 3], &w_data, move |graph, w| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &x_data), + ); + x.layer_norm_last_dim_fused::<1, 2>(&w, None, eps) + .flatten_all() + .sum() + }) + .await; + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [1, 3], &b_data, move |graph, b| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &x_data), + ); + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &w_data), + ); + x.layer_norm_last_dim_fused::<1, 2>(&w, Some(&b), eps) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_linear_forward_matches_inference() { + for device in test_devices().await { + let weight_data = [ + 0.5f32, -1.0, 0.25, 2.0, 1.5, -0.75, 0.1, 0.4, -0.2, 0.9, -1.3, 0.6, + ]; + let bias_data = [0.3f32, -0.6, 1.1]; + let input_data = [ + 1.0f32, -2.0, 0.5, 0.25, 0.75, 1.5, -0.5, 2.0, -1.25, 0.4, 0.8, -0.3, 0.15, -0.9, 1.2, + 0.7, + ]; + let weight_bytes: Vec = weight_data + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect(); + let qweight = |device: &Device| { + crate::QMatrix::from_raw_bytes(device, [3, 4], &weight_bytes, fusor_gguf::GgmlType::F32) + .unwrap() + }; + let raw_bias = RawTensor::from_slice(&device, [3], &bias_data); + let inference = crate::layers::Linear::new(qweight(&device), Some(raw_bias)); + let inference_no_bias = crate::layers::Linear::::new(qweight(&device), None); + + let graph = Graph::new(); + let layer = layers::Linear::new( + Tensor::from_slice(&graph, &device, [3, 4], &weight_data), + Some(Tensor::from_slice(&graph, &device, [3], &bias_data)), + ); + assert_eq!(layer.in_features(), 4); + assert_eq!(layer.out_features(), 3); + + let raw_input = RawTensor::from_slice(&device, [2, 2, 4], &input_data); + let input = Tensor::constant_from_raw(&graph, raw_input.clone()); + let output = layer.forward(&input); + assert_eq!(output.shape(), [2, 2, 3]); + let expected = flatten(inference.forward(&raw_input)).await; + assert_slice_close(&flatten(output.raw().clone()).await, &expected); + + let raw_input_2d = RawTensor::from_slice(&device, [4, 4], &input_data); + let input_2d = Tensor::constant_from_raw(&graph, raw_input_2d.clone()); + let output_2d = layer.forward_2d(&input_2d); + assert_eq!(output_2d.shape(), [4, 3]); + let expected_2d = flatten(inference.forward_2d(&raw_input_2d)).await; + assert_slice_close(&flatten(output_2d.raw().clone()).await, &expected_2d); + + let layer_no_bias = layers::Linear::new( + Tensor::from_slice(&graph, &device, [3, 4], &weight_data), + None, + ); + let output_no_bias = layer_no_bias.forward(&input); + let expected_no_bias = flatten(inference_no_bias.forward(&raw_input)).await; + assert_slice_close( + &flatten(output_no_bias.raw().clone()).await, + &expected_no_bias, + ); + } +} + +#[tokio::test] +async fn test_autograd_layer_linear_weight_gradient() { + for device in test_devices().await { + let weight_data = [ + 0.5f32, -1.0, 0.25, 2.0, 1.5, -0.75, 0.1, 0.4, -0.2, 0.9, -1.3, 0.6, + ]; + let bias_data = [0.3f32, -0.6, 1.1]; + let input_data = [ + 1.0f32, -2.0, 0.5, 0.25, 0.75, 1.5, -0.5, 2.0, -1.25, 0.4, 0.8, -0.3, 0.15, -0.9, 1.2, + 0.7, + ]; + assert_gradient_matches_finite_difference( + &device, + [3, 4], + &weight_data, + |graph, weight| { + let bias = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&device, [3], &bias_data), + ); + let layer = layers::Linear::new(weight, Some(bias)); + let input = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&device, [2, 2, 4], &input_data), + ); + layer.forward(&input).sqr().flatten_all().sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_linear_bias_gradient() { + for device in test_devices().await { + let weight_data = [ + 0.5f32, -1.0, 0.25, 2.0, 1.5, -0.75, 0.1, 0.4, -0.2, 0.9, -1.3, 0.6, + ]; + let bias_data = [0.3f32, -0.6, 1.1]; + let input_data = [ + 1.0f32, -2.0, 0.5, 0.25, 0.75, 1.5, -0.5, 2.0, -1.25, 0.4, 0.8, -0.3, 0.15, -0.9, 1.2, + 0.7, + ]; + assert_gradient_matches_finite_difference(&device, [3], &bias_data, |graph, bias| { + let weight = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&device, [3, 4], &weight_data), + ); + let layer = layers::Linear::new(weight, Some(bias)); + let input = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&device, [2, 2, 4], &input_data), + ); + layer.forward(&input).sqr().flatten_all().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_linear_from_inference() { + for device in test_devices().await { + let weight_data = [ + 0.5f32, -1.0, 0.25, 2.0, 1.5, -0.75, 0.1, 0.4, -0.2, 0.9, -1.3, 0.6, + ]; + let bias_data = [0.3f32, -0.6, 1.1]; + let input_data = [ + 1.0f32, -2.0, 0.5, 0.25, 0.75, 1.5, -0.5, 2.0, -1.25, 0.4, 0.8, -0.3, 0.15, -0.9, 1.2, + 0.7, + ]; + let weight_bytes: Vec = weight_data + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect(); + let qweight = crate::QMatrix::from_raw_bytes( + &device, + [3, 4], + &weight_bytes, + fusor_gguf::GgmlType::F32, + ) + .unwrap(); + let raw_bias = RawTensor::from_slice(&device, [3], &bias_data); + let inference = crate::layers::Linear::new(qweight, Some(raw_bias)); + + let graph = Graph::new(); + let layer = layers::Linear::from_inference(&graph, &inference); + assert_eq!(layer.in_features(), 4); + assert_eq!(layer.out_features(), 3); + + let raw_input = RawTensor::from_slice(&device, [2, 2, 4], &input_data); + let input = Tensor::constant_from_raw(&graph, raw_input.clone()); + let output = layer.forward(&input); + let expected = flatten(inference.forward(&raw_input)).await; + assert_slice_close(&flatten(output.raw().clone()).await, &expected); + + let gradients = output.sqr().flatten_all().sum().backward().unwrap(); + let dweight = gradients.get(layer.weight()).unwrap(); + assert_eq!(dweight.shape(), [3, 4]); + let dbias = gradients.get(layer.bias().unwrap()).unwrap(); + let mut expected_dbias = [0.0f32; 3]; + for row in expected.chunks(3) { + for (acc, value) in expected_dbias.iter_mut().zip(row) { + *acc += 2.0 * value; + } + } + assert_slice_close(&flatten(dbias).await, &expected_dbias); + } +} + +#[tokio::test] +async fn test_autograd_layer_embedding_forward_parity() { + let weights: Vec = (0..12).map(|index| (index as f32 * 0.7).sin()).collect(); + for device in test_devices().await { + let table = RawTensor::from_slice(&device, [4, 3], &weights); + let raw_layer = crate::layers::Embedding::new_from_tensor(table.clone()); + let graph = Graph::new(); + let layer = + layers::Embedding::new_from_tensor(graph.leaf(table)); + assert_eq!(layer.num_embeddings(), 4); + assert_eq!(layer.embedding_dim(), 3); + + let indices: RawTensor<2, u32> = RawTensor::from_slice(&device, [2, 2], &[0, 2, 1, 3]); + let expected: RawTensor<3, f32> = raw_layer.forward(&indices); + let output = layer.forward(&indices); + assert_eq!(output.shape(), [2, 2, 3]); + assert_slice_close(&flatten(output.raw().clone()).await, &flatten(expected).await); + + let flat_indices: RawTensor<1, u32> = RawTensor::from_slice(&device, [3], &[2, 0, 2]); + let expected_flat: RawTensor<2, f32> = raw_layer.forward(&flat_indices); + let output_flat = layer.forward_1d(&flat_indices); + assert_eq!(output_flat.shape(), [3, 3]); + assert_slice_close( + &flatten(output_flat.raw().clone()).await, + &flatten(expected_flat).await, + ); + } +} + +#[tokio::test] +async fn test_autograd_layer_embedding_weight_gradient() { + let weights: Vec = (0..12).map(|index| (index as f32 * 0.7).sin()).collect(); + for device in test_devices().await { + let indices: RawTensor<2, u32> = RawTensor::from_slice(&device, [2, 2], &[0, 2, 1, 2]); + assert_gradient_matches_finite_difference(&device, [4, 3], &weights, |_graph, table| { + layers::Embedding::new_from_tensor(table) + .forward(&indices) + .sqr() + .flatten_all() + .sum() + }) + .await; + + let flat_indices: RawTensor<1, u32> = RawTensor::from_slice(&device, [3], &[3, 1, 3]); + assert_gradient_matches_finite_difference(&device, [4, 3], &weights, |_graph, table| { + layers::Embedding::new_from_tensor(table) + .forward_1d(&flat_indices) + .sqr() + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_embedding_from_inference() { + let weights: Vec = (0..12).map(|index| index as f32 * 0.25 - 1.0).collect(); + let bytes: Vec = weights + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect(); + for device in test_devices().await { + let quantized = crate::QMatrix::from_raw_bytes( + &device, + [4usize, 3], + &bytes, + fusor_gguf::GgmlType::F32, + ) + .unwrap(); + let quantized_layer = crate::layers::Embedding::::new(quantized); + let dense_layer = crate::layers::Embedding::::new_from_tensor(RawTensor::from_slice( + &device, + [4, 3], + &weights, + )); + for raw_layer in [&quantized_layer, &dense_layer] { + let indices: RawTensor<2, u32> = RawTensor::from_slice(&device, [2, 2], &[3, 0, 2, 2]); + let expected: RawTensor<3, f32> = raw_layer.forward(&indices); + + let graph = Graph::new(); + let layer = + layers::Embedding::from_inference(&graph, raw_layer); + assert_eq!(layer.num_embeddings(), 4); + assert_eq!(layer.embedding_dim(), 3); + assert_slice_close(&flatten(layer.embeddings().raw().clone()).await, &weights); + let output = layer.forward(&indices); + assert_slice_close(&flatten(output.raw().clone()).await, &flatten(expected).await); + + let loss = layer.forward(&indices).sqr().flatten_all().sum(); + let gradients = loss.backward().unwrap(); + assert!(gradients.get(layer.embeddings()).is_some()); + } + } +} + +#[tokio::test] +async fn test_autograd_layer_layer_norm_forward_matches_inference() { + for device in test_devices().await { + let weight_data = [0.5f32, 1.5, -1.0, 2.0]; + let bias_data = [0.1f32, -0.2, 0.3, 0.4]; + let input_data: Vec = (0..24).map(|index| (index as f32 * 0.7).sin()).collect(); + + let inference = crate::layers::LayerNorm::new( + RawTensor::from_slice(&device, [4], &weight_data), + Some(RawTensor::from_slice(&device, [4], &bias_data)), + 1e-5, + ); + let raw_input = RawTensor::from_slice(&device, [2, 3, 4], &input_data); + let raw_input_2d = RawTensor::from_slice(&device, [6, 4], &input_data); + let expected = flatten(inference.forward(&raw_input)).await; + let expected_2d = flatten(inference.forward_2d(&raw_input_2d)).await; + + let graph = Graph::new(); + let layer = layers::LayerNorm::new( + Tensor::from_slice(&graph, &device, [4], &weight_data), + Some(Tensor::from_slice(&graph, &device, [4], &bias_data)), + 1e-5, + ); + let input = Tensor::from_slice(&graph, &device, [2, 3, 4], &input_data); + let input_2d = Tensor::from_slice(&graph, &device, [6, 4], &input_data); + assert_slice_close(&flatten(layer.forward(&input).into_raw()).await, &expected); + assert_slice_close( + &flatten(layer.forward_fused(&input).into_raw()).await, + &expected, + ); + assert_slice_close( + &flatten(layer.forward_2d(&input_2d).into_raw()).await, + &expected_2d, + ); + } +} + +#[tokio::test] +async fn test_autograd_layer_layer_norm_nd_forward_matches_inference() { + for device in test_devices().await { + let weight_data = [0.5f32, 1.5, -1.0, 2.0]; + let bias_data = [0.1f32, -0.2, 0.3, 0.4]; + let axis_weight_data = [0.5f32, 1.5, -1.0]; + let axis_bias_data = [0.1f32, -0.2, 0.3]; + let input_data: Vec = (0..24).map(|index| (index as f32 * 0.7).sin()).collect(); + let raw_input = RawTensor::from_slice(&device, [2, 3, 4], &input_data); + + let inference = crate::layers::LayerNormNd::new( + RawTensor::from_slice(&device, [4], &weight_data), + Some(RawTensor::from_slice(&device, [4], &bias_data)), + 1e-5, + ); + let expected = flatten(inference.forward::<3, 2, _>(&raw_input)).await; + let raw_input_2d = RawTensor::from_slice(&device, [6, 4], &input_data); + let expected_2d = flatten(inference.forward_2d(&raw_input_2d)).await; + + let inference_axis = crate::layers::LayerNormNd::new_over_axis( + RawTensor::from_slice(&device, [3], &axis_weight_data), + Some(RawTensor::from_slice(&device, [3], &axis_bias_data)), + 1, + 1e-5, + ); + let expected_axis = flatten(inference_axis.forward::<3, 2, _>(&raw_input)).await; + + let graph = Graph::new(); + let input = Tensor::from_slice(&graph, &device, [2, 3, 4], &input_data); + let input_2d = Tensor::from_slice(&graph, &device, [6, 4], &input_data); + let layer = layers::LayerNormNd::new( + Tensor::from_slice(&graph, &device, [4], &weight_data), + Some(Tensor::from_slice(&graph, &device, [4], &bias_data)), + 1e-5, + ); + assert_slice_close( + &flatten(layer.forward::<3, 2>(&input).into_raw()).await, + &expected, + ); + assert_slice_close( + &flatten(layer.forward_fused(&input).into_raw()).await, + &expected, + ); + assert_slice_close( + &flatten(layer.forward_2d(&input_2d).into_raw()).await, + &expected_2d, + ); + + let layer_axis = layers::LayerNormNd::new_over_axis( + Tensor::from_slice(&graph, &device, [3], &axis_weight_data), + Some(Tensor::from_slice(&graph, &device, [3], &axis_bias_data)), + 1, + 1e-5, + ); + assert_slice_close( + &flatten(layer_axis.forward::<3, 2>(&input).into_raw()).await, + &expected_axis, + ); + assert_slice_close( + &flatten(layer_axis.forward_fused(&input).into_raw()).await, + &expected_axis, + ); + } +} + +#[tokio::test] +async fn test_autograd_layer_layer_norm_parameter_gradients() { + for device in test_devices().await { + let weight_data = [0.5f32, 1.5, -1.0, 2.0]; + let bias_data = [0.1f32, -0.2, 0.3, 0.4]; + let input_data: Vec = (0..24).map(|index| (index as f32 * 0.7).sin()).collect(); + + assert_gradient_matches_finite_difference(&device, [4], &weight_data, |graph, weight| { + let bias = Tensor::from_slice(graph, &device, [4], &bias_data); + let input = Tensor::from_slice(graph, &device, [2, 3, 4], &input_data); + layers::LayerNorm::new(weight, Some(bias), 1e-5) + .forward(&input) + .flatten_all() + .sum() + }) + .await; + + assert_gradient_matches_finite_difference(&device, [4], &bias_data, |graph, bias| { + let weight = Tensor::from_slice(graph, &device, [4], &weight_data); + let input = Tensor::from_slice(graph, &device, [2, 3, 4], &input_data); + layers::LayerNorm::new(weight, Some(bias), 1e-5) + .forward(&input) + .flatten_all() + .sum() + }) + .await; + + assert_gradient_matches_finite_difference(&device, [4], &weight_data, |graph, weight| { + let bias = Tensor::from_slice(graph, &device, [4], &bias_data); + let input = Tensor::from_slice(graph, &device, [2, 3, 4], &input_data); + layers::LayerNorm::new(weight, Some(bias), 1e-5) + .forward_fused(&input) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_layer_norm_nd_parameter_gradients() { + for device in test_devices().await { + let weight_data = [0.5f32, 1.5, -1.0]; + let bias_data = [0.1f32, -0.2, 0.3]; + let input_data: Vec = (0..24).map(|index| (index as f32 * 0.7).sin()).collect(); + + assert_gradient_matches_finite_difference(&device, [3], &weight_data, |graph, weight| { + let bias = Tensor::from_slice(graph, &device, [3], &bias_data); + let input = Tensor::from_slice(graph, &device, [2, 3, 4], &input_data); + layers::LayerNormNd::new_over_axis(weight, Some(bias), 1, 1e-5) + .forward::<3, 2>(&input) + .flatten_all() + .sum() + }) + .await; + + assert_gradient_matches_finite_difference(&device, [3], &bias_data, |graph, bias| { + let weight = Tensor::from_slice(graph, &device, [3], &weight_data); + let input = Tensor::from_slice(graph, &device, [2, 3, 4], &input_data); + layers::LayerNormNd::new_over_axis(weight, Some(bias), 1, 1e-5) + .forward::<3, 2>(&input) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_layer_norm_from_inference_roundtrip() { + for device in test_devices().await { + let weight_data = [0.5f32, 1.5, -1.0, 2.0]; + let bias_data = [0.1f32, -0.2, 0.3, 0.4]; + let axis_weight_data = [0.5f32, 1.5, -1.0]; + let input_data: Vec = (0..24).map(|index| (index as f32 * 0.7).sin()).collect(); + let raw_input = RawTensor::from_slice(&device, [2, 3, 4], &input_data); + + let inference = crate::layers::LayerNorm::new( + RawTensor::from_slice(&device, [4], &weight_data), + Some(RawTensor::from_slice(&device, [4], &bias_data)), + 1e-5, + ); + let expected = flatten(inference.forward(&raw_input)).await; + + let graph = Graph::new(); + let input = Tensor::constant_from_raw(&graph, raw_input.clone()); + let layer = layers::LayerNorm::from_inference(&graph, &inference); + let output = layer.forward(&input); + assert_slice_close(&flatten(output.raw().clone()).await, &expected); + + let gradients = output.flatten_all().sum().backward().unwrap(); + assert_eq!(gradients.get(layer.weight()).unwrap().shape(), [4]); + assert_eq!(gradients.get(layer.bias().unwrap()).unwrap().shape(), [4]); + assert!(gradients.get(&input).is_none()); + + let inference_nd = crate::layers::LayerNormNd::new( + RawTensor::from_slice(&device, [4], &weight_data), + Some(RawTensor::from_slice(&device, [4], &bias_data)), + 1e-5, + ); + let expected_nd = flatten(inference_nd.forward::<3, 2, _>(&raw_input)).await; + let layer_nd = layers::LayerNormNd::from_inference(&graph, &inference_nd); + assert_slice_close( + &flatten(layer_nd.forward::<3, 2>(&input).into_raw()).await, + &expected_nd, + ); + + let inference_axis = crate::layers::LayerNormNd::new_over_axis( + RawTensor::from_slice(&device, [3], &axis_weight_data), + None, + 1, + 1e-5, + ); + let expected_axis = flatten(inference_axis.forward::<3, 2, _>(&raw_input)).await; + let layer_axis = + layers::LayerNormNd::from_inference_over_axis(&graph, &inference_axis, 1); + assert_slice_close( + &flatten(layer_axis.forward::<3, 2>(&input).into_raw()).await, + &expected_axis, + ); + } +} + +#[tokio::test] +async fn test_autograd_layer_rms_norm_forward_parity() { + let weight_data = [0.5f32, 1.5, -0.75, 2.0]; + let bias_data = [0.1f32, -0.2, 0.3, -0.4]; + let input_data: Vec = (0..24).map(|i| (i as f32 * 0.37).sin()).collect(); + let residual_data: Vec = (0..24).map(|i| (i as f32 * 0.61).cos()).collect(); + for device in test_devices().await { + let raw_weight = RawTensor::from_slice(&device, [4], &weight_data); + let raw_bias = RawTensor::from_slice(&device, [4], &bias_data); + let raw_layer = + crate::layers::RmsNorm::new(raw_weight.clone(), Some(raw_bias.clone()), 1e-5); + + let graph = Graph::new(); + let layer = layers::RmsNorm::new( + graph.leaf(raw_weight), + Some(graph.leaf(raw_bias)), + 1e-5, + ); + + let input_2d = RawTensor::from_slice(&device, [6, 4], &input_data); + let expected = flatten(raw_layer.forward_2d(&input_2d)).await; + let output = layer.forward_2d(&Tensor::constant_from_raw(&graph, input_2d)); + assert_slice_close(&flatten(output.raw().clone()).await, &expected); + + let input_3d = RawTensor::from_slice(&device, [2, 3, 4], &input_data); + let expected = flatten(raw_layer.forward(&input_3d)).await; + let output = layer.forward(&Tensor::constant_from_raw(&graph, input_3d.clone())); + assert_slice_close(&flatten(output.raw().clone()).await, &expected); + + let input_4d = RawTensor::from_slice(&device, [2, 1, 3, 4], &input_data); + let expected = flatten(raw_layer.forward_4d(&input_4d)).await; + let output = layer.forward_4d(&Tensor::constant_from_raw(&graph, input_4d)); + assert_slice_close(&flatten(output.raw().clone()).await, &expected); + + let residual_3d = RawTensor::from_slice(&device, [2, 3, 4], &residual_data); + let expected = flatten(raw_layer.forward_residual_f32(&input_3d, &residual_3d)).await; + let output = layer.forward_residual( + &Tensor::constant_from_raw(&graph, input_3d), + &Tensor::constant_from_raw(&graph, residual_3d), + ); + assert_slice_close(&flatten(output.raw().clone()).await, &expected); + } +} + +#[tokio::test] +async fn test_autograd_layer_rms_norm_weight_gradient() { + let weight_data = [0.5f32, 1.5, -0.75, 2.0]; + let bias_data = [0.1f32, -0.2, 0.3, -0.4]; + let input_data: Vec = (0..24).map(|i| (i as f32 * 0.37).sin() + 0.25).collect(); + let residual_data: Vec = (0..24).map(|i| (i as f32 * 0.61).cos()).collect(); + for device in test_devices().await { + let input_2d = RawTensor::from_slice(&device, [6, 4], &input_data); + assert_gradient_matches_finite_difference(&device, [4], &weight_data, |graph, weight| { + let layer = layers::RmsNorm::new(weight, None, 1e-5); + layer + .forward_2d(&Tensor::constant_from_raw(graph, input_2d.clone())) + .sqr() + .flatten_all() + .sum() + }) + .await; + + let input_3d = RawTensor::from_slice(&device, [2, 3, 4], &input_data); + let bias = RawTensor::from_slice(&device, [4], &bias_data); + assert_gradient_matches_finite_difference(&device, [4], &weight_data, |graph, weight| { + let layer = layers::RmsNorm::new( + weight, + Some(Tensor::constant_from_raw(graph, bias.clone())), + 1e-5, + ); + layer + .forward(&Tensor::constant_from_raw(graph, input_3d.clone())) + .sqr() + .flatten_all() + .sum() + }) + .await; + + let residual_3d = RawTensor::from_slice(&device, [2, 3, 4], &residual_data); + assert_gradient_matches_finite_difference(&device, [4], &weight_data, |graph, weight| { + let layer = layers::RmsNorm::new( + weight, + Some(Tensor::constant_from_raw(graph, bias.clone())), + 1e-5, + ); + layer + .forward_residual( + &Tensor::constant_from_raw(graph, input_3d.clone()), + &Tensor::constant_from_raw(graph, residual_3d.clone()), + ) + .sqr() + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_rms_norm_bias_gradient() { + let weight_data = [0.5f32, 1.5, -0.75, 2.0]; + let bias_data = [0.1f32, -0.2, 0.3, -0.4]; + let input_data: Vec = (0..24).map(|i| (i as f32 * 0.37).sin() + 0.25).collect(); + let residual_data: Vec = (0..24).map(|i| (i as f32 * 0.61).cos()).collect(); + for device in test_devices().await { + let input_3d = RawTensor::from_slice(&device, [2, 3, 4], &input_data); + let weight = RawTensor::from_slice(&device, [4], &weight_data); + assert_gradient_matches_finite_difference(&device, [4], &bias_data, |graph, bias| { + let layer = layers::RmsNorm::new( + Tensor::constant_from_raw(graph, weight.clone()), + Some(bias), + 1e-5, + ); + layer + .forward(&Tensor::constant_from_raw(graph, input_3d.clone())) + .sqr() + .flatten_all() + .sum() + }) + .await; + + let residual_3d = RawTensor::from_slice(&device, [2, 3, 4], &residual_data); + assert_gradient_matches_finite_difference(&device, [4], &bias_data, |graph, bias| { + let layer = layers::RmsNorm::new( + Tensor::constant_from_raw(graph, weight.clone()), + Some(bias), + 1e-5, + ); + layer + .forward_residual( + &Tensor::constant_from_raw(graph, input_3d.clone()), + &Tensor::constant_from_raw(graph, residual_3d.clone()), + ) + .sqr() + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_rms_norm_from_inference() { + let weight_data = [0.5f32, 1.5, -0.75, 2.0]; + let bias_data = [0.1f32, -0.2, 0.3, -0.4]; + let input_data: Vec = (0..24).map(|i| (i as f32 * 0.37).sin()).collect(); + for device in test_devices().await { + let raw_layer = crate::layers::RmsNorm::new( + RawTensor::from_slice(&device, [4], &weight_data), + Some(RawTensor::from_slice(&device, [4], &bias_data)), + 1e-5, + ); + let graph = Graph::new(); + let layer = layers::RmsNorm::from_inference(&graph, &raw_layer); + assert_eq!(layer.eps(), raw_layer.eps()); + + let input = RawTensor::from_slice(&device, [2, 3, 4], &input_data); + let expected = flatten(raw_layer.forward(&input)).await; + let output = layer.forward(&Tensor::constant_from_raw(&graph, input)); + assert_slice_close(&flatten(output.raw().clone()).await, &expected); + + let gradients = output.flatten_all().sum().backward().unwrap(); + let weight_gradient = gradients.get(layer.weight()).unwrap(); + let bias_gradient = gradients.get(layer.bias().unwrap()).unwrap(); + assert_eq!(weight_gradient.shape(), [4]); + assert_eq!(bias_gradient.shape(), [4]); + } +} + +#[tokio::test] +async fn test_autograd_layer_conv1d_forward_parity() { + for device in test_devices().await { + let w_data: Vec = (0..12).map(|i| (i as f32 * 0.37).cos()).collect(); + let b_data = [0.3f32, -0.7]; + let x_data: Vec = (0..15).map(|i| (i as f32 * 0.61).sin()).collect(); + let config = crate::layers::Conv1dConfig { + padding: 1, + stride: 2, + ..Default::default() + }; + + let graph = Graph::new(); + let layer = layers::Conv1d::new( + Tensor::from_slice(&graph, &device, [2, 3, 2], &w_data), + Some(Tensor::from_slice(&graph, &device, [2], &b_data)), + config, + ); + let x: Tensor<3> = Tensor::from_slice(&graph, &device, [1, 3, 5], &x_data); + let output = layer.forward(&x); + + let raw_layer = crate::layers::Conv1d::new( + RawTensor::from_slice(&device, [2, 3, 2], &w_data), + Some(RawTensor::from_slice(&device, [2], &b_data)), + config, + ); + let expected = raw_layer.forward(&RawTensor::from_slice(&device, [1, 3, 5], &x_data)); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + assert_eq!(layer.in_channels(), 3); + assert_eq!(layer.out_channels(), 2); + assert_eq!(layer.kernel_size(), 2); + assert_eq!(layer.config().stride, 2); + } +} + +#[tokio::test] +async fn test_autograd_layer_conv1d_parameter_gradients() { + for device in test_devices().await { + let w_data: Vec = (0..12).map(|i| (i as f32 * 0.37).cos()).collect(); + let b_data = [0.3f32, -0.7]; + let x_data: Vec = (0..15).map(|i| (i as f32 * 0.61).sin()).collect(); + let config = crate::layers::Conv1dConfig { + padding: 1, + stride: 2, + ..Default::default() + }; + + let graph = Graph::new(); + let layer = layers::Conv1d::new( + Tensor::from_slice(&graph, &device, [2, 3, 2], &w_data), + Some(Tensor::from_slice(&graph, &device, [2], &b_data)), + config, + ); + let x = Tensor::constant_from_raw( + &graph, + RawTensor::from_slice(&device, [1, 3, 5], &x_data), + ); + let gradients = layer.forward(&x).flatten_all().sum().backward().unwrap(); + assert!(gradients.get(layer.weight()).is_some()); + assert!(gradients.get(layer.bias().unwrap()).is_some()); + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + let b_fd = b_data; + assert_gradient_matches_finite_difference( + &device, + [2, 3, 2], + &w_data, + move |graph, w| { + let layer = layers::Conv1d::new( + w, + Some(Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2], &b_fd), + )), + config, + ); + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3, 5], &x_fd), + ); + let out = layer.forward(&x); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }, + ) + .await; + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + let w_fd = w_data.clone(); + assert_gradient_matches_finite_difference(&device, [2], &b_data, move |graph, b| { + let layer = layers::Conv1d::new( + Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3, 2], &w_fd), + ), + Some(b), + config, + ); + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3, 5], &x_fd), + ); + let out = layer.forward(&x); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_conv1d_from_inference() { + for device in test_devices().await { + let w_data: Vec = (0..12).map(|i| (i as f32 * 0.37).cos()).collect(); + let b_data = [0.3f32, -0.7]; + let x_data: Vec = (0..15).map(|i| (i as f32 * 0.61).sin()).collect(); + let config = crate::layers::Conv1dConfig { + padding: 2, + stride: 2, + ..Default::default() + }; + let raw_layer = crate::layers::Conv1d::new( + RawTensor::from_slice(&device, [2, 3, 2], &w_data), + Some(RawTensor::from_slice(&device, [2], &b_data)), + config, + ); + + let graph = Graph::new(); + let layer = layers::Conv1d::from_inference(&graph, &raw_layer); + assert_slice_close(&flatten(layer.weight().raw().clone()).await, &w_data); + assert_slice_close( + &flatten(layer.bias().unwrap().raw().clone()).await, + &b_data, + ); + + let x: Tensor<3> = Tensor::from_slice(&graph, &device, [1, 3, 5], &x_data); + let output = layer.forward(&x); + let expected = raw_layer.forward(&RawTensor::from_slice(&device, [1, 3, 5], &x_data)); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + + let raw_no_bias = crate::layers::Conv1d::new( + RawTensor::from_slice(&device, [2, 3, 2], &w_data), + None, + crate::layers::Conv1dConfig::default(), + ); + let layer = layers::Conv1d::from_inference(&graph, &raw_no_bias); + assert!(layer.bias().is_none()); + let output = layer.forward(&x); + let expected = raw_no_bias.forward(&RawTensor::from_slice(&device, [1, 3, 5], &x_data)); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + } +} + +#[tokio::test] +async fn test_autograd_layer_conv_nd_forward_parity() { + for device in test_devices().await { + { + let x_data: Vec = (0..18).map(|i| (i as f32 * 0.41).sin()).collect(); + let w_data: Vec = (0..16).map(|i| (i as f32 * 0.23).cos()).collect(); + let b_data = [0.5f32, -1.0]; + let config = crate::layers::ConvNdConfig { + padding: [1, 1], + stride: [2, 2], + groups: 1, + }; + + let raw_layer = crate::layers::ConvNd::<2, 4, f32>::new( + RawTensor::from_slice(&device, [2, 2, 2, 2], &w_data), + Some(RawTensor::from_slice(&device, [2], &b_data)), + config, + ); + let expected = + raw_layer.forward(&RawTensor::from_slice(&device, [1, 2, 3, 3], &x_data)); + + let graph = Graph::new(); + let layer = layers::ConvNd::<2, 4>::new( + Tensor::from_slice(&graph, &device, [2, 2, 2, 2], &w_data), + Some(Tensor::from_slice(&graph, &device, [2], &b_data)), + config, + ); + let x: Tensor<4> = Tensor::from_slice(&graph, &device, [1, 2, 3, 3], &x_data); + let output = layer.forward(&x); + + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + } + + { + let x_data: Vec = (0..16).map(|i| (i as f32 * 0.57).sin()).collect(); + let w_data: Vec = (0..16).map(|i| (i as f32 * 0.31).cos()).collect(); + let b_data = [0.2f32, -0.4, 0.6, -0.8]; + let config = crate::layers::ConvNdConfig { + padding: [1], + stride: [2], + groups: 2, + }; + + let raw_layer = crate::layers::ConvNd::<1, 3, f32>::new( + RawTensor::from_slice(&device, [4, 2, 2], &w_data), + Some(RawTensor::from_slice(&device, [4], &b_data)), + config, + ); + let expected = raw_layer.forward(&RawTensor::from_slice(&device, [1, 4, 4], &x_data)); + + let graph = Graph::new(); + let layer = layers::ConvNd::<1, 3>::new( + Tensor::from_slice(&graph, &device, [4, 2, 2], &w_data), + Some(Tensor::from_slice(&graph, &device, [4], &b_data)), + config, + ); + let x: Tensor<3> = Tensor::from_slice(&graph, &device, [1, 4, 4], &x_data); + let output = layer.forward(&x); + + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + } + } +} + +#[tokio::test] +async fn test_autograd_layer_conv_nd_parameter_gradients() { + for device in test_devices().await { + let x_data: Vec = (0..16).map(|i| (i as f32 * 0.57).sin()).collect(); + let w_data: Vec = (0..16).map(|i| (i as f32 * 0.31).cos()).collect(); + let b_data = [0.2f32, -0.4, 0.6, -0.8]; + let config = crate::layers::ConvNdConfig { + padding: [1], + stride: [2], + groups: 2, + }; + + { + let graph = Graph::new(); + let layer = layers::ConvNd::<1, 3>::new( + Tensor::from_slice(&graph, &device, [4, 2, 2], &w_data), + Some(Tensor::from_slice(&graph, &device, [4], &b_data)), + config, + ); + let x = Tensor::constant_from_raw( + &graph, + RawTensor::from_slice(&device, [1, 4, 4], &x_data), + ); + let gradients = layer.forward(&x).flatten_all().sum().backward().unwrap(); + assert_eq!(gradients.get(layer.weight()).unwrap().shape(), [4, 2, 2]); + assert_eq!(gradients.get(layer.bias().unwrap()).unwrap().shape(), [4]); + } + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + assert_gradient_matches_finite_difference( + &device, + [4, 2, 2], + &w_data, + move |graph, w| { + let layer = layers::ConvNd::<1, 3>::new( + w, + Some(Tensor::from_slice(graph, &fd_device, [4], &b_data)), + config, + ); + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 4, 4], &x_fd), + ); + let out = layer.forward(&x); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }, + ) + .await; + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + let w_fd = w_data.clone(); + assert_gradient_matches_finite_difference(&device, [4], &b_data, move |graph, b| { + let layer = layers::ConvNd::<1, 3>::new( + Tensor::from_slice(graph, &fd_device, [4, 2, 2], &w_fd), + Some(b), + config, + ); + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 4, 4], &x_fd), + ); + let out = layer.forward(&x); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_conv_nd_from_inference() { + for device in test_devices().await { + let x_data: Vec = (0..16).map(|i| (i as f32 * 0.57).sin()).collect(); + let w_data: Vec = (0..16).map(|i| (i as f32 * 0.31).cos()).collect(); + let b_data = [0.2f32, -0.4, 0.6, -0.8]; + let config = crate::layers::ConvNdConfig { + padding: [1], + stride: [2], + groups: 2, + }; + let raw_layer = crate::layers::ConvNd::<1, 3, f32>::new( + RawTensor::from_slice(&device, [4, 2, 2], &w_data), + Some(RawTensor::from_slice(&device, [4], &b_data)), + config, + ); + + let graph = Graph::new(); + let layer = layers::ConvNd::<1, 3>::from_inference(&graph, &raw_layer); + assert_slice_close(&flatten(layer.weight().raw().clone()).await, &w_data); + assert_slice_close(&flatten(layer.bias().unwrap().raw().clone()).await, &b_data); + + let x_raw = RawTensor::from_slice(&device, [1, 4, 4], &x_data); + let expected = raw_layer.forward(&x_raw); + let output = layer.forward(&Tensor::constant_from_raw(&graph, x_raw)); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + + let gradients = output.flatten_all().sum().backward().unwrap(); + assert!(gradients.get(layer.weight()).is_some()); + assert!(gradients.get(layer.bias().unwrap()).is_some()); + + let raw_no_bias = crate::layers::ConvNd::<1, 3, f32>::new( + RawTensor::from_slice(&device, [4, 2, 2], &w_data), + None, + config, + ); + let imported = layers::ConvNd::<1, 3>::from_inference(&graph, &raw_no_bias); + assert!(imported.bias().is_none()); + let x_raw = RawTensor::from_slice(&device, [1, 4, 4], &x_data); + let expected = raw_no_bias.forward(&x_raw); + let output = imported.forward(&Tensor::constant_from_raw(&graph, x_raw)); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + } +} + +/// End-to-end training with the trainable autograd layers: the same 2-16-2 +/// XOR MLP as [`test_train_xor_classifier`], but built from two +/// `layers::Linear` layers with gradients fetched through the +/// layer's `weight()`/`bias()` parameter handles. +#[tokio::test] +async fn test_train_xor_with_layers() { + const SAMPLES: usize = 64; + const HIDDEN: usize = 16; + const STEPS: usize = 500; + const LEARNING_RATE: f32 = 1.0; + + let (features, labels, w1_init, w2_init) = xor_training_data(HIDDEN); + + for (device, name) in test_devices().await.into_iter().zip(["cpu", "gpu"]) { + let inputs = RawTensor::from_slice(&device, [SAMPLES, 2], &features); + let targets = RawTensor::from_slice(&device, [SAMPLES], &labels); + + // Linear stores weights as (out_features, in_features). + let mut w1 = RawTensor::from_slice(&device, [HIDDEN, 2], &w1_init); + let mut b1 = RawTensor::zeros(&device, [HIDDEN]); + let mut w2 = RawTensor::from_slice(&device, [2, HIDDEN], &w2_init); + let mut b2 = RawTensor::zeros(&device, [2]); + + let mut final_loss = f32::INFINITY; + for step in 0..STEPS { + let graph = Graph::new(); + let x = Tensor::constant_from_raw(&graph, inputs.clone()); + let layer1 = layers::Linear::new( + graph.leaf(w1.clone()), + Some(graph.leaf(b1.clone())), + ); + let layer2 = layers::Linear::new( + graph.leaf(w2.clone()), + Some(graph.leaf(b2.clone())), + ); + + let hidden = layer1.forward_2d(&x).relu(); + let logits = layer2.forward_2d(&hidden); + // Numerically stable cross-entropy: log softmax via log-sum-exp + // so a saturated class cannot underflow to log(0). + let shifted = logits.sub_::<2, 2>(&logits.max_keepdim::<1>(1)); + let log_sum_exp = shifted.exp().sum_keepdim(1).log(); + let label_log_probs = shifted.sub_::<2, 2>(&log_sum_exp).gather_last(&targets); + let loss: Tensor<0> = label_log_probs.sum().mul_scalar(-1.0 / SAMPLES as f32); + + let loss_value = flatten(loss.raw().clone()).await[0]; + let gradients = loss.backward().unwrap().into_detached(); + let dw1 = gradients.get(layer1.weight()).unwrap(); + let db1 = gradients.get(layer1.bias().unwrap()).unwrap(); + let dw2 = gradients.get(layer2.weight()).unwrap(); + let db2 = gradients.get(layer2.bias().unwrap()).unwrap(); + + w1 = (w1 - dw1 * LEARNING_RATE).to_concrete(); + b1 = (b1 - db1 * LEARNING_RATE).to_concrete(); + w2 = (w2 - dw2 * LEARNING_RATE).to_concrete(); + b2 = (b2 - db2 * LEARNING_RATE).to_concrete(); + + final_loss = loss_value; + if step % 100 == 0 { + eprintln!("[{name}] step {step}: loss {loss_value:.4}"); + } + } + eprintln!("[{name}] final loss {final_loss:.4}"); + + let graph = Graph::new(); + let x = Tensor::constant_from_raw(&graph, inputs.clone()); + let layer1 = layers::Linear::new( + Tensor::constant_from_raw(&graph, w1.clone()), + Some(Tensor::constant_from_raw(&graph, b1.clone())), + ); + let layer2 = layers::Linear::new( + Tensor::constant_from_raw(&graph, w2.clone()), + Some(Tensor::constant_from_raw(&graph, b2.clone())), + ); + let hidden = layer1.forward_2d(&x).relu(); + let logits = layer2.forward_2d(&hidden); + let logits = logits.raw().clone().as_slice().await.unwrap().to_vec2(); + let correct = logits + .iter() + .zip(&labels) + .filter(|(row, label)| u32::from(row[1] > row[0]) == **label) + .count(); + eprintln!("[{name}] accuracy {correct}/{SAMPLES}"); + + assert!( + final_loss < 0.1, + "training did not converge: final loss {final_loss}", + ); + assert_eq!(correct, SAMPLES, "classifier misclassified training points"); + } +} diff --git a/fusor-ml/fusor/src/autograd/view.rs b/fusor-ml/fusor/src/autograd/view.rs new file mode 100644 index 000000000..069623ad5 --- /dev/null +++ b/fusor-ml/fusor/src/autograd/view.rs @@ -0,0 +1,888 @@ +use std::ops::Range; + +use crate::{Dim, Layout}; +use fusor_types::{SlidingWindow, StrideSpec}; + +use super::*; + +impl Tensor { + pub fn reshape(&self, shape: [usize; OUT]) -> Tensor { + let input_shape = self.shape(); + let value = self.value.reshape(shape).to_concrete(); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "reshape")?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(gradient.reshape(input_shape).to_concrete()), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn transpose(&self, dim0: usize, dim1: usize) -> Self { + let value = self.value.transpose(dim0, dim1).to_concrete(); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "transpose")?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(gradient.transpose(dim0, dim1).to_concrete()), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn permute(&self, axes: [usize; R]) -> Self { + let value = self.value.permute(axes).to_concrete(); + let input_id = self.handle.id; + let mut inverse = [0usize; R]; + for (index, axis) in axes.iter().copied().enumerate() { + inverse[axis] = index; + } + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "permute")?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(gradient.permute(inverse).to_concrete()), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn slice(&self, slices: [Range; R]) -> Self { + let input_shape = self.shape(); + let value = self.value.slice(slices.clone()).to_concrete(); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "slice")?; + let zeros = RawTensor::zeros(&gradient.device(), input_shape); + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(zeros.slice_assign(slices.clone(), &gradient).to_concrete()), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn broadcast_as(&self, shape: [usize; OUT]) -> Tensor { + let input_shape = self.shape(); + let value = self.value.broadcast_as(shape).to_concrete(); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "broadcast_as")?; + let reduced = reduce_broadcast_gradient(gradient, input_shape)?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: reduced, + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn expand(&self, shape: [usize; OUT]) -> Tensor { + self.broadcast_as(shape) + } + + pub fn flatten_all(&self) -> Tensor<1> { + self.reshape([self.shape().iter().product()]) + } + + pub fn flatten_last_n(&self) -> Tensor + where + crate::gpu::Tensor: crate::gpu::SmallerRank, + { + let shape = self.shape(); + let new_shape: [usize; OUT] = std::array::from_fn(|i| { + if i < R - 1 - FROM_END { + shape[i] + } else if i == R - 1 - FROM_END { + shape[R - 1 - FROM_END..].iter().product() + } else { + 1 + } + }); + self.reshape(new_shape) + } + + pub fn flatten_first_n(&self) -> Tensor + where + crate::gpu::Tensor: crate::gpu::SmallerRank, + { + let shape = self.shape(); + let new_shape: [usize; OUT] = std::array::from_fn(|i| { + if i == 0 { + shape[..=FROM_START].iter().product() + } else { + shape[i + FROM_START] + } + }); + self.reshape(new_shape) + } + + pub fn narrow(&self, dim: impl Dim, start: usize, length: usize) -> Self { + let dim = dim.resolve(); + let shape = self.shape(); + let slices: [Range; R] = std::array::from_fn(|axis| { + if axis == dim { + start..start + length + } else { + 0..shape[axis] + } + }); + self.slice(slices) + } + + pub fn chunk(&self, chunks: usize, dim: impl Dim) -> Vec { + let dim = dim.resolve(); + let shape = self.shape(); + let dim_size = shape[dim]; + let chunk_size = dim_size.div_ceil(chunks); + + let mut result = Vec::with_capacity(chunks); + let mut start = 0; + while start < dim_size { + let length = chunk_size.min(dim_size - start); + result.push(self.narrow(dim, start, length)); + start += length; + } + result + } + + pub fn repeat(&self, repeats: [usize; R]) -> Self { + let input_shape = self.shape(); + let value = self.value.repeat(repeats).to_concrete(); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "repeat")?; + let total: usize = gradient.shape().iter().product(); + let mut flat = gradient.reshape([total]).to_concrete(); + for axis in (0..R).rev() { + if repeats[axis] == 1 { + continue; + } + let before: usize = (0..axis) + .map(|dim| repeats[dim] * input_shape[dim]) + .product(); + let after: usize = input_shape[axis + 1..].iter().product(); + flat = flat + .reshape([before, repeats[axis], input_shape[axis], after]) + .to_concrete() + .sum::<3>(1) + .reshape([before * input_shape[axis] * after]) + .to_concrete(); + } + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(flat.reshape(input_shape).to_concrete()), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn resize(&self, new_shape: [usize; R]) -> Self { + let input_shape = self.shape(); + let value = self.value.resize(new_shape).to_concrete(); + let input_id = self.handle.id; + let copy_shape = std::array::from_fn(|axis| input_shape[axis].min(new_shape[axis])); + let copy_slices = copy_shape.map(|size| 0..size); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "resize")?; + let patch = gradient.slice(copy_slices.clone()).to_concrete(); + let zeros = RawTensor::zeros(&gradient.device(), input_shape); + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new( + zeros + .slice_assign(copy_slices.clone(), &patch) + .to_concrete(), + ), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn restride(&self, specs: [StrideSpec; OUT]) -> Tensor { + let input_shape = self.shape(); + let value = self.value.restride(specs).to_concrete(); + let input_id = self.handle.id; + let output_shape: [usize; OUT] = specs.map(|spec| spec.size); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "restride")?; + let reduced = reduce_restride_gradient(&gradient, &specs, [0; R], input_shape) + .unwrap_or_else(|| { + scatter_restride_gradient( + &gradient, + output_shape, + input_shape, + |output_index| restride_input_index(specs, output_index), + ) + }); + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(reduced), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn restride_layout(&self, new_layout: Layout) -> Tensor { + assert_eq!(new_layout.rank(), OUT, "restride_layout rank mismatch"); + let input_shape = self.shape(); + let value = self.value.restride_layout(new_layout.clone()).to_concrete(); + let input_id = self.handle.id; + let output_shape: [usize; OUT] = std::array::from_fn(|axis| new_layout.shape()[axis]); + let input_strides = Layout::continuous_strides(&input_shape); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "restride_layout")?; + let reduced = layout_restride_specs(&new_layout, input_shape) + .and_then(|(specs, offsets)| { + reduce_restride_gradient(&gradient, &specs, offsets, input_shape) + }) + .unwrap_or_else(|| { + scatter_restride_gradient( + &gradient, + output_shape, + input_shape, + |output_index| { + let linear = new_layout.linear_index(&output_index); + contiguous_index_from_linear::(linear, &input_strides) + }, + ) + }); + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(reduced), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn squeeze_dims( + &self, + axes: [usize; DIFF], + ) -> Tensor + where + crate::gpu::Tensor: crate::gpu::SmallerRank, + { + let shape = self.shape(); + for &axis in &axes { + assert_eq!( + shape[axis], 1, + "Squeeze dimension {} must have size 1", + axis + ); + } + let mut sorted_axes = axes; + sorted_axes.sort_unstable(); + let mut input_axis = 0; + let mut axis_index = 0; + let specs: [StrideSpec; OUT] = std::array::from_fn(|_| { + while axis_index < DIFF && input_axis == sorted_axes[axis_index] { + input_axis += 1; + axis_index += 1; + } + let spec = StrideSpec::dim(input_axis, shape[input_axis]); + input_axis += 1; + spec + }); + self.restride(specs) + } + + pub fn unsqueeze_dims( + &self, + axes: [usize; DIFF], + ) -> Tensor + where + crate::gpu::Tensor: crate::gpu::LargerRank, + { + let shape = self.shape(); + let mut sorted_axes = axes; + sorted_axes.sort_unstable(); + let mut input_axis = 0; + let mut axis_index = 0; + let specs: [StrideSpec; OUT] = std::array::from_fn(|output_axis| { + if axis_index < DIFF && output_axis == sorted_axes[axis_index] { + axis_index += 1; + StrideSpec::dim_with(0, 1, 0) + } else { + let spec = StrideSpec::dim(input_axis, shape[input_axis]); + input_axis += 1; + spec + } + }); + self.restride(specs) + } + + pub fn squeeze(&self, dim: usize) -> Tensor + where + crate::gpu::Tensor: crate::gpu::SmallerRank<1, OUT, f32>, + { + self.squeeze_dims::<1, OUT>([dim]) + } + + pub fn unsqueeze(&self, dim: usize) -> Tensor + where + crate::gpu::Tensor: crate::gpu::LargerRank<1, OUT, f32>, + { + self.unsqueeze_dims::<1, OUT>([dim]) + } + + pub fn sliding_window_view( + &self, + windows: [SlidingWindow; DIFF], + ) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank, + { + let shape = self.shape(); + let mut sorted_windows = windows; + sorted_windows.sort_by_key(|window| window.axis); + let specs: [StrideSpec; R2] = std::array::from_fn(|out_i| { + if out_i < R { + if let Some(window) = sorted_windows.iter().find(|window| window.axis == out_i) { + let positions = (shape[out_i] - window.window_size) / window.step + 1; + StrideSpec::dim_with(out_i, positions, window.step) + } else { + StrideSpec::dim(out_i, shape[out_i]) + } + } else { + let window = &sorted_windows[out_i - R]; + StrideSpec::dim(window.axis, window.window_size) + } + }); + self.restride(specs) + } + + pub fn pad_axis(&self, axis: usize, padding: usize) -> Self { + self.pad_with_zeros(axis, padding, padding) + } + + pub fn pad_with_zeros(&self, axis: usize, left: usize, right: usize) -> Self { + if left == 0 && right == 0 { + return self.clone(); + } + let shape = self.shape(); + let mut padded_shape = shape; + padded_shape[axis] += left + right; + let padded = Self::constant_from_raw( + &self.graph(), + RawTensor::zeros(&self.device(), padded_shape), + ); + let slices: [Range; R] = std::array::from_fn(|dim| { + if dim == axis { + left..left + shape[axis] + } else { + 0..shape[dim] + } + }); + padded.slice_assign(slices, self) + } + + pub fn slice_assign(&self, slices: [Range; R], value: &Self) -> Self { + assert_same_graph(self, value); + + let output = self + .value + .slice_assign(slices.clone(), &value.value) + .to_concrete(); + let input_id = self.handle.id; + let value_id = value.handle.id; + let slice_shape = slices + .clone() + .map(|range| range.end.saturating_sub(range.start)); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "slice_assign")?; + let zeros = RawTensor::zeros(&gradient.device(), slice_shape); + Ok(vec![ + BackwardTarget { + node: input_id, + gradient: Box::new(gradient.slice_assign(slices.clone(), &zeros).to_concrete()), + }, + BackwardTarget { + node: value_id, + gradient: Box::new(gradient.slice(slices.clone()).to_concrete()), + }, + ]) + }); + self.emit_op( + output, + vec![self.handle.clone(), value.handle.clone()], + Some(backward), + ) + } + + pub fn stack( + tensors: impl IntoIterator, + dim: usize, + ) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank<1, OUT, f32>, + { + let tensors: Vec = tensors.into_iter().collect(); + assert!(!tensors.is_empty(), "stack requires at least one tensor"); + + let graph = tensors[0].handle.graph.clone(); + let input_shape = tensors[0].shape(); + let raw = tensors + .iter() + .map(|tensor| { + assert!( + Arc::ptr_eq(&graph, &tensor.handle.graph), + "cannot mix autograd tensors from different graphs" + ); + assert_eq!( + tensor.shape(), + input_shape, + "stack requires matching shapes" + ); + tensor.value.unsqueeze_dims::<1, OUT>([dim]).to_concrete() + }) + .collect::>(); + let value = RawTensor::cat(raw, dim); + let parents = tensors + .iter() + .map(|tensor| tensor.handle.clone()) + .collect::>(); + let parent_ids = parents.iter().map(|parent| parent.id).collect::>(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "stack")?; + let mut targets = Vec::with_capacity(parent_ids.len()); + for (index, &parent_id) in parent_ids.iter().enumerate() { + let slices: [Range; OUT] = std::array::from_fn(|axis| { + if axis == dim { + index..index + 1 + } else { + 0..gradient.shape()[axis] + } + }); + let grad = gradient.slice(slices).reshape(input_shape).to_concrete(); + targets.push(BackwardTarget { + node: parent_id, + gradient: Box::new(grad), + }); + } + Ok(targets) + }); + let id = graph.add_node( + parents.iter().map(|parent| parent.id).collect(), + Some(backward), + parents + .iter() + .any(|parent| parent.graph.requires_grad(parent.id)), + ); + Tensor { + value, + handle: NodeHandle { graph, id }, + } + } + pub fn cat(tensors: Vec, dim: usize) -> Self { + assert!(!tensors.is_empty(), "cat requires at least one tensor"); + let graph = tensors[0].handle.graph.clone(); + let raw = tensors + .iter() + .map(|tensor| tensor.value.clone()) + .collect::>(); + let value = RawTensor::cat(raw, dim); + let parents = tensors + .iter() + .map(|tensor| tensor.handle.clone()) + .collect::>(); + let parent_ids = parents.iter().map(|parent| parent.id).collect::>(); + let slices = tensors + .iter() + .scan(0usize, |offset, tensor| { + let start = *offset; + let length = tensor.shape()[dim]; + *offset += length; + Some(start..start + length) + }) + .collect::>(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "cat")?; + let mut targets = Vec::with_capacity(parent_ids.len()); + for (&parent_id, slice) in parent_ids.iter().zip(slices.iter()) { + let ranges: [Range; R] = std::array::from_fn(|axis| { + if axis == dim { + slice.clone() + } else { + 0..gradient.shape()[axis] + } + }); + targets.push(BackwardTarget { + node: parent_id, + gradient: Box::new(gradient.slice(ranges).to_concrete()), + }); + } + Ok(targets) + }); + let id = graph.add_node( + parents.iter().map(|parent| parent.id).collect(), + Some(backward), + parents + .iter() + .any(|parent| parent.graph.requires_grad(parent.id)), + ); + Tensor { + value, + handle: NodeHandle { graph, id }, + } + } +} + +fn for_each_index(limits: [usize; R], mut visitor: impl FnMut([usize; R])) { + if limits.contains(&0) { + return; + } + + let mut index = [0; R]; + loop { + visitor(index); + + let mut axis = R; + loop { + if axis == 0 { + return; + } + axis -= 1; + index[axis] += 1; + if index[axis] < limits[axis] { + break; + } + index[axis] = 0; + } + } +} + +fn restride_input_index( + specs: [StrideSpec; OUT], + output_index: [usize; OUT], +) -> [usize; R] { + let mut input_index = [0; R]; + for axis in 0..OUT { + let spec = specs[axis]; + input_index[spec.input_dim] += spec.offset + output_index[axis] * spec.multiplier; + } + input_index +} + +fn contiguous_index_from_linear( + mut linear: usize, + strides: &[usize], +) -> [usize; R] { + let mut input_index = [0; R]; + for axis in 0..R { + input_index[axis] = linear / strides[axis]; + linear %= strides[axis]; + } + input_index +} + +/// Grouped view of a restride's specs for one input axis: at most one strided +/// "position" run and one unit-stride "window" run plus a constant offset, so +/// the forward map factors per input axis as +/// `input = offset + position * step + window`. +#[derive(Clone, Copy)] +struct RestrideRuns { + offset: usize, + /// `(output_axis, step, count)` of the strided run. + position: Option<(usize, usize, usize)>, + /// `(output_axis, size)` of the unit-stride run. + window: Option<(usize, usize)>, +} + +impl RestrideRuns { + fn counts(&self) -> (usize, usize, usize) { + let (_, step, positions) = self.position.unwrap_or((0, 1, 1)); + let (_, window) = self.window.unwrap_or((0, 1)); + (positions, step, window) + } + + fn output_len(&self) -> usize { + let (positions, _, window) = self.counts(); + positions * window + } + + fn fold_len(&self) -> usize { + let (positions, step, window) = self.counts(); + (positions - 1) * step + window + } + + /// True when the runs enumerate the input axis exactly once in row-major + /// order, so the gradient maps back with a plain reshape. + fn is_reshape(&self, size: usize) -> bool { + let (positions, step, window) = self.counts(); + self.offset == 0 && self.fold_len() == size && (positions == 1 || step == window) + } +} + +fn group_restride_runs( + specs: &[StrideSpec], + base_offsets: [usize; IN], + input_shape: [usize; IN], +) -> Option<[RestrideRuns; IN]> { + let mut runs: [RestrideRuns; IN] = std::array::from_fn(|axis| RestrideRuns { + offset: base_offsets[axis], + position: None, + window: None, + }); + for (output_axis, spec) in specs.iter().enumerate() { + let axis = &mut runs[spec.input_dim]; + axis.offset += spec.offset; + if spec.size == 1 { + continue; + } + if spec.multiplier == 1 && axis.window.is_none() { + axis.window = Some((output_axis, spec.size)); + } else if spec.multiplier >= 1 && axis.position.is_none() { + axis.position = Some((output_axis, spec.multiplier, spec.size)); + } else { + return None; + } + } + runs.iter() + .zip(&input_shape) + .all(|(axis, &size)| axis.offset + axis.fold_len() <= size) + .then_some(runs) +} + +/// Backward of a position/window restride as compiled graph ops: group the +/// output axes per input axis, canonicalize their order with one permute, +/// then fold each axis with [`fold_restride_axis`]. Returns `None` for +/// restrides that do not factor into per-axis runs; those go through the +/// host-loop fallback. +fn reduce_restride_gradient( + gradient: &RawTensor, + specs: &[StrideSpec], + base_offsets: [usize; IN], + input_shape: [usize; IN], +) -> Option> { + let output_shape = gradient.shape(); + if input_shape.contains(&0) || output_shape.contains(&0) { + return None; + } + let runs = group_restride_runs(specs, base_offsets, input_shape)?; + + let mut order = Vec::with_capacity(OUT); + for axis in &runs { + if let Some((output_axis, _, _)) = axis.position { + order.push(output_axis); + } + if let Some((output_axis, _)) = axis.window { + order.push(output_axis); + } + } + // Size-1 output axes reshape away wherever they sit, so only the run axes + // decide whether the gradient needs a permute into canonical order. + let canonical = if order.windows(2).all(|pair| pair[0] < pair[1]) { + gradient.clone() + } else { + let mut in_runs = [false; OUT]; + for &axis in &order { + in_runs[axis] = true; + } + order.extend((0..OUT).filter(|&axis| !in_runs[axis])); + let permutation: [usize; OUT] = order + .as_slice() + .try_into() + .expect("every output axis appears in the permutation once"); + gradient.permute(permutation).to_concrete() + }; + + let mut flat = canonical + .reshape([output_shape.iter().product()]) + .to_concrete(); + let mut after = 1usize; + for dim in (0..IN).rev() { + let size = input_shape[dim]; + if !runs[dim].is_reshape(size) { + let before = runs[..dim].iter().map(RestrideRuns::output_len).product(); + let (positions, step, window) = runs[dim].counts(); + flat = fold_restride_axis( + flat, + before, + positions, + step, + window, + runs[dim].offset, + size, + after, + ); + } + after *= size; + } + Some(flat.reshape(input_shape).to_concrete()) +} + +/// Scatter-add one folded axis of the gradient with padded views and a +/// reduce: `out[offset + p*step + w] += g[.., p, w, ..]` for every position +/// `p` and window element `w`, with the surrounding axes flattened into +/// `before` and `after` batch extents. +#[allow(clippy::too_many_arguments)] +fn fold_restride_axis( + flat: RawTensor<1, f32>, + before: usize, + positions: usize, + step: usize, + window: usize, + offset: usize, + size: usize, + after: usize, +) -> RawTensor<1, f32> { + let fold_len = (positions - 1) * step + window; + let block = flat + .reshape([before, positions, window, after]) + .to_concrete(); + let folded: RawTensor<3, f32> = if positions == 1 { + block.reshape([before, window, after]).to_concrete() + } else if step >= window { + // Injective: interleave the windows with zeros and trim the overhang. + block + .pad_with_zeros(2, 0, step - window) + .reshape([before, positions * step, after]) + .to_concrete() + .narrow(1usize, 0, fold_len) + .to_concrete() + } else { + // Overlapping: reverse the window axis (`u = window - 1 - w`), + // right-pad each window row to `step * window` elements and left-pad + // `(window - 1) * window` zeros; the affine view + // `f(v, u) = v*window + u*(window + 1)` then reads `g[p, w]` exactly + // when `p*step + w == v` and a zero cell otherwise, so one reduce + // over `u` folds every overlapping window. + let reversed: Vec = (0..window as u32).rev().collect(); + let indices = RawTensor::from_slice(&block.device(), [window], &reversed); + block + .index_select(2, &indices) + .pad_with_zeros(2, 0, (step - 1) * window) + .reshape([before, positions * step * window, after]) + .to_concrete() + .pad_with_zeros(1, (window - 1) * window, window * (window - step)) + .restride([ + StrideSpec::dim(0, before), + StrideSpec::dim_with(1, fold_len, window), + StrideSpec::dim_with(1, window, window + 1), + StrideSpec::dim(2, after), + ]) + .to_concrete() + .sum::<3>(2) + }; + folded + .pad_with_zeros(1, offset, size - offset - fold_len) + .reshape([before * size * after]) + .to_concrete() +} + +/// Express a layout over a contiguous input as per-output-axis stride specs +/// plus per-input-axis base offsets. Returns `None` when a stride does not +/// decompose into a single input dimension or when an axis' reach could carry +/// into the next dimension, where per-axis factoring would diverge from the +/// layout's linear indexing. +fn layout_restride_specs( + layout: &Layout, + input_shape: [usize; IN], +) -> Option<(Vec, [usize; IN])> { + if input_shape.contains(&0) { + return None; + } + let input_strides = Layout::continuous_strides(&input_shape); + let mut reach = [0usize; IN]; + let mut specs = Vec::with_capacity(layout.rank()); + for (&size, &stride) in layout.shape().iter().zip(layout.strides()) { + if size == 1 || stride == 0 { + specs.push(StrideSpec::dim_with(0, size, 0)); + continue; + } + let dim = (0..IN).find(|&dim| input_strides[dim] <= stride)?; + if stride % input_strides[dim] != 0 { + return None; + } + let multiplier = stride / input_strides[dim]; + reach[dim] += (size - 1) * multiplier; + specs.push(StrideSpec::dim_with(dim, size, multiplier)); + } + let mut offsets = [0usize; IN]; + let mut offset = layout.offset(); + for dim in 0..IN { + offsets[dim] = offset / input_strides[dim]; + offset %= input_strides[dim]; + reach[dim] += offsets[dim]; + } + reach + .iter() + .zip(&input_shape) + .all(|(&reach, &size)| reach < size) + .then_some((specs, offsets)) +} + +/// Host-loop fallback for the restride patterns [`reduce_restride_gradient`] +/// cannot factor into per-axis runs. +fn scatter_restride_gradient( + gradient: &RawTensor, + output_shape: [usize; OUT], + input_shape: [usize; IN], + input_index: impl Fn([usize; OUT]) -> [usize; IN], +) -> RawTensor { + let mut input_gradient = RawTensor::zeros(&gradient.device(), input_shape); + for_each_index(output_shape, |output_index| { + let input_index = input_index(output_index); + let output_slices: [Range; OUT] = + std::array::from_fn(|axis| output_index[axis]..output_index[axis] + 1); + let patch = gradient.slice(output_slices).reshape([1; IN]).to_concrete(); + let target: [Range; IN] = + std::array::from_fn(|axis| input_index[axis]..input_index[axis] + 1); + let current = input_gradient.slice(target.clone()).to_concrete(); + let updated = (current + patch).to_concrete(); + input_gradient = input_gradient.slice_assign(target, &updated).to_concrete(); + }); + input_gradient +} + +fn reduce_broadcast_gradient( + gradient: RawTensor, + input_shape: [usize; IN], +) -> Result> { + let output_shape = gradient.shape(); + let mut aligned_input_shape = [1usize; OUT]; + for axis in 0..IN { + aligned_input_shape[OUT - IN + axis] = input_shape[axis]; + } + + for axis in 0..OUT { + let output_dim = output_shape[axis]; + let input_dim = aligned_input_shape[axis]; + if input_dim != 1 && input_dim != output_dim { + return Err(Error::msg("incompatible broadcast gradient shape")); + } + } + + if aligned_input_shape == output_shape { + if IN == OUT { + return Ok(Box::new(gradient)); + } + return Ok(Box::new(gradient.reshape(input_shape).to_concrete())); + } + + // Sum the axes the forward broadcast expanded, one compiled reduce per axis. + let mut remaining = output_shape; + let mut flat = gradient + .reshape([output_shape.iter().product()]) + .to_concrete(); + for axis in 0..OUT { + if aligned_input_shape[axis] != 1 || remaining[axis] == 1 { + continue; + } + let before: usize = remaining[..axis].iter().product(); + let after: usize = remaining[axis + 1..].iter().product(); + flat = flat + .reshape([before, remaining[axis], after]) + .to_concrete() + .sum::<2>(1) + .reshape([before * after]) + .to_concrete(); + remaining[axis] = 1; + } + Ok(Box::new(flat.reshape(input_shape).to_concrete())) +} diff --git a/fusor-ml/fusor/src/composite/comparison.rs b/fusor-ml/fusor/src/composite/comparison.rs index 95ce32db3..f6232f324 100644 --- a/fusor-ml/fusor/src/composite/comparison.rs +++ b/fusor-ml/fusor/src/composite/comparison.rs @@ -27,7 +27,7 @@ macro_rules! scalar_cmp { }; } -/// Emit a tensor-tensor comparison method that runs on CPU only. +/// Emit a tensor-tensor comparison method that dispatches CPU/GPU. macro_rules! tensor_cmp { ($(#[$meta:meta])* $method:ident, $op:ident, $cpu_method:ident) => { $(#[$meta])* @@ -36,7 +36,11 @@ macro_rules! tensor_cmp { B2: Fusion, $op: SimdBinaryOp, { - self.dispatch_cpu_only_pair(rhs, |a, b| a.as_ref().$cpu_method(b.as_ref()).to_concrete()) + self.dispatch_pair_concrete( + rhs, + |a, b| a.as_ref().$cpu_method(b.as_ref()).to_concrete(), + |a, b| a.$method(b), + ) } }; } @@ -63,7 +67,6 @@ where /// Element-wise equality comparison between two tensors. /// /// Returns 1.0 where elements are equal, 0.0 otherwise. - /// Note: GPU comparison is only available for CPU tensors at this time. eq_tensor, EqOp, eq ); @@ -71,7 +74,6 @@ where /// Element-wise inequality comparison between two tensors. /// /// Returns 1.0 where elements are not equal, 0.0 otherwise. - /// Note: GPU comparison is only available for CPU tensors at this time. ne_tensor, NeOp, ne ); @@ -79,7 +81,6 @@ where /// Element-wise less-than comparison between two tensors. /// /// Returns 1.0 where self < rhs, 0.0 otherwise. - /// Note: GPU comparison is only available for CPU tensors at this time. lt_tensor, LtOp, lt ); @@ -87,7 +88,6 @@ where /// Element-wise less-than-or-equal comparison between two tensors. /// /// Returns 1.0 where self <= rhs, 0.0 otherwise. - /// Note: GPU comparison is only available for CPU tensors at this time. lte_tensor, LteOp, lte ); @@ -95,7 +95,6 @@ where /// Element-wise greater-than comparison between two tensors. /// /// Returns 1.0 where self > rhs, 0.0 otherwise. - /// Note: GPU comparison is only available for CPU tensors at this time. gt_tensor, GtOp, gt ); @@ -103,7 +102,6 @@ where /// Element-wise greater-than-or-equal comparison between two tensors. /// /// Returns 1.0 where self >= rhs, 0.0 otherwise. - /// Note: GPU comparison is only available for CPU tensors at this time. gte_tensor, GteOp, gte ); diff --git a/fusor-ml/fusor/src/composite/index.rs b/fusor-ml/fusor/src/composite/index.rs index 2018ca101..1f2d7a82f 100644 --- a/fusor-ml/fusor/src/composite/index.rs +++ b/fusor-ml/fusor/src/composite/index.rs @@ -48,7 +48,7 @@ impl From for IndexOp { } impl IndexOp { - fn to_range(&self, dim_size: usize) -> Range { + pub(crate) fn to_range(&self, dim_size: usize) -> Range { match self { IndexOp::Full => 0..dim_size, IndexOp::Range(r) => r.clone(), @@ -58,7 +58,7 @@ impl IndexOp { } } - fn removes_dim(&self) -> bool { + pub(crate) fn removes_dim(&self) -> bool { matches!(self, IndexOp::Index(_)) } } @@ -88,7 +88,7 @@ where } } -fn removed_dim(removes: [bool; R]) -> usize { +pub(crate) fn removed_dim(removes: [bool; R]) -> usize { let num_removes = removes.iter().filter(|&&removed| removed).count(); assert!( num_removes == 1, diff --git a/fusor-ml/fusor/src/gpu.rs b/fusor-ml/fusor/src/gpu.rs index 615368c84..45ba65498 100644 --- a/fusor-ml/fusor/src/gpu.rs +++ b/fusor-ml/fusor/src/gpu.rs @@ -334,6 +334,36 @@ impl Tensor { Tensor::from_core(self.inner.mte::(rhs)) } + #[inline] + pub fn eq_tensor(&self, other: &Self) -> Tensor { + Tensor::from_core(self.inner.eq_tensor::(other.as_core())) + } + + #[inline] + pub fn ne_tensor(&self, other: &Self) -> Tensor { + Tensor::from_core(self.inner.ne_tensor::(other.as_core())) + } + + #[inline] + pub fn lt_tensor(&self, other: &Self) -> Tensor { + Tensor::from_core(self.inner.lt_tensor::(other.as_core())) + } + + #[inline] + pub fn lte_tensor(&self, other: &Self) -> Tensor { + Tensor::from_core(self.inner.lte_tensor::(other.as_core())) + } + + #[inline] + pub fn gt_tensor(&self, other: &Self) -> Tensor { + Tensor::from_core(self.inner.gt_tensor::(other.as_core())) + } + + #[inline] + pub fn gte_tensor(&self, other: &Self) -> Tensor { + Tensor::from_core(self.inner.gte_tensor::(other.as_core())) + } + #[inline] pub fn cast(&self) -> Tensor where diff --git a/fusor-ml/fusor/src/layers/conv.rs b/fusor-ml/fusor/src/layers/conv.rs index 9786fd554..2d1f826e3 100644 --- a/fusor-ml/fusor/src/layers/conv.rs +++ b/fusor-ml/fusor/src/layers/conv.rs @@ -98,6 +98,16 @@ where pub fn out_channels(&self) -> usize { self.out_channels } + + /// Get the weight tensor. + pub fn weight(&self) -> &Tensor> { + &self.weight + } + + /// Get the bias tensor if present. + pub fn bias(&self) -> Option<&Tensor<1, D, Concrete>> { + self.bias.as_ref() + } } impl ConvNd diff --git a/fusor-ml/fusor/src/layers/conv1d.rs b/fusor-ml/fusor/src/layers/conv1d.rs index 6e4d90cc9..daac3b57f 100644 --- a/fusor-ml/fusor/src/layers/conv1d.rs +++ b/fusor-ml/fusor/src/layers/conv1d.rs @@ -122,4 +122,14 @@ where pub fn kernel_size(&self) -> usize { self.kernel_size } + + /// Get the weight tensor. + pub fn weight(&self) -> &Tensor<3, D, Concrete> { + &self.weight + } + + /// Get the bias tensor if present. + pub fn bias(&self) -> Option<&Tensor<1, D, Concrete>> { + self.bias.as_ref() + } } diff --git a/fusor-ml/fusor/src/layers/embedding.rs b/fusor-ml/fusor/src/layers/embedding.rs index 73c2f3bf0..48f7a9dd6 100644 --- a/fusor-ml/fusor/src/layers/embedding.rs +++ b/fusor-ml/fusor/src/layers/embedding.rs @@ -115,6 +115,11 @@ impl Embedding { .expect("dense embeddings unavailable for this embedding table") } + /// Get the dense embedding table, if this layer holds one. + pub fn dense_embeddings(&self) -> Option<&Tensor<2, T>> { + self.embeddings.as_ref() + } + /// Get the number of embeddings. pub fn num_embeddings(&self) -> usize { self.num_embeddings diff --git a/fusor-ml/fusor/src/lib.rs b/fusor-ml/fusor/src/lib.rs index 1d3360641..9c131ba31 100644 --- a/fusor-ml/fusor/src/lib.rs +++ b/fusor-ml/fusor/src/lib.rs @@ -11,6 +11,7 @@ #[cfg(not(any(feature = "cpu", feature = "gpu")))] compile_error!("fusor requires at least one backend feature: `cpu` or `gpu`."); +pub mod autograd; pub mod cache; mod composite; mod cpu;