AttributeError: module 'jax' has no attribute 'Array'

[SongDoHou]

I've encountered with different versions of JAX. Your assistance in resolving these would be greatly appreciated.

Issue with JAX version 0.3.13:
Error: AttributeError: module 'jax' has no attribute 'Array'.
Context: "from jaxman.planner.rl_planner.agent.core import restore_agent"

Issue after upgrading to JAX version 0.4.8:
Error: TypeError: split accepts a single key, but was given a key array of shape (2,) != (). Use jax.vmap for batching.
Code Snippet: This issue occurs at key_start_goal = jax.random.split(jax.random.PRNGKey(46)) in jaxman/env/navigation/instance.py:88.

Given these issues, I would greatly appreciate if you could advise on the appropriate versions of JAX and any related packages that would be compatible with the code snippet provided.

[tadashiK]

Thank you for raising this issue.

@HikaruAsano will look into details, but would you provide more contexts in the meantime? For example, what command did you use to run which script? Also, would you provide a full error message so that we can closely check the issue?

As far as I quickly checked, we do not explicitly use "Array" in the jaxman library, which seems to be causing the error in your environment.

[SongDoHou]

I found that the issues was the chex version.
It works well with '0.1.3' version. Thank you.