Repository: coreylowman/dfdx (language: rust)

Issue:
Add matmul(matrix, vector)
Here's the code for it.

```rust
use dfdx::prelude::*;
use Tensor1D as V;
use Tensor2D as M;

fn matmul<const N0: usize, const N1: usize>(lns: M<N0, N1>, rhs: V<N1>) -> V<N0> {
    let rhs = rhs.reshape::<Rank2<N1, 1>>();
    let res: M<N0, 1> = lns.matmul(rhs);
    res.reshape()
}
```


Discussion / hints:
This should be fairly close in implementation to the vector * matrix implementation. Any interest in contributing this? I can help out with the cuda stuff
I looked at here: https://github.com/coreylowman/dfdx/tree/main/src/tensor_ops/matmul

The problem is that I don't understand the complex type class or how to use `sgemm`. Is it good enough to use `Tensor::reshape` in my impl?

Fix patch:
diff --git a/src/tensor_ops/matmul/cpu_kernel.rs b/src/tensor_ops/matmul/cpu_kernel.rs
index e8e235392..4a08face8 100644
--- a/src/tensor_ops/matmul/cpu_kernel.rs
+++ b/src/tensor_ops/matmul/cpu_kernel.rs
@@ -160,117 +160,6 @@ impl MatMulImpl<f64> for Cpu {
     }
 }
 
-impl<E: Dtype> super::VecVecKernel<E> for Cpu
-where
-    Self: MatMulImpl<E>,
-{
-    fn forward<M: Dim, N: Dim>(
-        &self,
-        lhs: &Tensor<(M,), E, Self>,
-        rhs: &Tensor<(N,), E, Self>,
-    ) -> Result<Tensor<(M, N), E, Self>, Self::Err> {
-        let m = lhs.shape.0;
-        let n = rhs.shape.0;
-        let mut out = self.try_zeros_like(&(m, n))?;
-        Self::matmul(
-            (m, Const::<1>, n),
-            lhs.data.as_ptr(),
-            [lhs.strides[0], 0],
-            rhs.data.as_ptr(),
-            [0, rhs.strides[0]],
-            Arc::get_mut(&mut out.data).unwrap().as_mut_ptr(),
-            out.strides,
-        );
-        Ok(out)
-    }
-
-    fn backward<M: Dim, N: Dim>(
-        &self,
-        lhs: &Tensor<(M,), E, Self>,
-        grad_lhs: &mut Self::Vec<E>,
-        rhs: &Tensor<(N,), E, Self>,
-        grad_rhs: &mut Self::Vec<E>,
-        grad_out: &Self::Vec<E>,
-    ) -> Result<(), Self::Err> {
-        let m = lhs.shape.0;
-        let k = Const::<1>;
-        let n = rhs.shape.0;
-        Self::matmul(
-            (m, n, k),
-            grad_out.as_ptr(),
-            [n.size(), 1],
-            rhs.data.as_ptr(),
-            [rhs.strides[0], 0],
-            grad_lhs.as_mut_ptr(),
-            [lhs.strides[0], 0],
-        );
-        Self::matmul(
-            (k, m, n),
-            lhs.data.as_ptr(),
-            [0, lhs.strides[0]],
-            grad_out.as_ptr(),
-            [n.size(), 1],
-            grad_rhs.as_mut_ptr(),
-            [0, rhs.strides[0]],
-        );
-        Ok(())
-    }
-}
-
-impl<E: Dtype> super::VecMatKernel<E> for Cpu
-where
-    Self: MatMulImpl<E>,
-{
-    fn forward<K: Dim, N: Dim>(
-        &self,
-        lhs: &Tensor<(K,), E, Self>,
-        rhs: &Tensor<(K, N), E, Self>,
-    ) -> Result<Tensor<(N,), E, Self>, Self::Err> {
-        let (k, n) = rhs.shape;
-        let mut out = self.try_zeros_like(&(n,))?;
-        Self::matmul(
-            (Const::<1>, k, n),
-            lhs.data.as_ptr(),
-            [0, lhs.strides[0]],
-            rhs.data.as_ptr(),
-            rhs.strides,
-            Arc::get_mut(&mut out.data).unwrap().as_mut_ptr(),
-            [0, 1],
-        );
-        Ok(out)
-    }
-    fn backward<K: Dim, N: Dim>(
-        &self,
-        lhs: &Tensor<(K,), E, Self>,
-        grad_lhs: &mut Self::Vec<E>,
-        rhs: &Tensor<(K, N), E, Self>,
-        grad_rhs: &mut Self::Vec<E>,
-        grad_out: &Self::Vec<E>,
-    ) -> Result<(), Self::Err> {
-        let m = Const::<1>;
-        let (k, n) = rhs.shape;
-        Self::matmul(
-            (m, n, k),
-            grad_out.as_ptr(),
-            [0, 1],
-            rhs.data.as_ptr(),
-            [rhs.strides[1], rhs.strides[0]],
-            grad_lhs.as_mut_ptr(),
-            [0, lhs.strides[0]],
-        );
-        Self::matmul(
-            (k, m, n),
-            lhs.data.as_ptr(),
-            [lhs.strides[0], 0],
-            grad_out.as_ptr(),
-            [0, 1],
-            grad_rhs.as_mut_ptr(),
-            rhs.strides,
-        );
-        Ok(())
-    }
-}
-
 impl<E: Dtype> super::MatMatKernel<E> for Cpu
 where
     Self: MatMulImpl<E>,
