From 867642d9a3fc32358e45b83abbf61a949a0c69ef Mon Sep 17 00:00:00 2001
From: Marcel Luethi <marcel.luethi@inf.ethz.ch>
Date: Tue, 21 Jul 2026 16:34:08 +0200
Subject: [PATCH] use tensordot to have proper generalized outer product

---
 .../tensor/tensorops/ContractionOps.scala     |  6 +----
 .../tensor/TensorOpsContractionSuite.scala    | 22 +++++++++++++++++++
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/core/src/main/scala/dimwit/tensor/tensorops/ContractionOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/ContractionOps.scala
index e38429f..f29ff4d 100644
--- a/core/src/main/scala/dimwit/tensor/tensorops/ContractionOps.scala
+++ b/core/src/main/scala/dimwit/tensor/tensorops/ContractionOps.scala
@@ -26,11 +26,7 @@ object ContractionOps:
         primeConcat: PrimeConcat[T, OtherShape],
         labels: Labels[primeConcat.Out]
     ): Tensor[primeConcat.Out, V] = Tensor(
-      // Jax outer product flattens, reshape required
-      Jax.jnp.reshape(
-        Jax.jnp.outer(tensor.jaxValue, other.jaxValue),
-        (tensor.shape.dimensions ++ other.shape.dimensions).toPythonProxy
-      )
+      Jax.jnp.tensordot(tensor.jaxValue, other.jaxValue, axes = 0) // generalized outer product
     )
 
     /** Computes the dot product of this tensor with another tensor along the specified axis.
diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsContractionSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsContractionSuite.scala
index 262337f..87e0332 100644
--- a/core/src/test/scala/dimwit/tensor/TensorOpsContractionSuite.scala
+++ b/core/src/test/scala/dimwit/tensor/TensorOpsContractionSuite.scala
@@ -73,3 +73,25 @@ class TensorOpsContractionSuite extends DimwitTest:
           Array(10.0f, 20.0f, 20.0f, 40.0f)
         )
       )
+
+    it("Tensor2[A, B] and Tensor2[C, D] to Tensor4[A, B, C, D]"):
+      val mAB = m1
+      val mCD = m2.relabelAll((Axis[C], Axis[D]))
+
+      val res = mAB.outerProduct(mCD)
+
+      res.shape.labels shouldBe List("A", "B", "C", "D")
+      res should approxEqual(
+        Tensor.like(res).fromArray(
+          Array(
+            10.0f, 20.0f,
+            30.0f, 40.0f,
+            20.0f, 40.0f,
+            60.0f, 80.0f,
+            30.0f, 60.0f,
+            90.0f, 120.0f,
+            40.0f, 80.0f,
+            120.0f, 160.0f
+          )
+        )
+      )