flax

1f3cf2bd - Removes jraph dependency

Removes jraph dependency This was added in #1470 and currently impedes the build. The problem is that jraph does not yet have a stable release.