Add impls of ProgramNode for various math operations

Author: ihincks

crates/providers/src/math_nodes/binary.rs

@@ -0,0 +1,243 @@
+// 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::{MissingCallError, ProgramNode};
+use crate::tensor::{DTypeLike, Tensor, TensorType, promotion};
+use std::borrow::Cow;
+use std::sync::OnceLock;
+
+/// Shared input type spec for all elementwise binary nodes: two broadcastable tensors `x` and `y`.
+fn elementwise_binary_input_types() -> &'static DataTree<TensorType> {
+    static LOCK: OnceLock<DataTree<TensorType>> = OnceLock::new();
+    LOCK.get_or_init(|| {
+        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.
+fn elementwise_binary_output_types() -> &'static DataTree<TensorType> {
+    static LOCK: OnceLock<DataTree<TensorType>> = OnceLock::new();
+    LOCK.get_or_init(|| {
+        DataTree::new_leaf(TensorType {
+            dtype: DTypeLike::Promotion(
+                vec![DTypeLike::Var("x".into()), DTypeLike::Var("y".into())].into(),
+            ),
+            shape: vec![],
+            broadcastable: true,
+        })
+    })
+}
+
+/// Extract `x` and `y` from `args`, promote dtypes, and apply `op` element-wise.
+fn binary_elementwise_call(
+    args: &DataTree<Tensor>,
+    op: impl Fn(&Tensor, &Tensor) -> Tensor,
+) -> Result<DataTree<Tensor>, MissingCallError> {
+    let DataTree::Leaf(x) = args.get_by_str_key("x").expect("missing input x") else {
+        panic!("expected leaf at x");
+    };
+    let DataTree::Leaf(y) = args.get_by_str_key("y").expect("missing input y") else {
+        panic!("expected leaf at y");
+    };
+    let out_dtype = promotion(x.dtype(), y.dtype());
+
+    // Use copy-on-write smart pointer to avoid cloning when promotion is unnecessary
+    let x = if x.dtype() == out_dtype {
+        Cow::Borrowed(x)
+    } else {
+        Cow::Owned(x.clone().cast(out_dtype))
+    };
+    let y = if y.dtype() == out_dtype {
+        Cow::Borrowed(y)
+    } else {
+        Cow::Owned(y.clone().cast(out_dtype))
+    };
+    Ok(DataTree::new_leaf(op(x.as_ref(), y.as_ref())))
+}
+
+/// 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 {
+            fn name(&self) -> &'static str {
+                $node_name
+            }
+            fn namespace(&self) -> &'static str {
+                "math"
+            }
+            fn input_types(&self) -> &DataTree<TensorType> {
+                elementwise_binary_input_types()
+            }
+            fn output_types(&self) -> &DataTree<TensorType> {
+                elementwise_binary_output_types()
+            }
+            fn implements_call(&self) -> bool {
+                true
+            }
+            fn call(&self, args: &DataTree<Tensor>) -> Result<DataTree<Tensor>, MissingCallError> {
+                binary_elementwise_call(args, $call_fn)
+            }
+        }
+    };
+}
+
+elementwise_binary_node!(Add, "add", |x, y| x + y);
+elementwise_binary_node!(Subtract, "subtract", |x, y| x - y);
+elementwise_binary_node!(Multiply, "multiply", |x, y| x * y);
+elementwise_binary_node!(Divide, "divide", |x, y| x / y);
+elementwise_binary_node!(Remainder, "remainder", |x, y| x % y);
+elementwise_binary_node!(Power, "power", |x, y| x.pow(y));
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::tensor::{DType, Tensor};
+
+    fn args(x: Tensor, y: Tensor) -> DataTree<Tensor> {
+        let mut tree = DataTree::new();
+        tree.insert_leaf("x", x);
+        tree.insert_leaf("y", y);
+        tree
+    }
+
+    #[test]
+    fn test_add_same_dtype() {
+        let result = Add.call(&args(Tensor::from([1.0_f64, 2.0, 3.0]), Tensor::from([4.0_f64, 5.0, 6.0]))).unwrap();
+        let DataTree::Leaf(Tensor::F64(arr)) = result else {
+            panic!("expected f64 leaf")
+        };
+        assert_eq!(arr.as_slice().unwrap(), &[5.0, 7.0, 9.0]);
+    }
+
+    #[test]
+    fn test_add_promotes_dtype() {
+        let result = Add.call(&args(Tensor::from([1.0_f32, 2.0]), Tensor::from([3.0_f64, 4.0]))).unwrap();
+        let DataTree::Leaf(tensor) = result else {
+            panic!("expected leaf")
+        };
+        assert_eq!(tensor.dtype(), DType::F64);
+        let Tensor::F64(arr) = tensor else {
+            panic!("expected f64")
+        };
+        assert_eq!(arr.as_slice().unwrap(), &[4.0, 6.0]);
+    }
+
+    #[test]
+    fn test_add_broadcasts_1d_scalar() {
+        // shape [3] + shape [1] -> shape [3]
+        let result = Add.call(&args(Tensor::from([1.0_f64, 2.0, 3.0]), Tensor::from([10.0_f64]))).unwrap();
+        let DataTree::Leaf(Tensor::F64(arr)) = result else {
+            panic!("expected f64 leaf")
+        };
+        assert_eq!(arr.as_slice().unwrap(), &[11.0, 12.0, 13.0]);
+    }
+
+    #[test]
+    fn test_add_broadcasts_2d_with_1d() {
+        // shape [2, 3] + shape [3] -> shape [2, 3]
+        use ndarray::arr2;
+        let x = Tensor::F64(arr2(&[[1.0_f64, 2.0, 3.0], [4.0, 5.0, 6.0]]).into_dyn());
+        let y = Tensor::from([10.0_f64, 20.0, 30.0]);
+        let result = Add.call(&args(x, y)).unwrap();
+        let DataTree::Leaf(Tensor::F64(arr)) = result else {
+            panic!("expected f64 leaf")
+        };
+        let expected = arr2(&[[11.0_f64, 22.0, 33.0], [14.0, 25.0, 36.0]]).into_dyn();
+        assert_eq!(arr, expected);
+    }
+
+    #[test]
+    fn test_subtract() {
+        let result = Subtract
+            .call(&args(Tensor::from([5.0_f64, 6.0, 7.0]), Tensor::from([1.0_f64, 2.0, 3.0])))
+            .unwrap();
+        let DataTree::Leaf(Tensor::F64(arr)) = result else {
+            panic!()
+        };
+        assert_eq!(arr.as_slice().unwrap(), &[4.0, 4.0, 4.0]);
+    }
+
+    #[test]
+    fn test_multiply() {
+        let result = Multiply
+            .call(&args(Tensor::from([2.0_f64, 3.0, 4.0]), Tensor::from([10.0_f64, 10.0, 10.0])))
+            .unwrap();
+        let DataTree::Leaf(Tensor::F64(arr)) = result else {
+            panic!()
+        };
+        assert_eq!(arr.as_slice().unwrap(), &[20.0, 30.0, 40.0]);
+    }
+
+    #[test]
+    fn test_divide() {
+        let result = Divide
+            .call(&args(Tensor::from([10.0_f64, 9.0, 8.0]), Tensor::from([2.0_f64, 3.0, 4.0])))
+            .unwrap();
+        let DataTree::Leaf(Tensor::F64(arr)) = result else {
+            panic!()
+        };
+        assert_eq!(arr.as_slice().unwrap(), &[5.0, 3.0, 2.0]);
+    }
+
+    #[test]
+    fn test_remainder() {
+        let result = Remainder
+            .call(&args(Tensor::from([7.0_f64, 8.0, 9.0]), Tensor::from([3.0_f64, 3.0, 3.0])))
+            .unwrap();
+        let DataTree::Leaf(Tensor::F64(arr)) = result else {
+            panic!()
+        };
+        assert_eq!(arr.as_slice().unwrap(), &[1.0, 2.0, 0.0]);
+    }
+
+    #[test]
+    fn test_power() {
+        let result = Power
+            .call(&args(Tensor::from([2.0_f64, 3.0, 4.0]), Tensor::from([3.0_f64, 2.0, 1.0])))
+            .unwrap();
+        let DataTree::Leaf(Tensor::F64(arr)) = result else {
+            panic!()
+        };
+        assert_eq!(arr.as_slice().unwrap(), &[8.0, 9.0, 4.0]);
+    }
+
+    #[test]
+    fn test_power_broadcasts() {
+        // shape [3] ** shape [1] -> shape [3]
+        let result = Power.call(&args(Tensor::from([2.0_f64, 3.0, 4.0]), Tensor::from([2.0_f64]))).unwrap();
+        let DataTree::Leaf(Tensor::F64(arr)) = result else {
+            panic!()
+        };
+        assert_eq!(arr.as_slice().unwrap(), &[4.0, 9.0, 16.0]);
+    }
+}