feat(deep_causality_tensor): Updated EinSumOp for better ergonomics.

Signed-off-by: Marvin Hansen <[email protected]>

Author: marvin-hansen

deep_causality_tensor/README.md

@@ -122,11 +122,7 @@ fn main() {
     let vec2_tensor = CausalTensor::new(vec2_data, vec2_shape).unwrap();
     
     // Execute the Einstein summation for dot product 
-    let result_dot_prod = CausalTensor::ein_sum(&EinSumOp::dot_prod(
-        vec1_tensor.clone(),
-        vec2_tensor.clone(),
-    ))
-    .unwrap();
+    let result_dot_prod = CausalTensor::ein_sum(&EinSumOp::dot_prod(vec1_tensor, vec2_tensor)).unwrap();
     println!("Result of Dot Product:\n{:?}", result_dot_prod);
 }
 ```

deep_causality_tensor/examples/ein_sum_causal_tensor.rs

@@ -18,8 +18,7 @@ fn main() {
     println!("LHS Tensor:\n{:?}", lhs_tensor);
     println!("RHS Tensor:\n{:?}", rhs_tensor);
 
-    let result_mat_mul =
-        CausalTensor::ein_sum(&EinSumOp::mat_mul(lhs_tensor.clone(), rhs_tensor.clone())).unwrap();
+    let result_mat_mul = CausalTensor::ein_sum(&EinSumOp::mat_mul(lhs_tensor, rhs_tensor)).unwrap();
     println!("Result of Matrix Multiplication:\n{:?}", result_mat_mul);
     let expected_mat_mul = CausalTensor::new(vec![19.0, 22.0, 43.0, 50.0], vec![2, 2]).unwrap();
     assert_eq!(result_mat_mul, expected_mat_mul);
@@ -37,11 +36,8 @@ fn main() {
     println!("Vector 1:\n{:?}", vec1_tensor);
     println!("Vector 2:\n{:?}", vec2_tensor);
 
-    let result_dot_prod = CausalTensor::ein_sum(&EinSumOp::dot_prod(
-        vec1_tensor.clone(),
-        vec2_tensor.clone(),
-    ))
-    .unwrap();
+    let result_dot_prod =
+        CausalTensor::ein_sum(&EinSumOp::dot_prod(vec1_tensor, vec2_tensor)).unwrap();
     println!("Result of Dot Product:\n{:?}", result_dot_prod);
     let expected_dot_prod = CausalTensor::new(vec![32.0], vec![]).unwrap();
     assert_eq!(result_dot_prod, expected_dot_prod);
@@ -53,7 +49,7 @@ fn main() {
     let trace_tensor = CausalTensor::new(trace_data, trace_shape).unwrap();
 
     println!("Tensor for Trace:\n{:?}", trace_tensor);
-    let result_trace = CausalTensor::ein_sum(&EinSumOp::trace(trace_tensor.clone(), 0, 1)).unwrap();
+    let result_trace = CausalTensor::ein_sum(&EinSumOp::trace(trace_tensor, 0, 1)).unwrap();
     println!("Result of Trace (axes 0, 1):\n{:?}", result_trace);
     let expected_trace = CausalTensor::new(vec![5.0], vec![]).unwrap();
     assert_eq!(result_trace, expected_trace);
@@ -72,8 +68,8 @@ fn main() {
     println!("RHS Tensor for Element-wise Product:\n{:?}", ew_rhs_tensor);
 
     let result_ew_prod = CausalTensor::ein_sum(&EinSumOp::element_wise_product(
-        ew_lhs_tensor.clone(),
-        ew_rhs_tensor.clone(),
+        ew_lhs_tensor,
+        ew_rhs_tensor,
     ))
     .unwrap();
     println!("Result of Element-wise Product:\n{:?}", result_ew_prod);
@@ -100,11 +96,9 @@ fn main() {
     println!("LHS Tensor for Batch MatMul:\n{:?}", bmm_lhs_tensor);
     println!("RHS Tensor for Batch MatMul:\n{:?}", bmm_rhs_tensor);
 
-    let result_bmm = CausalTensor::ein_sum(&EinSumOp::batch_mat_mul(
-        bmm_lhs_tensor.clone(),
-        bmm_rhs_tensor.clone(),
-    ))
-    .unwrap();
+    let result_bmm =
+        CausalTensor::ein_sum(&EinSumOp::batch_mat_mul(bmm_lhs_tensor, bmm_rhs_tensor)).unwrap();
+
     println!("Result of Batch Matrix Multiplication:\n{:?}", result_bmm);
     let expected_bmm = CausalTensor::new(
         vec![

deep_causality_tensor/src/types/causal_tensor/from/mod.rs

@@ -19,3 +19,10 @@ impl<'a, T: Clone> From<&'a T> for CausalTensor<T> {
             .expect("Failed to create scalar CausalTensor from &T")
     }
 }
+
+impl<'a, T: Clone> From<&'a CausalTensor<T>> for CausalTensor<T> {
+    /// Creates a new `CausalTensor` by cloning an existing `CausalTensor` reference.
+    fn from(item: &'a CausalTensor<T>) -> Self {
+        item.clone()
+    }
+}

deep_causality_tensor/src/types/causal_tensor/op_tensor_ein_sum/ein_sum_op.rs

@@ -79,14 +79,14 @@ impl<T> EinSumOp<T> {
     /// # Returns
     ///
     /// An `EinSumAST<T>` node representing the contraction operation.
-    pub fn contraction(
-        lhs: CausalTensor<T>,
-        rhs: CausalTensor<T>,
+    pub fn contraction<L: Into<CausalTensor<T>>, R: Into<CausalTensor<T>>>(
+        lhs: L,
+        rhs: R,
         lhs_axes: Vec<usize>,
         rhs_axes: Vec<usize>,
     ) -> EinSumAST<T> {
-        let lhs_leaf = EinSumOp::tensor_source(lhs);
-        let rhs_leaf = EinSumOp::tensor_source(rhs);
+        let lhs_leaf = EinSumOp::tensor_source(lhs.into());
+        let rhs_leaf = EinSumOp::tensor_source(rhs.into());
         EinSumAST::with_children(
             EinSumOp::Contraction { lhs_axes, rhs_axes },
             vec![lhs_leaf, rhs_leaf],
@@ -105,8 +105,8 @@ impl<T> EinSumOp<T> {
     /// # Returns
     ///
     /// An `EinSumAST<T>` node representing the reduction operation.
-    pub fn reduction(operand: CausalTensor<T>, axes: Vec<usize>) -> EinSumAST<T> {
-        let operand_leaf = EinSumOp::tensor_source(operand);
+    pub fn reduction<O: Into<CausalTensor<T>>>(operand: O, axes: Vec<usize>) -> EinSumAST<T> {
+        let operand_leaf = EinSumOp::tensor_source(operand.into());
         EinSumAST::with_children(EinSumOp::Reduction { axes }, vec![operand_leaf])
     }
 
@@ -122,9 +122,12 @@ impl<T> EinSumOp<T> {
     /// # Returns
     ///
     /// An `EinSumAST<T>` node representing the matrix multiplication operation.
-    pub fn mat_mul(lhs: CausalTensor<T>, rhs: CausalTensor<T>) -> EinSumAST<T> {
-        let lhs_leaf = EinSumOp::tensor_source(lhs);
-        let rhs_leaf = EinSumOp::tensor_source(rhs);
+    pub fn mat_mul<L: Into<CausalTensor<T>>, R: Into<CausalTensor<T>>>(
+        lhs: L,
+        rhs: R,
+    ) -> EinSumAST<T> {
+        let lhs_leaf = EinSumOp::tensor_source(lhs.into());
+        let rhs_leaf = EinSumOp::tensor_source(rhs.into());
         EinSumAST::with_children(EinSumOp::MatMul, vec![lhs_leaf, rhs_leaf])
     }
 
@@ -140,9 +143,12 @@ impl<T> EinSumOp<T> {
     /// # Returns
     ///
     /// An `EinSumAST<T>` node representing the dot product operation.
-    pub fn dot_prod(lhs: CausalTensor<T>, rhs: CausalTensor<T>) -> EinSumAST<T> {
-        let lhs_leaf = EinSumOp::tensor_source(lhs);
-        let rhs_leaf = EinSumOp::tensor_source(rhs);
+    pub fn dot_prod<L: Into<CausalTensor<T>>, R: Into<CausalTensor<T>>>(
+        lhs: L,
+        rhs: R,
+    ) -> EinSumAST<T> {
+        let lhs_leaf = EinSumOp::tensor_source(lhs.into());
+        let rhs_leaf = EinSumOp::tensor_source(rhs.into());
         EinSumAST::with_children(EinSumOp::DotProd, vec![lhs_leaf, rhs_leaf])
     }
 
@@ -159,8 +165,8 @@ impl<T> EinSumOp<T> {
     /// # Returns
     ///
     /// An `EinSumAST<T>` node representing the trace operation.
-    pub fn trace(operand: CausalTensor<T>, axes1: usize, axes2: usize) -> EinSumAST<T> {
-        let operand_leaf = EinSumOp::tensor_source(operand);
+    pub fn trace<O: Into<CausalTensor<T>>>(operand: O, axes1: usize, axes2: usize) -> EinSumAST<T> {
+        let operand_leaf = EinSumOp::tensor_source(operand.into());
         EinSumAST::with_children(EinSumOp::Trace { axes1, axes2 }, vec![operand_leaf])
     }
 
@@ -176,9 +182,12 @@ impl<T> EinSumOp<T> {
     /// # Returns
     ///
     /// An `EinSumAST<T>` node representing the tensor product operation.
-    pub fn tensor_product(lhs: CausalTensor<T>, rhs: CausalTensor<T>) -> EinSumAST<T> {
-        let lhs_leaf = EinSumOp::tensor_source(lhs);
-        let rhs_leaf = EinSumOp::tensor_source(rhs);
+    pub fn tensor_product<L: Into<CausalTensor<T>>, R: Into<CausalTensor<T>>>(
+        lhs: L,
+        rhs: R,
+    ) -> EinSumAST<T> {
+        let lhs_leaf = EinSumOp::tensor_source(lhs.into());
+        let rhs_leaf = EinSumOp::tensor_source(rhs.into());
         EinSumAST::with_children(EinSumOp::TensorProduct, vec![lhs_leaf, rhs_leaf])
     }
 
@@ -194,9 +203,12 @@ impl<T> EinSumOp<T> {
     /// # Returns
     ///
     /// An `EinSumAST<T>` node representing the element-wise product operation.
-    pub fn element_wise_product(lhs: CausalTensor<T>, rhs: CausalTensor<T>) -> EinSumAST<T> {
-        let lhs_leaf = EinSumOp::tensor_source(lhs);
-        let rhs_leaf = EinSumOp::tensor_source(rhs);
+    pub fn element_wise_product<L: Into<CausalTensor<T>>, R: Into<CausalTensor<T>>>(
+        lhs: L,
+        rhs: R,
+    ) -> EinSumAST<T> {
+        let lhs_leaf = EinSumOp::tensor_source(lhs.into());
+        let rhs_leaf = EinSumOp::tensor_source(rhs.into());
         EinSumAST::with_children(EinSumOp::ElementWiseProduct, vec![lhs_leaf, rhs_leaf])
     }
 
@@ -212,8 +224,8 @@ impl<T> EinSumOp<T> {
     /// # Returns
     ///
     /// An `EinSumAST<T>` node representing the transpose operation.
-    pub fn transpose(operand: CausalTensor<T>, new_order: Vec<usize>) -> EinSumAST<T> {
-        let operand_leaf = EinSumOp::tensor_source(operand);
+    pub fn transpose<O: Into<CausalTensor<T>>>(operand: O, new_order: Vec<usize>) -> EinSumAST<T> {
+        let operand_leaf = EinSumOp::tensor_source(operand.into());
         EinSumAST::with_children(EinSumOp::Transpose { new_order }, vec![operand_leaf])
     }
 
@@ -230,12 +242,12 @@ impl<T> EinSumOp<T> {
     /// # Returns
     ///
     /// An `EinSumAST<T>` node representing the diagonal extraction operation.
-    pub fn diagonal_extraction(
-        operand: CausalTensor<T>,
+    pub fn diagonal_extraction<O: Into<CausalTensor<T>>>(
+        operand: O,
         axes1: usize,
         axes2: usize,
     ) -> EinSumAST<T> {
-        let operand_leaf = EinSumOp::tensor_source(operand);
+        let operand_leaf = EinSumOp::tensor_source(operand.into());
         EinSumAST::with_children(
             EinSumOp::DiagonalExtraction { axes1, axes2 },
             vec![operand_leaf],
@@ -254,9 +266,12 @@ impl<T> EinSumOp<T> {
     /// # Returns
     ///
     /// An `EinSumAST<T>` node representing the batch matrix multiplication operation.
-    pub fn batch_mat_mul(lhs: CausalTensor<T>, rhs: CausalTensor<T>) -> EinSumAST<T> {
-        let lhs_leaf = EinSumOp::tensor_source(lhs);
-        let rhs_leaf = EinSumOp::tensor_source(rhs);
+    pub fn batch_mat_mul<L: Into<CausalTensor<T>>, R: Into<CausalTensor<T>>>(
+        lhs: L,
+        rhs: R,
+    ) -> EinSumAST<T> {
+        let lhs_leaf = EinSumOp::tensor_source(lhs.into());
+        let rhs_leaf = EinSumOp::tensor_source(rhs.into());
         EinSumAST::with_children(EinSumOp::BatchMatMul, vec![lhs_leaf, rhs_leaf])
     }
 }

deep_causality_tensor/tests/types/causal_tensor/op_tensor_ein_sum_tests.rs

@@ -21,7 +21,19 @@ fn test_ein_sum_mat_mul() {
     let rhs = utils_tests::matrix_tensor(vec![5.0, 6.0, 7.0, 8.0], 2, 2);
     let expected = utils_tests::matrix_tensor(vec![19.0, 22.0, 43.0, 50.0], 2, 2);
 
-    let ast = EinSumOp::mat_mul(lhs, rhs);
+    let ast = EinSumOp::<f64>::mat_mul(lhs, rhs);
+    let result = CausalTensor::ein_sum(&ast).unwrap();
+    assert_eq!(result, expected);
+}
+
+#[test]
+fn test_ein_sum_mat_mul_with_references() {
+    let lhs_owned = utils_tests::matrix_tensor(vec![1.0, 2.0, 3.0, 4.0], 2, 2);
+    let rhs_owned = utils_tests::matrix_tensor(vec![5.0, 6.0, 7.0, 8.0], 2, 2);
+    let expected = utils_tests::matrix_tensor(vec![19.0, 22.0, 43.0, 50.0], 2, 2);
+
+    // Pass references to the EinSumOp::mat_mul method
+    let ast = EinSumOp::<f64>::mat_mul(&lhs_owned, &rhs_owned);
     let result = CausalTensor::ein_sum(&ast).unwrap();
     assert_eq!(result, expected);
 }
@@ -59,7 +71,7 @@ fn test_ein_sum_trace() {
     let operand = utils_tests::matrix_tensor(vec![1.0, 2.0, 3.0, 4.0], 2, 2);
     let expected = utils_tests::scalar_tensor(5.0);
 
-    let ast = EinSumOp::trace(operand, 0, 1);
+    let ast = EinSumOp::<f64>::trace(operand, 0, 1);
     let result = CausalTensor::ein_sum(&ast).unwrap();
     assert_eq!(result, expected);
 }
@@ -142,7 +154,7 @@ fn test_ein_sum_error_propagation() {
     let lhs = utils_tests::vector_tensor(vec![1.0, 2.0]);
     let rhs = utils_tests::matrix_tensor(vec![5.0, 6.0, 7.0, 8.0], 2, 2);
 
-    let ast = EinSumOp::mat_mul(lhs, rhs);
+    let ast = EinSumOp::<f64>::mat_mul(lhs, rhs);
     let err = CausalTensor::ein_sum(&ast).unwrap_err();
     assert!(matches!(
         err,
@@ -154,7 +166,7 @@ fn test_ein_sum_error_propagation() {
 
     // Test Trace with invalid axes, expecting error from trace
     let operand = utils_tests::matrix_tensor(vec![1.0; 4], 2, 2);
-    let ast = EinSumOp::trace(operand, 0, 0);
+    let ast = EinSumOp::<f64>::trace(operand, 0, 0);
     let err = CausalTensor::ein_sum(&ast).unwrap_err();
     assert!(matches!(
         err,

deep_causality_tensor/tests/types/causal_tensor/op_tensor_reduction_tests.rs

@@ -26,7 +26,7 @@ fn test_sum_axes_2d() {
 fn test_sum_axes_full_reduction() {
     let tensor = CausalTensor::new(vec![1, 2, 3, 4, 5, 6], vec![2, 3]).unwrap();
     let sum_all = tensor.sum_axes(&[]).unwrap();
-    assert_eq!(sum_all.shape(), &[]);
+    assert_eq!(sum_all.shape(), &[] as &[usize]);
     assert_eq!(sum_all.as_slice(), &[21]);
 }
 
@@ -60,7 +60,7 @@ fn test_sum_axes_empty_tensor() {
 
     // Full reduction of an empty tensor should result in a scalar 0
     let sum_all = tensor.sum_axes(&[]).unwrap();
-    assert_eq!(sum_all.shape(), &[]);
+    assert_eq!(sum_all.shape(), &[] as &[usize]);
     assert_eq!(sum_all.as_slice(), &[0]);
 }
 
@@ -92,7 +92,7 @@ fn test_mean_axes_2d() {
 fn test_mean_axes_full_reduction() {
     let tensor = CausalTensor::new(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
     let mean_all = tensor.mean_axes(&[]).unwrap();
-    assert_eq!(mean_all.shape(), &[]);
+    assert_eq!(mean_all.shape(), &[] as &[usize]);
     assert_eq!(mean_all.as_slice(), &[3.5]);
 }
 

deep_causality_tensor/tests/types/causal_tensor/op_tensor_shape_tests.rs

@@ -165,6 +165,6 @@ fn test_permute_axes_scalar() {
     let tensor = CausalTensor::new(vec![42], vec![]).unwrap();
     // Permuting a 0-dim tensor
     let permuted = tensor.permute_axes(&[]).unwrap();
-    assert_eq!(permuted.shape(), &[]);
+    assert_eq!(permuted.shape(), &[] as &[usize]);
     assert_eq!(permuted, tensor);
 }

deep_causality_tensor/tests/types/causal_tensor/op_tensor_tensor_tests.rs

@@ -183,7 +183,7 @@ fn test_binary_op_scalar_tensors() {
     let t1 = CausalTensor::new(vec![10.0], vec![]).unwrap();
     let t2 = CausalTensor::new(vec![2.0], vec![]).unwrap();
     let expected_data = vec![20.0];
-    let expected_shape = vec![];
+    let expected_shape: Vec<usize> = vec![];
 
     let result = (&t1 * &t2).unwrap();
     assert_eq!(result.as_slice(), expected_data.as_slice());