distrax>=0.1.7
dm-haiku>=0.0.16
jaxlib>=0.8.11
jax>=0.8.1
optax>=0.1.7
tfp-nightly>=0.26.0.dev20250828
