diff --git a/Cargo.lock b/Cargo.lock index 13ce31fdb889..56de753bb7a4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2317,6 +2317,7 @@ dependencies = [ "hashbrown 0.15.5", "ndarray", "num-complex", + "num-traits", "rustworkx-core", "thiserror 2.0.19", ] diff --git a/crates/providers/Cargo.toml b/crates/providers/Cargo.toml index 0d397fc6f360..d9b98284e533 100644 --- a/crates/providers/Cargo.toml +++ b/crates/providers/Cargo.toml @@ -8,6 +8,9 @@ license.workspace = true [lib] name = "qiskit_providers" +[dev-dependencies] +num-traits.workspace = true# = "0.2" + [dependencies] rustworkx-core.workspace = true num-complex.workspace = true @@ -26,4 +29,4 @@ workspace = true [dependencies.ndarray] workspace = true -features = ["rayon", "approx"] \ No newline at end of file +features = ["rayon", "approx"] diff --git a/crates/providers/src/lib.rs b/crates/providers/src/lib.rs index 0312449be42b..a1d36395d417 100644 --- a/crates/providers/src/lib.rs +++ b/crates/providers/src/lib.rs @@ -11,6 +11,7 @@ // that they have been altered from the originals. mod data_tree; +pub mod math_nodes; mod program_node; mod store; pub mod tensor; diff --git a/crates/providers/src/math_nodes/binary.rs b/crates/providers/src/math_nodes/binary.rs new file mode 100644 index 000000000000..e56a69e9b3f9 --- /dev/null +++ b/crates/providers/src/math_nodes/binary.rs @@ -0,0 +1,269 @@ +// This code is part of Qiskit. +// +// (C) Copyright IBM 2026 +// +// This code is licensed under the Apache License, Version 2.0. You may +// obtain a copy of this license in the LICENSE.txt file in the root directory +// of this source tree or at https://www.apache.org/licenses/LICENSE-2.0. +// +// Any modifications or derivative works of this code must retain this +// copyright notice, and modified files need to carry a notice indicating +// that they have been altered from the originals. + +use crate::data_tree::DataTree; +use crate::program_node::ProgramNode; +use crate::tensor::{DTypeLike, Tensor, TensorType, promotion}; +use crate::unpack_tensor_args; +use std::sync::LazyLock; + +/// Shared input type spec for all elementwise binary nodes: two broadcastable tensors `x` and `y`. +static INPUT_TYPES: LazyLock> = LazyLock::new(|| { + let mut types = DataTree::with_capacity(2); + types.insert_leaf( + "x", + TensorType { + dtype: DTypeLike::Var("x".into()), + shape: vec![], + broadcastable: true, + }, + ); + types.insert_leaf( + "y", + TensorType { + dtype: DTypeLike::Var("y".into()), + shape: vec![], + broadcastable: true, + }, + ); + types +}); + +/// Shared output type spec for all elementwise binary nodes: a single tensor of the promoted dtype. +static OUTPUT_TYPES: LazyLock> = LazyLock::new(|| { + DataTree::new_leaf(TensorType { + dtype: DTypeLike::Promotion( + vec![DTypeLike::Var("x".into()), DTypeLike::Var("y".into())].into(), + ), + shape: vec![], + broadcastable: true, + }) +}); + +/// Generate a [`ProgramNode`] struct for an elementwise binary operation. +macro_rules! elementwise_binary_node { + ($name:ident, $node_name:literal, $call_fn:expr) => { + #[doc = concat!("Elementwise `", $node_name, "` of two broadcastable tensors.")] + pub struct $name; + + impl ProgramNode for $name { + type CallError = super::MathNodeError; + + fn name(&self) -> &str { + $node_name + } + fn namespace(&self) -> &str { + "qiskit" + } + fn input_types(&self) -> &DataTree { + &INPUT_TYPES + } + fn output_types(&self) -> &DataTree { + &OUTPUT_TYPES + } + fn implements_call(&self) -> bool { + true + } + fn call_flat(&self, args: &[Tensor]) -> Result, Self::CallError> { + unpack_tensor_args!(args, [x, y]); + let out_dtype = promotion(x.dtype(), y.dtype()); + let x = x.clone().cast(out_dtype); + let y = y.clone().cast(out_dtype); + Ok(vec![$call_fn(&x, &y)?]) + } + } + }; +} + +elementwise_binary_node!(Add, "add", Tensor::add_tensor); +elementwise_binary_node!(Subtract, "subtract", Tensor::sub_tensor); +elementwise_binary_node!(Multiply, "multiply", Tensor::mul_tensor); +elementwise_binary_node!(Divide, "divide", Tensor::div_tensor); +elementwise_binary_node!(Remainder, "remainder", Tensor::rem_tensor); +elementwise_binary_node!(Power, "power", Tensor::pow); + +#[cfg(test)] +mod tests { + use super::*; + use crate::math_nodes::MathNodeError; + use crate::program_node::{CallError, CallInputError, ProgramNodeExt}; + use crate::tensor::{DType, Tensor}; + + #[test] + fn test_add_same_dtype() { + let result = Add + .call_flat(&[ + Tensor::from([1.0_f64, 2.0, 3.0]), + Tensor::from([4.0_f64, 5.0, 6.0]), + ]) + .unwrap(); + assert_eq!(result.len(), 1); + let Tensor::F64(arr) = &result[0] else { + panic!("expected f64") + }; + assert_eq!(arr.as_slice().unwrap(), &[5.0, 7.0, 9.0]); + } + + #[test] + fn test_add_promotes_dtype() { + let result = Add + .call_flat(&[Tensor::from([1.0_f32, 2.0]), Tensor::from([3.0_f64, 4.0])]) + .unwrap(); + assert_eq!(result[0].dtype(), DType::F64); + let Tensor::F64(arr) = &result[0] else { + panic!("expected f64") + }; + assert_eq!(arr.as_slice().unwrap(), &[4.0, 6.0]); + } + + #[test] + fn test_add_broadcasts_2d_with_1d() { + use ndarray::arr2; + let x = Tensor::F64( + arr2(&[[1.0_f64, 2.0, 3.0], [4.0, 5.0, 6.0]]) + .into_dyn() + .into_shared(), + ); + let y = Tensor::from([10.0_f64, 20.0, 30.0]); + let result = Add.call_flat(&[x, y]).unwrap(); + let Tensor::F64(arr) = &result[0] else { + panic!("expected f64") + }; + let expected = arr2(&[[11.0_f64, 22.0, 33.0], [14.0, 25.0, 36.0]]) + .into_dyn() + .into_shared(); + assert_eq!(arr, &expected); + } + + #[test] + fn test_subtract() { + let result = Subtract + .call_flat(&[ + Tensor::from([5.0_f64, 6.0, 7.0]), + Tensor::from([1.0_f64, 2.0, 3.0]), + ]) + .unwrap(); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + assert_eq!(arr.as_slice().unwrap(), &[4.0, 4.0, 4.0]); + } + + #[test] + fn test_multiply() { + let result = Multiply + .call_flat(&[ + Tensor::from([2.0_f64, 3.0, 4.0]), + Tensor::from([10.0_f64, 10.0, 10.0]), + ]) + .unwrap(); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + assert_eq!(arr.as_slice().unwrap(), &[20.0, 30.0, 40.0]); + } + + #[test] + fn test_divide() { + let result = Divide + .call_flat(&[ + Tensor::from([10.0_f64, 9.0, 8.0]), + Tensor::from([2.0_f64, 3.0, 4.0]), + ]) + .unwrap(); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + assert_eq!(arr.as_slice().unwrap(), &[5.0, 3.0, 2.0]); + } + + #[test] + fn test_remainder() { + let result = Remainder + .call_flat(&[ + Tensor::from([7.0_f64, 8.0, 9.0]), + Tensor::from([3.0_f64, 3.0, 3.0]), + ]) + .unwrap(); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + assert_eq!(arr.as_slice().unwrap(), &[1.0, 2.0, 0.0]); + } + + #[test] + fn test_power() { + let result = Power + .call_flat(&[ + Tensor::from([2.0_f64, 3.0, 4.0]), + Tensor::from([3.0_f64, 2.0, 1.0]), + ]) + .unwrap(); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + for (a, b) in arr.as_slice().unwrap().iter().zip(&[8.0_f64, 9.0, 4.0]) { + assert!(approx::abs_diff_eq!(a, b, epsilon = 1e-12)); + } + } + + #[test] + fn test_call_missing_input_errors() { + let mut tree = DataTree::new(); + tree.insert_leaf("x", Tensor::from([1.0_f64])); + let err = Add.call(&tree).unwrap_err(); + assert!(matches!( + err, + CallError::::Input(CallInputError::MissingInput { + ref key, + }) if key == "y" + )); + } + + #[test] + fn test_call_branch_where_leaf_expected_errors() { + let mut tree = DataTree::new(); + tree.insert_leaf("x", Tensor::from([1.0_f64])); + tree.insert_branch("y", DataTree::new()); + let err = Add.call(&tree).unwrap_err(); + assert!(matches!( + err, + CallError::::Input(CallInputError::ExpectedLeaf { + ref key, + }) if key == "y" + )); + } + + #[test] + fn test_add_call_end_to_end() { + let mut tree = DataTree::new(); + tree.insert_leaf("x", Tensor::from([1.0_f64, 2.0, 3.0])); + tree.insert_leaf("y", Tensor::from([4.0_f64, 5.0, 6.0])); + let result = Add.call(&tree).unwrap(); + let Tensor::F64(arr) = result.unwrap_leaf() else { + panic!("expected f64") + }; + assert_eq!(arr.as_slice().unwrap(), &[5.0, 7.0, 9.0]); + } + + #[test] + fn test_add_wrong_arity_errors() { + let err = Add.call_flat(&[Tensor::from([1.0_f64])]).unwrap_err(); + assert_eq!( + err, + MathNodeError::Input(CallInputError::WrongArity { + expected: 2, + actual: 1, + }) + ); + } +} diff --git a/crates/providers/src/math_nodes/bitwise.rs b/crates/providers/src/math_nodes/bitwise.rs new file mode 100644 index 000000000000..ed7c5ab53e57 --- /dev/null +++ b/crates/providers/src/math_nodes/bitwise.rs @@ -0,0 +1,344 @@ +// This code is part of Qiskit. +// +// (C) Copyright IBM 2026 +// +// This code is licensed under the Apache License, Version 2.0. You may +// obtain a copy of this license in the LICENSE.txt file in the root directory +// of this source tree or at https://www.apache.org/licenses/LICENSE-2.0. +// +// Any modifications or derivative works of this code must retain this +// copyright notice, and modified files need to carry a notice indicating +// that they have been altered from the originals. + +use crate::data_tree::DataTree; +use crate::program_node::{CallInputError, ProgramNode}; +use crate::tensor::{DType, DTypeLike, Tensor, TensorType, broadcast_shape}; +use crate::unpack_tensor_args; +use ndarray::Axis; +use std::sync::LazyLock; + +/// Shared input type spec for binary bitwise nodes +static INPUT_TYPES: LazyLock> = LazyLock::new(|| { + let mut types = DataTree::with_capacity(2); + types.insert_leaf( + "x", + TensorType { + dtype: DTypeLike::Concrete(DType::Bit), + shape: vec![], + broadcastable: true, + }, + ); + types.insert_leaf( + "y", + TensorType { + dtype: DTypeLike::Concrete(DType::Bit), + shape: vec![], + broadcastable: true, + }, + ); + types +}); + +/// A single broadcastable `Bit` leaf — used for unary inputs and all bitwise outputs. +static LEAF_TYPE: LazyLock> = LazyLock::new(|| { + DataTree::new_leaf(TensorType { + dtype: DTypeLike::Concrete(DType::Bit), + shape: vec![], + broadcastable: true, + }) +}); + +/// Construct an `UnexpectedDType` error for a slice element that did not match +/// the schema's required dtype. +fn unexpected_dtype(key: &str, actual: &Tensor) -> CallInputError { + CallInputError::UnexpectedDType { + key: key.into(), + expected: DType::Bit.to_string(), + actual: actual.dtype(), + } +} + +/// Generate a [`ProgramNode`] struct for an elementwise binary bitwise operation on `Bit` tensors. +macro_rules! bitwise_binary_node { + ($name:ident, $node_name:literal, $call_fn:expr) => { + #[doc = concat!("Elementwise `", $node_name, "` of two broadcastable `Bit` tensors.")] + pub struct $name; + + impl ProgramNode for $name { + type CallError = super::MathNodeError; + + fn name(&self) -> &str { + $node_name + } + fn namespace(&self) -> &str { + "qiskit" + } + fn input_types(&self) -> &DataTree { + &INPUT_TYPES + } + fn output_types(&self) -> &DataTree { + &LEAF_TYPE + } + fn implements_call(&self) -> bool { + true + } + fn call_flat(&self, args: &[Tensor]) -> Result, Self::CallError> { + unpack_tensor_args!(args, [x, y]); + let Tensor::Bit(x_arr) = x else { + return Err(unexpected_dtype("x", x).into()); + }; + let Tensor::Bit(y_arr) = y else { + return Err(unexpected_dtype("y", y).into()); + }; + broadcast_shape(x_arr.shape(), y_arr.shape())?; + Ok(vec![Tensor::Bit($call_fn(x_arr, y_arr).into_shared())]) + } + } + }; +} + +bitwise_binary_node!(BitwiseAnd, "bitwise_and", |x, y| x & y); +bitwise_binary_node!(BitwiseOr, "bitwise_or", |x, y| x | y); +bitwise_binary_node!(BitwiseXor, "bitwise_xor", |x, y| x ^ y); + +/// Elementwise bitwise NOT of a broadcastable `Bit` tensor. +pub struct BitwiseNot; + +impl ProgramNode for BitwiseNot { + type CallError = super::MathNodeError; + + fn name(&self) -> &str { + "bitwise_not" + } + fn namespace(&self) -> &str { + "qiskit" + } + fn input_types(&self) -> &DataTree { + &LEAF_TYPE + } + fn output_types(&self) -> &DataTree { + &LEAF_TYPE + } + fn implements_call(&self) -> bool { + true + } + fn call_flat(&self, args: &[Tensor]) -> Result, Self::CallError> { + unpack_tensor_args!(args, [x]); + let Tensor::Bit(arr) = x else { + return Err(unexpected_dtype("", x).into()); + }; + Ok(vec![Tensor::Bit(arr.mapv(|b| b ^ 1).into_shared())]) + } +} + +/// XOR-reduction of a `Bit` tensor along a specified axis, removing that axis. +/// +/// The parity of a sequence of bits is 1 if an odd number of bits are 1, and 0 otherwise, +/// which is equivalent to XOR-folding the sequence. The output has one fewer dimension than +/// the input, with the reduction axis removed. +pub struct Parity { + axis: usize, +} + +impl Parity { + /// Construct a `Parity` node that reduces along `axis`. + pub fn new(axis: usize) -> Self { + Self { axis } + } +} + +impl ProgramNode for Parity { + type CallError = super::MathNodeError; + + fn name(&self) -> &str { + "parity" + } + fn namespace(&self) -> &str { + "qiskit" + } + fn input_types(&self) -> &DataTree { + &LEAF_TYPE + } + fn output_types(&self) -> &DataTree { + &LEAF_TYPE + } + fn implements_call(&self) -> bool { + true + } + fn call_flat(&self, args: &[Tensor]) -> Result, Self::CallError> { + unpack_tensor_args!(args, [x]); + super::check_axis(self.axis, x.shape().len())?; + let Tensor::Bit(arr) = x else { + return Err(unexpected_dtype("", x).into()); + }; + Ok(vec![Tensor::Bit( + arr.fold_axis(Axis(self.axis), 0u8, |&acc, &b| acc ^ b) + .into_shared(), + )]) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::math_nodes::MathNodeError; + use crate::program_node::{CallError, CallInputError, ProgramNodeExt}; + use ndarray::{arr1, arr2}; + + fn bit(data: &[u8]) -> Tensor { + Tensor::Bit(arr1(data).into_dyn().into_shared()) + } + + #[test] + fn test_bitwise_and() { + let result = BitwiseAnd + .call_flat(&[bit(&[1, 0, 1, 1]), bit(&[1, 1, 0, 1])]) + .unwrap(); + let Tensor::Bit(arr) = &result[0] else { + panic!("expected Bit leaf"); + }; + assert_eq!(arr.as_slice().unwrap(), &[1, 0, 0, 1]); + } + + #[test] + fn test_bitwise_or() { + let result = BitwiseOr + .call_flat(&[bit(&[1, 0, 1, 0]), bit(&[0, 1, 0, 1])]) + .unwrap(); + let Tensor::Bit(arr) = &result[0] else { + panic!("expected Bit leaf"); + }; + assert_eq!(arr.as_slice().unwrap(), &[1, 1, 1, 1]); + } + + #[test] + fn test_bitwise_xor() { + let result = BitwiseXor + .call_flat(&[bit(&[1, 0, 1, 1]), bit(&[1, 1, 0, 1])]) + .unwrap(); + let Tensor::Bit(arr) = &result[0] else { + panic!("expected Bit leaf"); + }; + assert_eq!(arr.as_slice().unwrap(), &[0, 1, 1, 0]); + } + + #[test] + fn test_bitwise_and_broadcasts() { + // shape [3] & shape [1] -> shape [3] + let result = BitwiseAnd.call_flat(&[bit(&[1, 0, 1]), bit(&[1])]).unwrap(); + let Tensor::Bit(arr) = &result[0] else { + panic!("expected Bit leaf"); + }; + assert_eq!(arr.as_slice().unwrap(), &[1, 0, 1]); + } + + #[test] + fn test_bitwise_not() { + let result = BitwiseNot.call_flat(&[bit(&[1, 0, 1, 0])]).unwrap(); + let Tensor::Bit(arr) = &result[0] else { + panic!("expected Bit leaf"); + }; + assert_eq!(arr.as_slice().unwrap(), &[0, 1, 0, 1]); + } + + #[test] + fn test_parity_axis0() { + // [[1,0,1],[0,1,1],[0,0,0]] axis 0 → [1, 1, 0] + let x = Tensor::Bit( + arr2(&[[1u8, 0, 1], [0, 1, 1], [0, 0, 0]]) + .into_dyn() + .into_shared(), + ); + let result = Parity::new(0).call_flat(&[x]).unwrap(); + let Tensor::Bit(arr) = &result[0] else { + panic!("expected Bit leaf"); + }; + assert_eq!(arr.as_slice().unwrap(), &[1, 1, 0]); + } + + #[test] + fn test_bitwise_and_wrong_dtype_errors() { + let err = BitwiseAnd + .call_flat(&[Tensor::from([1.0_f64]), bit(&[1])]) + .unwrap_err(); + assert_eq!( + err, + MathNodeError::Input(CallInputError::UnexpectedDType { + key: "x".to_string(), + expected: "Bit".to_string(), + actual: DType::F64, + }) + ); + } + + #[test] + fn test_bitwise_and_wrong_arity_errors() { + let err = BitwiseAnd.call_flat(&[bit(&[1, 0])]).unwrap_err(); + assert_eq!( + err, + MathNodeError::Input(CallInputError::WrongArity { + expected: 2, + actual: 1, + }) + ); + } + + #[test] + fn test_bitwise_not_wrong_arity_errors() { + let err = BitwiseNot + .call_flat(&[bit(&[1, 0]), bit(&[0, 1])]) + .unwrap_err(); + assert_eq!( + err, + MathNodeError::Input(CallInputError::WrongArity { + expected: 1, + actual: 2, + }) + ); + } + + #[test] + fn test_bitwise_and_shape_mismatch_errors() { + let err = BitwiseAnd + .call_flat(&[bit(&[1, 0, 1]), bit(&[1, 0, 1, 1])]) + .unwrap_err(); + assert_eq!( + err, + MathNodeError::Tensor(crate::tensor::TensorError::ShapeMismatch { + lhs: vec![3], + rhs: vec![4], + }) + ); + } + + #[test] + fn test_call_branch_where_leaf_expected_errors() { + let mut tree = DataTree::new(); + tree.insert_leaf("x", bit(&[1, 0])); + let err = BitwiseNot.call(&tree).unwrap_err(); + assert!(matches!( + err, + CallError::::Input(CallInputError::ExpectedLeaf { + ref key, + }) if key.is_empty() + )); + } + + #[test] + fn test_bitwise_and_call_end_to_end() { + let mut tree = DataTree::new(); + tree.insert_leaf("x", bit(&[1, 0, 1, 1])); + tree.insert_leaf("y", bit(&[1, 1, 0, 1])); + let result = BitwiseAnd.call(&tree).unwrap(); + let Tensor::Bit(arr) = result.unwrap_leaf() else { + panic!("expected Bit leaf"); + }; + assert_eq!(arr.as_slice().unwrap(), &[1, 0, 0, 1]); + } + + #[test] + fn test_parity_axis_out_of_bounds_errors() { + let err = Parity::new(1).call_flat(&[bit(&[1, 0, 1])]).unwrap_err(); + assert_eq!(err, MathNodeError::InvalidAxis { axis: 1, ndim: 1 }); + } +} diff --git a/crates/providers/src/math_nodes/mod.rs b/crates/providers/src/math_nodes/mod.rs new file mode 100644 index 000000000000..60670e58c969 --- /dev/null +++ b/crates/providers/src/math_nodes/mod.rs @@ -0,0 +1,45 @@ +// This code is part of Qiskit. +// +// (C) Copyright IBM 2026 +// +// This code is licensed under the Apache License, Version 2.0. You may +// obtain a copy of this license in the LICENSE.txt file in the root directory +// of this source tree or at https://www.apache.org/licenses/LICENSE-2.0. +// +// Any modifications or derivative works of this code must retain this +// copyright notice, and modified files need to carry a notice indicating +// that they have been altered from the originals. + +mod binary; +mod bitwise; +mod reduction; + +pub use binary::*; +pub use bitwise::*; +pub use reduction::*; + +use crate::program_node::CallInputError; +use crate::tensor::TensorError; +use thiserror::Error; + +/// Errors returned by [`crate::program_node::ProgramNode`] implementations in this module. +#[derive(Debug, Clone, PartialEq, Eq, Error)] +pub enum MathNodeError { + /// The input tree did not match the contract declared by `input_types`. + #[error(transparent)] + Input(#[from] CallInputError), + /// A tensor operation failed (dtype or shape mismatch). + #[error(transparent)] + Tensor(#[from] TensorError), + /// The requested axis was out of bounds for the tensor's number of dimensions. + #[error("axis {axis} is out of bounds for tensor with {ndim} dimension(s)")] + InvalidAxis { axis: usize, ndim: usize }, +} + +/// Validate that `axis` is a valid axis index for a tensor with `ndim` dimensions. +pub(crate) fn check_axis(axis: usize, ndim: usize) -> Result<(), MathNodeError> { + if axis >= ndim { + return Err(MathNodeError::InvalidAxis { axis, ndim }); + } + Ok(()) +} diff --git a/crates/providers/src/math_nodes/reduction.rs b/crates/providers/src/math_nodes/reduction.rs new file mode 100644 index 000000000000..6ef70a77e903 --- /dev/null +++ b/crates/providers/src/math_nodes/reduction.rs @@ -0,0 +1,873 @@ +// This code is part of Qiskit. +// +// (C) Copyright IBM 2026 +// +// This code is licensed under the Apache License, Version 2.0. You may +// obtain a copy of this license in the LICENSE.txt file in the root directory +// of this source tree or at https://www.apache.org/licenses/LICENSE-2.0. +// +// Any modifications or derivative works of this code must retain this +// copyright notice, and modified files need to carry a notice indicating +// that they have been altered from the originals. + +use crate::data_tree::DataTree; +use crate::program_node::ProgramNode; +use crate::tensor::{DType, DTypeLike, Tensor, TensorType}; +use crate::unpack_tensor_args; +use ndarray::{ArrayBase, ArrayD, Axis, Data, IxDyn, NdFloat, Zip}; +use num_complex::Complex; +use std::sync::LazyLock; + +/// Shared input type spec for reduction nodes: a single broadcastable tensor of any dtype. +static INPUT_TYPES: LazyLock> = LazyLock::new(|| { + DataTree::new_leaf(TensorType { + dtype: DTypeLike::Var("x".into()), + shape: vec![], + broadcastable: true, + }) +}); + +/// Shared output type spec for reduction nodes: a single broadcastable tensor of any dtype. +static OUTPUT_TYPES: LazyLock> = LazyLock::new(|| { + DataTree::new_leaf(TensorType { + dtype: DTypeLike::Var("out".into()), + shape: vec![], + broadcastable: true, + }) +}); + +/// Smallest output length at which the slice traversal in [`sum_sq_deviations`] is +/// faster than the lane traversal. +/// +/// The slice traversal runs one [`Zip`] per position along the reduced axis, so its +/// fixed per-slice cost is amortized only once the output holds enough elements. The +/// crossover is therefore set by the ratio of that setup cost to the per-element +/// work, and measurement on `C64` data puts it near 16. The per-element work is +/// smaller on a target with wider vectors, so the crossover there should be larger. +/// +/// The exact value matters little, and any of 8, 16 or 32 would do. Some threshold is +/// still needed. The output of a one-dimensional reduction is a single element, and +/// the slice traversal on it is more than forty times slower. +const MIN_SLICE_OUTPUT_LEN: usize = 16; + +/// Mean of `a` along `axis`, with that axis removed. +/// +/// A zero-length axis gives a NaN mean. +fn complex_mean(a: &ArrayBase, axis: Axis) -> ArrayD> +where + A: NdFloat, + S: Data>, +{ + let n = A::from(a.len_of(axis)).expect("an axis length converts to a float"); + a.sum_axis(axis) + .mapv_into(|c| Complex::new(c.re / n, c.im / n)) +} + +/// Sum of the squared moduli of the deviations of `a` from its mean along `axis`, +/// with that axis removed. +/// +/// Peak memory is independent of the length of the reduced axis. +/// The sum is accumulated in `f64` for both `f32` and `f64` to hold the summation error +/// below the rounding error of `f32` inputs. +fn sum_sq_deviations(a: &ArrayBase, axis: Axis) -> ArrayD +where + A: NdFloat + Into, + S: Data>, +{ + let mean = complex_mean(a, axis); + // This function contains two implementations. One does a reduction over each lane, + // which is fast when lanes ar contigous in memory. The other accumulates over + // slices, which is fast when slices are contigous in memory. We inspect stride + // information to form a heuristic about which one to choose: the difference in + // speed can be as much as 40x. + let reduced_stride = a.strides()[axis.index()].unsigned_abs(); + let fastest_stride = a + .shape() + .iter() + .zip(a.strides()) + .filter_map(|(&len, &stride)| (len > 1).then_some(stride.unsigned_abs())) + .min() + .unwrap_or(1); + if mean.len() >= MIN_SLICE_OUTPUT_LEN && reduced_stride > fastest_stride { + // Perform an accumulation, one slice at a time. + let mut accumulated = ArrayD::::zeros(mean.raw_dim()); + for slice in a.axis_iter(axis) { + Zip::from(&mut accumulated) + .and(&slice) + .and(&mean) + .for_each(|total, &x, &m| *total += (x - m).norm_sqr().into()); + } + accumulated + } else { + // Perform a reduction of each lane separately. + Zip::from(a.lanes(axis)) + .and(&mean) + .map_collect(|lane, &m| lane.iter().map(|&x| (x - m).norm_sqr().into()).sum()) + } +} + +/// Mean of a tensor along a specified axis, removing that axis. +/// +/// Integer inputs are cast to `F64` before computing the mean. `F32` inputs +/// produce `F32` output; all other float and integer types produce `F64`. +/// Complex inputs (`C64`, `C128`) preserve their complex dtype. +pub struct Mean { + axis: usize, +} + +impl Mean { + /// Construct a `Mean` node that reduces along `axis`. + pub fn new(axis: usize) -> Self { + Self { axis } + } +} + +impl ProgramNode for Mean { + type CallError = super::MathNodeError; + + fn name(&self) -> &str { + "mean" + } + fn namespace(&self) -> &str { + "qiskit" + } + fn input_types(&self) -> &DataTree { + &INPUT_TYPES + } + fn output_types(&self) -> &DataTree { + &OUTPUT_TYPES + } + fn implements_call(&self) -> bool { + true + } + fn call_flat(&self, args: &[Tensor]) -> Result, Self::CallError> { + unpack_tensor_args!(args, [x]); + super::check_axis(self.axis, x.shape().len())?; + let result = match x { + Tensor::F32(a) => Tensor::F32(a.mean_axis(Axis(self.axis)).unwrap().into_shared()), + Tensor::F64(a) => Tensor::F64(a.mean_axis(Axis(self.axis)).unwrap().into_shared()), + Tensor::C64(a) => { + let n = a.shape()[self.axis] as f32; + Tensor::C64((a.sum_axis(Axis(self.axis)) / Complex::new(n, 0.0)).into_shared()) + } + Tensor::C128(a) => { + let n = a.shape()[self.axis] as f64; + Tensor::C128((a.sum_axis(Axis(self.axis)) / Complex::new(n, 0.0)).into_shared()) + } + other => { + let Tensor::F64(a) = other.clone().cast(DType::F64) else { + unreachable!("Value cast as F64 can't be another dtype") + }; + Tensor::F64(a.mean_axis(Axis(self.axis)).unwrap().into_shared()) + } + }; + Ok(vec![result]) + } +} + +/// Variance of a tensor along a specified axis, removing that axis. +/// +/// The `ddof` (delta degrees of freedom) parameter adjusts the divisor: the result +/// is divided by `n - ddof` where `n` is the number of elements along the axis. +/// Use `ddof=0` for population variance and `ddof=1` for sample variance. +/// +/// Integer inputs are cast to `F64`. `F32` produces `F32`; all other real types +/// produce `F64`. Complex inputs (`C64`, `C128`) produce real output (`F32`, `F64` +/// respectively), computed as the mean squared modulus of the deviations. +pub struct Variance { + axis: usize, + ddof: f64, +} + +impl Variance { + /// Construct a `Variance` node that reduces along `axis` with degrees-of-freedom + /// correction `ddof`. + pub fn new(axis: usize, ddof: f64) -> Self { + Self { axis, ddof } + } +} + +impl ProgramNode for Variance { + type CallError = super::MathNodeError; + + fn name(&self) -> &str { + "variance" + } + fn namespace(&self) -> &str { + "qiskit" + } + fn input_types(&self) -> &DataTree { + &INPUT_TYPES + } + fn output_types(&self) -> &DataTree { + &OUTPUT_TYPES + } + fn implements_call(&self) -> bool { + true + } + fn call_flat(&self, args: &[Tensor]) -> Result, Self::CallError> { + unpack_tensor_args!(args, [x]); + super::check_axis(self.axis, x.shape().len())?; + let result = match x { + Tensor::F32(a) => { + Tensor::F32(a.var_axis(Axis(self.axis), self.ddof as f32).into_shared()) + } + Tensor::F64(a) => Tensor::F64(a.var_axis(Axis(self.axis), self.ddof).into_shared()), + Tensor::C64(a) => { + let denom = a.shape()[self.axis] as f64 - self.ddof; + let var = sum_sq_deviations(a, Axis(self.axis)); + Tensor::F32(var.mapv(|total| (total / denom) as f32).into_shared()) + } + Tensor::C128(a) => { + let denom = a.shape()[self.axis] as f64 - self.ddof; + let var = sum_sq_deviations(a, Axis(self.axis)); + Tensor::F64(var.mapv_into(|total| total / denom).into_shared()) + } + other => { + let Tensor::F64(a) = other.clone().cast(DType::F64) else { + unreachable!("Value cast as F64 can't be another dtype") + }; + Tensor::F64(a.var_axis(Axis(self.axis), self.ddof).into_shared()) + } + }; + Ok(vec![result]) + } +} + +/// Standard deviation of a tensor along a specified axis, removing that axis. +/// +/// This is the square root of [`Variance`]. See that type for details on `ddof`, +/// output dtypes, and complex handling. +pub struct Std { + axis: usize, + ddof: f64, +} + +impl Std { + /// Construct a `Std` node that reduces along `axis` with degrees-of-freedom + /// correction `ddof`. + pub fn new(axis: usize, ddof: f64) -> Self { + Self { axis, ddof } + } +} + +impl ProgramNode for Std { + type CallError = super::MathNodeError; + + fn name(&self) -> &str { + "std" + } + fn namespace(&self) -> &str { + "qiskit" + } + fn input_types(&self) -> &DataTree { + &INPUT_TYPES + } + fn output_types(&self) -> &DataTree { + &OUTPUT_TYPES + } + fn implements_call(&self) -> bool { + true + } + fn call_flat(&self, args: &[Tensor]) -> Result, Self::CallError> { + unpack_tensor_args!(args, [x]); + super::check_axis(self.axis, x.shape().len())?; + let result = match x { + Tensor::F32(a) => { + Tensor::F32(a.std_axis(Axis(self.axis), self.ddof as f32).into_shared()) + } + Tensor::F64(a) => Tensor::F64(a.std_axis(Axis(self.axis), self.ddof).into_shared()), + Tensor::C64(a) => { + let denom = a.shape()[self.axis] as f64 - self.ddof; + let var = sum_sq_deviations(a, Axis(self.axis)); + Tensor::F32( + var.mapv(|total| (total / denom).sqrt() as f32) + .into_shared(), + ) + } + Tensor::C128(a) => { + let denom = a.shape()[self.axis] as f64 - self.ddof; + let var = sum_sq_deviations(a, Axis(self.axis)); + Tensor::F64(var.mapv_into(|total| (total / denom).sqrt()).into_shared()) + } + other => { + let Tensor::F64(a) = other.clone().cast(DType::F64) else { + unreachable!("Value cast as F64 can't be another dtype") + }; + Tensor::F64(a.std_axis(Axis(self.axis), self.ddof).into_shared()) + } + }; + Ok(vec![result]) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::math_nodes::MathNodeError; + use crate::program_node::{CallError, CallInputError, ProgramNodeExt}; + use crate::tensor::{DType, Tensor}; + use ndarray::{ArrayView, ShapeBuilder, arr2}; + use num_complex::Complex; + use num_traits::{Float, NumCast, Signed, abs, cast}; + use std::ops::Sub; + + fn approx_eq_slice<'a, T>(a: &'a [T], b: &'a [T]) + where + T: Float + NumCast + std::fmt::Display, + &'a T: Sub<&'a T>, + <&'a T as Sub>::Output: Signed + Float, + { + assert_eq!(a.len(), b.len(), "slice lengths differ"); + for (x, y) in a.iter().zip(b.iter()) { + assert!(abs(x - y) < cast(1e-10).unwrap(), "{x} != {y}"); + } + } + + // --- Mean tests --- + + #[test] + fn test_mean_f64_axis0() { + // [[1,2,3],[4,5,6]] along axis 0 → [2.5, 3.5, 4.5] + let x = Tensor::F64( + arr2(&[[1.0_f64, 2.0, 3.0], [4.0, 5.0, 6.0]]) + .into_dyn() + .into_shared(), + ); + let result = Mean::new(0).call_flat(&[x]).unwrap(); + let Tensor::F64(arr) = &result[0] else { + panic!("expected F64 leaf"); + }; + approx_eq_slice(arr.as_slice().unwrap(), &[2.5, 3.5, 4.5]); + } + + #[test] + fn test_mean_f32_axis0() { + // [[1,2,3],[4,5,6]] along axis 0 → [2.5, 3.5, 4.5] + let x = Tensor::F32( + arr2(&[[1.0_f32, 2.0, 3.0], [4.0, 5.0, 6.0]]) + .into_dyn() + .into_shared(), + ); + let result = Mean::new(0).call_flat(&[x]).unwrap(); + let Tensor::F32(arr) = &result[0] else { + panic!("expected F32 leaf"); + }; + approx_eq_slice(arr.as_slice().unwrap(), &[2.5, 3.5, 4.5]); + } + + #[test] + fn test_mean_i32_casts_to_f64() { + let x = Tensor::from([1_i32, 2, 3, 4]); + let result = Mean::new(0).call_flat(&[x]).unwrap(); + assert_eq!( + result[0].dtype(), + DType::F64, + "integer input should produce F64 mean" + ); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[2.5]); + } + + #[test] + fn test_mean_c128() { + let data: Vec> = vec![ + Complex::new(1.0, 2.0), + Complex::new(3.0, 4.0), + Complex::new(5.0, 6.0), + ]; + let x = Tensor::C128(ndarray::Array1::from(data).into_dyn().into_shared()); + let result = Mean::new(0).call_flat(&[x]).unwrap(); + let Tensor::C128(arr) = &result[0] else { + panic!("expected C128 leaf"); + }; + let v = arr.as_slice().unwrap()[0]; + assert!((v.re - 3.0).abs() < 1e-10); + assert!((v.im - 4.0).abs() < 1e-10); + } + + #[test] + fn test_mean_c64() { + let data: Vec> = vec![ + Complex::new(1.0, 2.0), + Complex::new(3.0, 4.0), + Complex::new(5.0, 6.0), + ]; + let x = Tensor::C64(ndarray::Array1::from(data).into_dyn().into_shared()); + let result = Mean::new(0).call_flat(&[x]).unwrap(); + let Tensor::C64(arr) = &result[0] else { + panic!("expected C64 leaf"); + }; + let v = arr.as_slice().unwrap()[0]; + assert!((v.re - 3.0).abs() < 1e-10); + assert!((v.im - 4.0).abs() < 1e-10); + } + + #[test] + fn test_large_strided_mean() { + let raw_array: [Complex; 48] = [ + 2.0.into(), + 4.0.into(), + 4.0.into(), + 4.0.into(), + 5.0.into(), + 5.0.into(), + 7.0.into(), + 9.0.into(), + 2.0.into(), + 4.0.into(), + 4.0.into(), + 4.0.into(), + 5.0.into(), + 5.0.into(), + 7.0.into(), + 9.0.into(), + 2.0.into(), + 4.0.into(), + 4.0.into(), + 4.0.into(), + 5.0.into(), + 7.0.into(), + 9.0.into(), + 1.0.into(), + 2.0.into(), + 4.0.into(), + 4.0.into(), + 4.0.into(), + 5.0.into(), + 5.0.into(), + 7.0.into(), + 9.0.into(), + 2.0.into(), + 4.0.into(), + 4.0.into(), + 4.0.into(), + 5.0.into(), + 5.0.into(), + 7.0.into(), + 9.0.into(), + 2.0.into(), + 4.0.into(), + 4.0.into(), + 4.0.into(), + 5.0.into(), + 7.0.into(), + 9.0.into(), + 1.0.into(), + ]; + // For shape that triggers accumulation over slice path + let strided = ArrayView::from_shape((4, 5, 4).strides((1, 4, 2)), &raw_array).unwrap(); + let x = Tensor::C128(strided.into_dyn().into_owned().into_shared()); + let result = Variance::new(0, 0.).call_flat(&[x]).unwrap(); + let Tensor::F64(arr) = &result[0] else { + panic!("Expected F64 leaf") + }; + approx_eq_slice( + arr.as_slice().unwrap(), + strided + .mapv(|x| x.re) + .var_axis(Axis(0), 0.) + .as_slice() + .unwrap(), + ); + } + + // --- Variance tests --- + + #[test] + fn test_variance_f64_ddof0() { + // [2, 4, 4, 4, 5, 5, 7, 9] — classic example, population variance = 4.0 + let x = Tensor::from([2.0_f64, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0]); + let result = Variance::new(0, 0.0).call_flat(&[x]).unwrap(); + let Tensor::F64(arr) = &result[0] else { + panic!("expected F64 leaf"); + }; + approx_eq_slice(arr.as_slice().unwrap(), &[4.0]); + } + + #[test] + fn test_variance_f64_ddof1() { + // Sample variance (ddof=1) of the same sequence + let x = Tensor::from([2.0_f64, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0]); + let result = Variance::new(0, 1.0).call_flat(&[x]).unwrap(); + let Tensor::F64(arr) = &result[0] else { + panic!("expected F64 leaf"); + }; + // sample variance = population variance * n / (n-1) = 4.0 * 8/7 + approx_eq_slice(arr.as_slice().unwrap(), &[4.0 * 8.0 / 7.0]); + } + + #[test] + fn test_variance_f32_ddof0() { + // [2, 4, 4, 4, 5, 5, 7, 9] — classic example, population variance = 4.0 + let x = Tensor::from([2.0_f32, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0]); + let result = Variance::new(0, 0.0).call_flat(&[x]).unwrap(); + let Tensor::F32(arr) = &result[0] else { + panic!("expected F32 leaf"); + }; + approx_eq_slice(arr.as_slice().unwrap(), &[4.0]); + } + + #[test] + fn test_variance_f32_ddof1() { + // Sample variance (ddof=1) of the same sequence + let x = Tensor::from([2.0_f32, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0]); + let result = Variance::new(0, 1.0).call_flat(&[x]).unwrap(); + let Tensor::F32(arr) = &result[0] else { + panic!("expected F32 leaf"); + }; + // sample variance = population variance * n / (n-1) = 4.0 * 8/7 + approx_eq_slice(arr.as_slice().unwrap(), &[4.0 * 8.0 / 7.0]); + } + + #[test] + fn test_variance_c128_returns_real() { + // [1+1i, 3+3i] — mean = 2+2i, deviations = [−1−i, 1+i], |.|^2 = [2, 2], var = 2.0 + let data: Vec> = vec![Complex::new(1.0, 1.0), Complex::new(3.0, 3.0)]; + let x = Tensor::C128(ndarray::Array1::from(data).into_dyn().into_shared()); + let result = Variance::new(0, 0.0).call_flat(&[x]).unwrap(); + assert_eq!( + result[0].dtype(), + DType::F64, + "C128 variance should return F64" + ); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[2.0]); + } + + #[test] + fn test_variance_c64_returns_real() { + // [1+1i, 3+3i] — mean = 2+2i, deviations = [−1−i, 1+i], |.|^2 = [2, 2], var = 2.0 + let data: Vec> = vec![Complex::new(1.0, 1.0), Complex::new(3.0, 3.0)]; + let x = Tensor::C64(ndarray::Array1::from(data).into_dyn().into_shared()); + let result = Variance::new(0, 0.0).call_flat(&[x]).unwrap(); + assert_eq!( + result[0].dtype(), + DType::F32, + "C64 variance should return F32" + ); + let Tensor::F32(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[2.0]); + } + + // --- Std tests --- + + #[test] + fn test_std_matches_sqrt_of_variance() { + // Verify std = sqrt(variance) numerically + let x = Tensor::from([1.0_f64, 3.0, 5.0, 7.0, 9.0]); + let var_result = Variance::new(0, 0.0) + .call_flat(std::slice::from_ref(&x)) + .unwrap(); + let std_result = Std::new(0, 0.0).call_flat(&[x]).unwrap(); + + let Tensor::F64(var_arr) = &var_result[0] else { + panic!() + }; + let Tensor::F64(std_arr) = &std_result[0] else { + panic!() + }; + + let var_val = var_arr.as_slice().unwrap()[0]; + let std_val = std_arr.as_slice().unwrap()[0]; + assert!((std_val - var_val.sqrt()).abs() < 1e-10); + } + + #[test] + fn test_std_c128_returns_real() { + let data: Vec> = vec![Complex::new(1.0, 1.0), Complex::new(3.0, 3.0)]; + let x = Tensor::C128(ndarray::Array1::from(data).into_dyn().into_shared()); + let result = Std::new(0, 0.0).call_flat(&[x]).unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "C128 std should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + // std = sqrt(2.0) + approx_eq_slice(arr.as_slice().unwrap(), &[2.0_f64.sqrt()]); + } + + #[test] + fn test_std_c64_returns_real() { + let data: Vec> = vec![Complex::new(1.0, 1.0), Complex::new(3.0, 3.0)]; + let x = Tensor::C64(ndarray::Array1::from(data).into_dyn().into_shared()); + let result = Std::new(0, 0.0).call_flat(&[x]).unwrap(); + assert_eq!(result[0].dtype(), DType::F32, "C64 std should return F32"); + let Tensor::F32(arr) = &result[0] else { + panic!() + }; + // std = sqrt(2.0) + approx_eq_slice(arr.as_slice().unwrap(), &[2.0_f32.sqrt()]); + } + + #[test] + fn test_i8_cast_to_float() { + let data: Vec = vec![1, 3]; + let x = Tensor::I8(ndarray::Array1::from(data).into_dyn().into_shared()); + let result = Std::new(0, 0.0) + .call_flat(std::slice::from_ref(&x)) + .unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 std should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[1.0]); + let result = Mean::new(0).call_flat(std::slice::from_ref(&x)).unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 mean should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[2.]); + let result = Variance::new(0, 0.).call_flat(&[x]).unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 mean should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[1.]) + } + + #[test] + fn test_i16_cast_to_float() { + let data: Vec = vec![1, 3]; + let x = Tensor::I16(ndarray::Array1::from(data).into_dyn().into_shared()); + let result = Std::new(0, 0.0) + .call_flat(std::slice::from_ref(&x)) + .unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 std should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[1.0]); + let result = Mean::new(0).call_flat(std::slice::from_ref(&x)).unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 mean should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[2.]); + let result = Variance::new(0, 0.).call_flat(&[x]).unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 mean should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[1.]) + } + #[test] + fn test_i32_cast_to_float() { + let data: Vec = vec![1, 3]; + let x = Tensor::I32(ndarray::Array1::from(data).into_dyn().into_shared()); + let result = Std::new(0, 0.0) + .call_flat(std::slice::from_ref(&x)) + .unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 std should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[1.0]); + let result = Mean::new(0).call_flat(std::slice::from_ref(&x)).unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 mean should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[2.]); + let result = Variance::new(0, 0.).call_flat(&[x]).unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 mean should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[1.]) + } + + #[test] + fn test_i64_cast_to_float() { + let data: Vec = vec![1, 3]; + let x = Tensor::I64(ndarray::Array1::from(data).into_dyn().into_shared()); + let result = Std::new(0, 0.0) + .call_flat(std::slice::from_ref(&x)) + .unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 std should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[1.0]); + let result = Mean::new(0).call_flat(std::slice::from_ref(&x)).unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 mean should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[2.]); + let result = Variance::new(0, 0.).call_flat(&[x]).unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 mean should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[1.]) + } + + #[test] + fn test_u8_cast_to_float() { + let data: Vec = vec![1, 3]; + let x = Tensor::U8(ndarray::Array1::from(data).into_dyn().into_shared()); + let result = Std::new(0, 0.0) + .call_flat(std::slice::from_ref(&x)) + .unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 std should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[1.0]); + let result = Mean::new(0).call_flat(std::slice::from_ref(&x)).unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 mean should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[2.]); + let result = Variance::new(0, 0.).call_flat(&[x]).unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 mean should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[1.]) + } + + #[test] + fn test_u16_cast_to_float() { + let data: Vec = vec![1, 3]; + let x = Tensor::U16(ndarray::Array1::from(data).into_dyn().into_shared()); + let result = Std::new(0, 0.0) + .call_flat(std::slice::from_ref(&x)) + .unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 std should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[1.0]); + let result = Mean::new(0).call_flat(std::slice::from_ref(&x)).unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 mean should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[2.]); + let result = Variance::new(0, 0.).call_flat(&[x]).unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 mean should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[1.]) + } + #[test] + fn test_u32_cast_to_float() { + let data: Vec = vec![1, 3]; + let x = Tensor::U32(ndarray::Array1::from(data).into_dyn().into_shared()); + let result = Std::new(0, 0.0) + .call_flat(std::slice::from_ref(&x)) + .unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 std should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[1.0]); + let result = Mean::new(0).call_flat(std::slice::from_ref(&x)).unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 mean should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[2.]); + let result = Variance::new(0, 0.).call_flat(&[x]).unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 mean should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[1.]) + } + + #[test] + fn test_u64_cast_to_float() { + let data: Vec = vec![1, 3]; + let x = Tensor::U64(ndarray::Array1::from(data).into_dyn().into_shared()); + let result = Std::new(0, 0.0) + .call_flat(std::slice::from_ref(&x)) + .unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 std should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[1.0]); + let result = Mean::new(0).call_flat(std::slice::from_ref(&x)).unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 mean should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[2.]); + let result = Variance::new(0, 0.).call_flat(&[x]).unwrap(); + assert_eq!(result[0].dtype(), DType::F64, "I64 mean should return F64"); + let Tensor::F64(arr) = &result[0] else { + panic!() + }; + approx_eq_slice(arr.as_slice().unwrap(), &[1.]) + } + + #[test] + fn test_call_branch_where_leaf_expected_errors() { + let mut tree = DataTree::new(); + tree.insert_leaf("x", Tensor::from([1.0_f64, 2.0])); + let err = Mean::new(0).call(&tree).unwrap_err(); + assert!(matches!( + err, + CallError::::Input(CallInputError::ExpectedLeaf { + ref key, + }) if key.is_empty() + )); + } + + #[test] + fn test_mean_wrong_arity_errors() { + let err = Mean::new(0) + .call_flat(&[Tensor::from([1.0_f64]), Tensor::from([2.0_f64])]) + .unwrap_err(); + assert_eq!( + err, + MathNodeError::Input(CallInputError::WrongArity { + expected: 1, + actual: 2, + }) + ); + } + + #[test] + fn test_mean_call_end_to_end() { + let tree = DataTree::new_leaf(Tensor::from([1.0_f64, 2.0, 3.0, 4.0])); + let result = Mean::new(0).call(&tree).unwrap(); + let Tensor::F64(arr) = result.unwrap_leaf() else { + panic!("expected F64 leaf"); + }; + approx_eq_slice(arr.as_slice().unwrap(), &[2.5]); + } + + // --- Axis validation --- + + #[test] + fn test_mean_axis_out_of_bounds_errors() { + let x = Tensor::from([1.0_f64, 2.0, 3.0]); + let err = Mean::new(1).call_flat(&[x]).unwrap_err(); + assert_eq!(err, MathNodeError::InvalidAxis { axis: 1, ndim: 1 }); + } + + #[test] + fn test_variance_axis_out_of_bounds_errors() { + let x = Tensor::from([1.0_f64, 2.0, 3.0]); + let err = Variance::new(1, 0.0).call_flat(&[x]).unwrap_err(); + assert_eq!(err, MathNodeError::InvalidAxis { axis: 1, ndim: 1 }); + } + + #[test] + fn test_std_axis_out_of_bounds_errors() { + let x = Tensor::from([1.0_f64, 2.0, 3.0]); + let err = Std::new(1, 0.0).call_flat(&[x]).unwrap_err(); + assert_eq!(err, MathNodeError::InvalidAxis { axis: 1, ndim: 1 }); + } +} diff --git a/crates/providers/src/program_node.rs b/crates/providers/src/program_node.rs index b7a80b39beb6..f58701777c55 100644 --- a/crates/providers/src/program_node.rs +++ b/crates/providers/src/program_node.rs @@ -14,6 +14,28 @@ use crate::data_tree::{ArityMismatch, DataTree, TreeMatchError}; use crate::tensor::{DType, Tensor, TensorType}; use thiserror::Error; +/// Destructure `$args: &[Tensor]` into the named bindings, returning +/// [`CallInputError::WrongArity`] if the slice length does not match the pattern. +/// +/// ```ignore +/// crate::unpack_tensor_args!(args, [x, y]); // expects exactly 2 +/// crate::unpack_tensor_args!(args, [x]); // expects exactly 1 +/// ``` +#[macro_export] +macro_rules! unpack_tensor_args { + ($args:ident, [$($x:ident),+]) => { + let [$($x),+] = $args else { + return Err($crate::program_node::CallInputError::WrongArity { + expected: $crate::unpack_tensor_args!(@count $($x),+), + actual: $args.len(), + } + .into()); + }; + }; + (@count $x:ident) => { 1usize }; + (@count $x:ident, $($rest:ident),+) => { 1usize + $crate::unpack_tensor_args!(@count $($rest),+) }; +} + /// Errors returned when a tree-shaped argument does not match [`ProgramNode::input_types`]. #[derive(Debug, Clone, PartialEq, Eq, Error)] pub enum CallInputError { @@ -29,6 +51,9 @@ pub enum CallInputError { expected: String, actual: DType, }, + + #[error("expected {expected} total inputs, got {actual}")] + WrongArity { expected: usize, actual: usize }, } impl From for CallInputError { diff --git a/crates/providers/src/tensor.rs b/crates/providers/src/tensor.rs index 21bb65e65cf4..482adfc3b353 100644 --- a/crates/providers/src/tensor.rs +++ b/crates/providers/src/tensor.rs @@ -12,6 +12,7 @@ use ndarray::{ArcArrayD, ArrayD, IxDyn, Zip}; use num_complex::{Complex32, Complex64}; + use std::fmt; use thiserror::Error; @@ -303,7 +304,7 @@ macro_rules! cast_complex { /// Compute the NumPy-style broadcast shape for two operand shapes, or /// return [`TensorError::ShapeMismatch`] if they are not broadcast-compatible. -fn broadcast_shape(a: &[usize], b: &[usize]) -> Result, TensorError> { +pub fn broadcast_shape(a: &[usize], b: &[usize]) -> Result, TensorError> { let ndim = a.len().max(b.len()); (0..ndim) .map(|i| {