diff --git a/src/tensor_ops/matmul/cuda_kernel.rs b/src/tensor_ops/matmul/cuda_kernel.rs
index 5a4618c79..c6aa739f2 100644
--- a/src/tensor_ops/matmul/cuda_kernel.rs
+++ b/src/tensor_ops/matmul/cuda_kernel.rs
@@ -161,156 +161,6 @@ impl Cuda {
     }
 }
 
-impl<E: Dtype> super::VecVecKernel<E> for Cuda
-where
-    CudaBlas: Gemm<E>,
-{
-    fn forward<M: Dim, N: Dim>(
-        &self,
-        lhs: &Tensor<(M,), E, Self>,
-        rhs: &Tensor<(N,), E, Self>,
-    ) -> Result<Tensor<(M, N), E, Self>, Self::Err> {
-        let (m,) = lhs.shape;
-        let (n,) = rhs.shape;
-        let k = Const::<1>;
-        let shape = (m, n);
-        let strides = shape.strides();
-        let mut storage = unsafe { self.alloc_empty::<E>(shape.num_elements()) }?;
-
-        unsafe {
-            self.gemm(
-                (m, k, n),
-                lhs.data.as_ref(),
-                [lhs.strides[0], 0],
-                rhs.data.as_ref(),
-                [0, rhs.strides[0]],
-                Default::default(),
-                &mut storage,
-                strides,
-            )
-        }?;
-
-        Ok(self.build_tensor(shape, strides, storage))
-    }
-
-    fn backward<M: Dim, N: Dim>(
-        &self,
-        lhs: &Tensor<(M,), E, Self>,
-        grad_lhs: &mut Self::Vec<E>,
-        rhs: &Tensor<(N,), E, Self>,
-        grad_rhs: &mut Self::Vec<E>,
-        grad_out: &Self::Vec<E>,
-    ) -> Result<(), Self::Err> {
-        let m = lhs.shape.0;
-        let k = Const::<1>;
-        let n = rhs.shape.0;
-        let strides = (m, n).strides();
-        self.par_stream.wait_for_default()?;
-        unsafe {
-            // grad_lhs += grad_out * rhs^T
-            self.blas.set_stream(Some(self.par_stream.as_ref()))?;
-            self.gemm(
-                (m, n, k),
-                grad_out,
-                strides,
-                rhs.data.as_ref(),
-                [rhs.strides[0], 0],
-                E::ONE,
-                grad_lhs,
-                [lhs.strides[0], 0],
-            )?;
-            self.blas.set_stream(None)?;
-
-            // grad_rhs += lhs^T * grad_out
-            self.gemm(
-                (k, m, n),
-                lhs.data.as_ref(),
-                [0, lhs.strides[0]],
-                grad_out,
-                strides,
-                E::ONE,
-                grad_rhs,
-                [0, rhs.strides[0]],
-            )?;
-        }
-        self.dev.wait_for(self.par_stream.as_ref())?;
-        Ok(())
-    }
-}
-
-impl<E: Dtype> super::VecMatKernel<E> for Cuda
-where
-    CudaBlas: Gemm<E>,
-{
-    fn forward<K: Dim, N: Dim>(
-        &self,
-        lhs: &Tensor<(K,), E, Self>,
-        rhs: &Tensor<(K, N), E, Self>,
-    ) -> Result<Tensor<(N,), E, Self>, Self::Err> {
-        let m = Const::<1>;
-        let (k, n) = rhs.shape;
-        let shape = (n,);
-        let strides = shape.strides();
-        let mut storage = unsafe { self.alloc_empty::<E>(shape.num_elements()) }?;
-
-        unsafe {
-            self.gemm(
-                (m, k, n),
-                lhs.data.as_ref(),
-                [0, lhs.strides[0]],
-                rhs.data.as_ref(),
-                rhs.strides,
-                Default::default(),
-                &mut storage,
-                [0, strides[0]],
-            )?;
-        }
-
-        Ok(self.build_tensor(shape, strides, storage))
-    }
-    fn backward<K: Dim, N: Dim>(
-        &self,
-        lhs: &Tensor<(K,), E, Self>,
-        grad_lhs: &mut Self::Vec<E>,
-        rhs: &Tensor<(K, N), E, Self>,
-        grad_rhs: &mut Self::Vec<E>,
-        grad_out: &Self::Vec<E>,
-    ) -> Result<(), Self::Err> {
-        let m = Const::<1>;
-        let (k, n) = rhs.shape;
-        self.par_stream.wait_for_default()?;
-        unsafe {
-            // grad_lhs += grad_out * rhs^T
-            self.blas.set_stream(Some(self.par_stream.as_ref()))?;
-            self.gemm(
-                (m, n, k),
-                grad_out,
-                [0, 1],
-                rhs.data.as_ref(),
-                [rhs.strides[1], rhs.strides[0]],
-                E::ONE,
-                grad_lhs,
-                [0, lhs.strides[0]],
-            )?;
-            self.blas.set_stream(None)?;
-
-            // grad_rhs += lhs^T * grad_out
-            self.gemm(
-                (k, m, n),
-                lhs.data.as_ref(),
-                [lhs.strides[0], 0],
-                grad_out,
-                [0, 1],
-                E::ONE,
-                grad_rhs,
-                rhs.strides,
-            )?;
-        }
-        self.dev.wait_for(self.par_stream.as_ref())?;
-        Ok(())
-    }
-}
-
 impl<E: Dtype> super::MatMatKernel<E> for Cuda
 where
     CudaBlas: Gemm<E>,
