carrasqu / JAX_GRU

A GRU written in JAX

JAX_GRU

This code is based on the implementation here: https://roberttlange.github.io/posts/2020/03/blog-post-10/