Add JAX train/predict support

trivoldus28 commented 2 years ago

Notes:

BufferPredict is factored out of the Tensorflow Predict and then used as a parent class for JAX Predict. This should make using BufferPredict easier for other implementations.
Train was written based off PyTorch implementation. Since JAX is rather low level and allowing many freedom for model implementation, to make things easier for us, we make a simple GenericJaxModel interface for models to follow in order to train or predict. A model implementing this interface will need to contain not only the forward model but also loss and update fn. Some examples can be found in https://github.com/funkelab/funlib.learn.jax

Pre-Merge Checklist:

[ ] if this PR adds a feature: request merge into the latest development branch (vX.Y-dev)
[ ] if this PR fixes a bug: request merge into the latest patch branch (patch-X.Y.Z)
[ ] PR branch name is short and describes the feature/bug addressed in this PR
[ ] commit messages are consistent
[ ] changes reviewed by another contributor
[ ] tests cover changes
[ ] all tests pass

trivoldus28 commented 2 years ago

@pattonw I tried to resolve much of your comments... can you take a second look? Thanks!

pattonw commented 2 years ago

Documentation is nice and easy to follow! Thanks

pattonw commented 2 years ago

Would still like to see tests for the Jax nodes to provide examples on how to use them.

pattonw commented 2 years ago

Looks good to me now. I think we're good to merge

pattonw commented 2 years ago

Actually 1 last thing, can you add the new nodes to the docs? That can be done in the api.rst file

trivoldus28 commented 2 years ago

Merging to v1.3-dev with @pattonw 's approval.

funkelab / gunpowder