Fixes to work on more recent versions of jax

[NathanHowell]

optix has been removed from jax and moved to a new optax package. I also relaxed some version constraints so I can install this with poetry.

[toshikwa]

Hi @NathanHowell

Thank you for your contribution :)
I will review this PR in a few days.

One thing, dm-haiku is being aggressively improved, and there is concern that it may not be backward compatible.
Also, I would like to specify the versions of jax and optax.

Anyway, I really appreciate your help :)

[NathanHowell]

@ku2482 I've updated the minimum versions for jax and optax. wrt Haiku, that's ok. if I end up using this library in a serious matter I can pin haiku, I'm currently using jax, jaxlib and (sometimes) trax.

[kew96]

Wanted to check on the status of this. Is there an issue with adding jaxlib>=0.1.59 to requirements.txt? I'm running into an issue and have had to modify this all locally so figured I'd try to help some people out.

[toshikwa]

@kew96

Thank you for your comment.

Is there an issue with adding jaxlib>=0.1.59 to requirements.txt? I'm running into an issue and have had to modify this all locally so figured I'd try to help some people out.

I haven't tested this pull request, but I want to make it compatible with the newer version of jax.
So if you are interested in it, please make new PR. If you need some help, please raise the issue and we can discuss there.

Anyway, I really appreciate your interest.