VectorOps.dot and MatrixOps.matMul are slightly faster

tribbloid · Mar 27, 2022 · 00c194d · 00c194d
1 parent bc90255
commit 00c194d
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 8 deletions.
diff --git a/core/src/main/scala/shapesafe/core/shape/Names.scala b/core/src/main/scala/shapesafe/core/shape/Names.scala
@@ -74,7 +74,11 @@ object Names extends Tuples with ApplyLiterals.ToNames {
   object Syntax extends Syntax
 
   val i = Names("i")
+  val ii = Names("i", "i")
+
   val ij = Names("i", "j")
   val jk = Names("j", "k")
+  val ijjk = Names("i", "j", "j", "k")
+
   val ik = Names("i", "k")
 }
diff --git a/core/src/main/scala/shapesafe/core/shape/ops/MatrixOps.scala b/core/src/main/scala/shapesafe/core/shape/ops/MatrixOps.scala
@@ -1,23 +1,33 @@
 package shapesafe.core.shape.ops
 
 import shapeless.Nat
-import shapesafe.core.shape.{Index, Indices, Names, Shape}
+import shapesafe.core.shape.Index.LtoR
+import shapesafe.core.shape.{Indices, Names, Shape}
 
 trait MatrixOps extends HasShape {
 
   def matMul[THAT <: Shape](that: THAT) = {
-    val s1 = shape :<<= Names.ij
-    val s2 = that :<<= Names.jk
+//    val s1 = shape :<<= Names.ij
+//    val s2 = that :<<= Names.jk
+//
+//    s1.einSum(s2) --> Names.ik
 
-    s1.einSum(s2) --> Names.ik
+    // TODO: above is slight slower, need to optimise
+
+    val outer = (shape >< that) :<<= Names.ijjk
+    outer.einSum --> Names.ik
   }
 
   def mat_*[THAT <: Shape](that: THAT) = {
     matMul[that.type](that)
   }
 
   def transpose = {
-    shape.requireNumDim(Nat._2).rearrangeBy(Indices & Index.LtoR(1) & Index.LtoR(0))
+    shape
+      .requireNumDim(Nat._2)
+      .rearrangeBy(
+        Indices & LtoR(1) & LtoR(0)
+      )
   }
 
   def `^T` = {

diff --git a/core/src/main/scala/shapesafe/core/shape/ops/VectorOps.scala b/core/src/main/scala/shapesafe/core/shape/ops/VectorOps.scala
@@ -5,10 +5,15 @@ import shapesafe.core.shape.{Names, Shape}
 trait VectorOps extends HasShape {
 
   def dot[THAT <: Shape](that: THAT) = {
-    val s1 = shape :<<= Names.i
-    val s2 = that :<<= Names.i
+//    val s1 = shape :<<= Names.i
+//    val s2 = that :<<= Names.i
+//
+//    (s1.einSum(s2) --> Names.Eye) >< Shape(1)
 
-    (s1.einSum(s2) --> Names.Eye) >< Shape(1)
+    // TODO: above is slight slower, need to optimise
+
+    val outer = (shape >< that) :<<= Names.ii
+    (outer.einSum --> Names.Eye) >< Shape._1
   }
 
   def cross[THAT <: Shape](that: THAT) = {