diff --git a/src/tensor_ops/matmul/mod.rs b/src/tensor_ops/matmul/mod.rs
index 058b4a792..a1fb3c7ef 100644
--- a/src/tensor_ops/matmul/mod.rs
+++ b/src/tensor_ops/matmul/mod.rs
@@ -6,10 +6,12 @@ pub(super) mod cpu_kernel;
 pub(super) mod cuda_kernel;
 
 use crate::{
-    shapes::{Dim, Dtype, Shape},
+    shapes::{Const, Dim, Dtype, Shape},
     tensor::{DeviceStorage, HasErr, Merge, PutTape, SplitTape, Tape, Tensor},
 };
 
+use super::reshape_to::{ReshapeKernel, ReshapeTo};
+
 /// Matrix * Matrix, Vector * Matrix, Vector * Vector, and broadcasted/batched versions.
 ///
 /// # Examples
@@ -108,74 +110,65 @@ fn try_binary_op<
     Ok(out.put_tape(tape))
 }
 
-pub trait VecVecKernel<E: Dtype>: DeviceStorage {
-    fn forward<M: Dim, N: Dim>(
+pub trait MatMatKernel<E: Dtype>: DeviceStorage {
+    fn forward<M: Dim, K: Dim, N: Dim>(
         &self,
-        lhs: &Tensor<(M,), E, Self>,
-        rhs: &Tensor<(N,), E, Self>,
+        lhs: &Tensor<(M, K), E, Self>,
+        rhs: &Tensor<(K, N), E, Self>,
     ) -> Result<Tensor<(M, N), E, Self>, Self::Err>;
 
-    fn backward<M: Dim, N: Dim>(
+    fn backward<M: Dim, K: Dim, N: Dim>(
         &self,
-        lhs: &Tensor<(M,), E, Self>,
+        lhs: &Tensor<(M, K), E, Self>,
         grad_lhs: &mut Self::Vec<E>,
-        rhs: &Tensor<(N,), E, Self>,
+        rhs: &Tensor<(K, N), E, Self>,
         grad_rhs: &mut Self::Vec<E>,
         grad_out: &Self::Vec<E>,
     ) -> Result<(), Self::Err>;
 }
 
-impl<M: Dim, N: Dim, E: Dtype, D: VecVecKernel<E>, T: Tape<E, D> + Merge<R>, R: Tape<E, D>>
+impl<M: Dim, N: Dim, E: Dtype, D, T: Tape<E, D> + Merge<R>, R: Tape<E, D>>
     TryMatMul<Tensor<(N,), E, D, R>> for Tensor<(M,), E, D, T>
+where
+    D: MatMatKernel<E> + ReshapeKernel<E>,
 {
     type Output = Tensor<(M, N), E, D, T>;
     fn try_matmul(self, rhs: Tensor<(N,), E, D, R>) -> Result<Self::Output, Self::Err> {
-        try_binary_op(self, rhs, D::forward, D::backward)
+        let m = self.shape.0;
+        let n = rhs.shape.0;
+        let lhs = self.try_reshape_like(&(m, Const::<1>)).unwrap()?;
+        lhs.try_matmul(rhs.try_reshape_like(&(Const::<1>, n)).unwrap()?)
     }
 }
 
-pub trait VecMatKernel<E: Dtype>: DeviceStorage {
-    fn forward<K: Dim, N: Dim>(
-        &self,
-        lhs: &Tensor<(K,), E, Self>,
-        rhs: &Tensor<(K, N), E, Self>,
-    ) -> Result<Tensor<(N,), E, Self>, Self::Err>;
-
-    fn backward<K: Dim, N: Dim>(
-        &self,
-        lhs: &Tensor<(K,), E, Self>,
-        grad_lhs: &mut Self::Vec<E>,
-        rhs: &Tensor<(K, N), E, Self>,
-        grad_rhs: &mut Self::Vec<E>,
-        grad_out: &Self::Vec<E>,
-    ) -> Result<(), Self::Err>;
-}
-
-impl<K: Dim, N: Dim, E: Dtype, D: VecMatKernel<E>, T: Tape<E, D> + Merge<R>, R: Tape<E, D>>
+impl<K: Dim, N: Dim, E: Dtype, D, T: Tape<E, D> + Merge<R>, R: Tape<E, D>>
     TryMatMul<Tensor<(K, N), E, D, R>> for Tensor<(K,), E, D, T>
+where
+    D: MatMatKernel<E> + ReshapeKernel<E>,
 {
     type Output = Tensor<(N,), E, D, T>;
     fn try_matmul(self, rhs: Tensor<(K, N), E, D, R>) -> Result<Self::Output, Self::Err> {
-        assert_eq!(self.shape.0, rhs.shape.0);
-        try_binary_op(self, rhs, D::forward, D::backward)
+        let k1 = self.shape.0;
+        let (k2, n) = rhs.shape;
+        assert_eq!(k1, k2);
+        let lhs = self.try_reshape_like(&(Const::<1>, k1)).unwrap()?;
+        lhs.try_matmul(rhs)?.try_reshape_like(&(n,)).unwrap()
     }
 }
 
-pub trait MatMatKernel<E: Dtype>: DeviceStorage {
-    fn forward<M: Dim, K: Dim, N: Dim>(
-        &self,
-        lhs: &Tensor<(M, K), E, Self>,
-        rhs: &Tensor<(K, N), E, Self>,
-    ) -> Result<Tensor<(M, N), E, Self>, Self::Err>;
-
-    fn backward<M: Dim, K: Dim, N: Dim>(
-        &self,
-        lhs: &Tensor<(M, K), E, Self>,
-        grad_lhs: &mut Self::Vec<E>,
-        rhs: &Tensor<(K, N), E, Self>,
-        grad_rhs: &mut Self::Vec<E>,
-        grad_out: &Self::Vec<E>,
-    ) -> Result<(), Self::Err>;
+impl<M: Dim, K: Dim, E: Dtype, D, T: Tape<E, D> + Merge<R>, R: Tape<E, D>>
+    TryMatMul<Tensor<(K,), E, D, R>> for Tensor<(M, K), E, D, T>
+where
+    D: MatMatKernel<E> + ReshapeKernel<E>,
+{
+    type Output = Tensor<(M,), E, D, T>;
+    fn try_matmul(self, rhs: Tensor<(K,), E, D, R>) -> Result<Self::Output, Self::Err> {
+        let (m, k1) = self.shape;
+        let k2 = rhs.shape.0;
+        assert_eq!(k1, k2);
+        let rhs = rhs.try_reshape_like(&(k2, Const::<1>)).unwrap()?;
+        self.try_matmul(rhs)?.try_reshape_like(&(m,)).unwrap()
+    }
 }
 
 impl<M: Dim, K: Dim, N: Dim, E: Dtype, D: MatMatKernel<E>, T, R> TryMatMul<Tensor<(K, N), E, D, R>>
diff --git a/src/tensor_ops/utilities/device.rs b/src/tensor_ops/utilities/device.rs
index 57a8cad74..c209bd9fd 100644
--- a/src/tensor_ops/utilities/device.rs
+++ b/src/tensor_ops/utilities/device.rs
@@ -41,9 +41,7 @@ pub trait Device<E: Dtype>:
     + super::super::roll::RollKernel<E>
 
     // matmuls
-    + super::super::matmul::VecMatKernel<E>
     + super::super::matmul::MatMatKernel<E>
-    + super::super::matmul::VecVecKernel<E>
     + super::super::matmul::MatMatBrKernel<E>
     + super::super::matmul::MatMatBatch3Kernel<E>
     + super::super::matmul::MatMatBatch4Kernel<E>


Test patch: