mctx
dm-haiku
optax
