Skip to content

WIP: A wrapper for jax.experimental.stax neural networks

Vicentini Filippo requested to merge github/fork/twesterhout/stax into v2.1

Created by: twesterhout

See test_machine.py for an example.

Still need to clean up the code a bit, but the PR is ready for feedback.

Merge request reports