jax update
update jax to >=0.6.1, and make neural dependencies non-optional for a quick fix. In future times, neural part will be moved to a different repository.
update jax to >=0.6.1, and make neural dependencies non-optional for a quick fix. In future times, neural part will be moved to a different repository.