jvp_matmul.dx

n = if dex_test_mode() then 10 else 500

m1 = rand_mat(n, n, randn, new_key 0)
m2 = rand_mat(n, n, randn, new_key 1)

def mmp'(m1:l=>m=>Float, m2:m=>n=>Float) -> l=>n=>Float given (l|Ix, m|Ix, n|Ix) =
  jvp (\m. m1 ** m) m2 m2

%bench "jvp_matmul"
res = mmp'(m1, m2)
>
> jvp_matmul
> Compile time: 82.255 ms
> Run time:     7.298 us 	(based on 1 run)