Add a benchmark of calling the Dex diagonal convolution from Jax.

google-research / dex-lang

Research language for array processing in the Haskell/ML family

BSD 3-Clause "New" or "Revised" License

1.58k stars 107 forks source link

Closed axch closed 2 years ago

axch commented 2 years ago

Add a baseline for it of doing the same thing in Jax (as best I can navigate the convolution docs).