Running JAX-MD molecular dynamics / particle simulation in the browser
- Minimal tflite (not working - issue): https://josephrocca.github.io/jax-md-web-test/minimal/tflite.html
- Minimal ONNX - awaiting tf2onnx support for PartitionedCall op (issue). Here's the conversion notebook.
- Here's the conversion notebook that goes from JAX to tflite via jax2tf and tf