JAX 0.7 appears to be too new for this project

[KyleVasulka]

AttributeError: jax.interpreters.xla.pytype_aval_mappings was deprecated in JAX v0.5.0 and removed in JAX v0.7.0. jax.core.pytype_aval_mappings can be used as a replacement in most cases.

Steps to reproduce:
Basically a fresh clone/install and run causes this

git clone ksim-gym
cd ksim-gym
python -m venv ./venv
. ./venv/bin/activate
pip install -r requirements.txt
python -m train run_mode=view

Note: I am on an m4 mac running on cpu.

I was able to get it to run by adding the following to my requirements.txt file

jax==0.6.2

[LinFei83]

I encountered the same error as you when running python -m train --help with JAX 0.7.

Following the error message, I located the file /home/lf/anaconda3/envs/orca_gym/lib/python3.11/site-packages/tensorflow_probability/python/internal/backend/jax/ops.py and changed lines 680 and 681 from:

jax.interpreters.xla.pytype_aval_mappings[NumpyVariable] = (
    jax.interpreters.xla.pytype_aval_mappings[onp.ndarray])

to the following:

jax.core.pytype_aval_mappings[NumpyVariable] = (
    jax.core.pytype_aval_mappings[onp.ndarray])

Note: I am using a GPU.

[LinFei83]

I encountered the same error as you when running python -m train --help with JAX 0.7.

Following the error message, I located the file /home/lf/anaconda3/envs/orca_gym/lib/python3.11/site-packages/tensorflow_probability/python/internal/backend/jax/ops.py and changed lines 680 and 681 from:

jax.interpreters.xla.pytype_aval_mappings[NumpyVariable] = (
jax.interpreters.xla.pytype_aval_mappings[onp.ndarray])
to the following:

jax.core.pytype_aval_mappings[NumpyVariable] = (
jax.core.pytype_aval_mappings[onp.ndarray])
Note: I am using a GPU.

This method may not completely solve the problem after being used. A new error occurred when running train.py

[LinFei83]

I encountered the same error as you when running python -m train --help with JAX 0.7.
Following the error message, I located the file /home/lf/anaconda3/envs/orca_gym/lib/python3.11/site-packages/tensorflow_probability/python/internal/backend/jax/ops.py and changed lines 680 and 681 from:
jax.interpreters.xla.pytype_aval_mappings[NumpyVariable] = (
jax.interpreters.xla.pytype_aval_mappings[onp.ndarray])
to the following:
jax.core.pytype_aval_mappings[NumpyVariable] = (
jax.core.pytype_aval_mappings[onp.ndarray])
Note: I am using a GPU.

This method may not completely solve the problem after being used. A new error occurred when running train.py

My latest solution is #52