absl-py
cvxpy
dm-tree
jax>=0.4.6
jaxlib>=0.4.6
numpy
optax
dm-haiku
einshape @ git+https://github.com/deepmind/einshape.git
ml_collections
