jax
flax
numpy
matplotlib
imageio

[dev]
black
pre-commit
