Nested vmap over HopperHop fails

Short script to reproduce:
```
import jax
import mujoco_playground

env = mujoco_playground.registry.load("HopperHop")

def reset_one_seed(seed_key):
return jax.vmap(env.reset)(jax.random.split(seed_key, 4))

keys = jax.random.split(jax.random.PRNGKey(0), 2)

jax.jit(jax.vmap(reset_one_seed)).lower(keys).compile()
```
Stack trace:
```
Traceback (most recent call last):
File "", line 1, in
File "PATH_TO_ENV/.venv/lib/python3.10/site-packages/jax/_src/stages.py", line 649, in compile
self._lowering.compile(**kw), # pytype: disable=wrong-keyword-args
File "PATH_TO_ENV/.venv/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py", line 2448, in compile
executable = UnloadedMeshExecutable.from_hlo(
File "PATH_TO_ENV/.venv/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py", line 2967, in from_hlo
xla_executable = _cached_compilation(
File "PATH_TO_ENV/.venv/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py", line 2758, in _cached_compilation
xla_executable = compiler.compile_or_get_cached(
File "PATH_TO_ENV/.venv/lib/python3.10/site-packages/jax/_src/compiler.py", line 470, in compile_or_get_cached
return _compile_and_write_cache(
File "PATH_TO_ENV/.venv/lib/python3.10/site-packages/jax/_src/compiler.py", line 687, in _compile_and_write_cache
executable = backend_compile(
File "PATH_TO_ENV/.venv/lib/python3.10/site-packages/jax/_src/profiler.py", line 334, in wrapper
return func(*args, **kwargs)
File "PATH_TO_ENV/.venv/lib/python3.10/site-packages/jax/_src/compiler.py", line 327, in backend_compile
raise e
File "PATH_TO_ENV/.venv/lib/python3.10/site-packages/jax/_src/compiler.py", line 321, in backend_compile
return backend.compile(built_c, compile_options=options)
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: during context [hlo verifier]: Expected instruction to have shape equal to f32[2,2,4,6,3], actual shape is f32[4,2,2,6,3]:
%dot.25 = f32[4,2,2,6,3]{4,3,1,0,2} dot(%reshape.2384, %reshape.2389), lhs_batch_dims={0,1,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,1,2}, rhs_contracting_dims={3}, metadata={op_name="jit(reset_one_seed)/jit(main)/vmap(vmap(forward))/dot_general" source_file="PATH_TO_ENV/.venv/lib/python3.10/site-packages/mujoco/mjx/_src/sensor.py" source_line=512}

Failed after simplification
```

On the other hand, on some other envs this script works, see e.g. this:
```
import jax
import mujoco_playground
env = mujoco_playground.registry.load("CheetahRun")

def reset_one_seed(seed_key):
return jax.vmap(env.reset)(jax.random.split(seed_key, 4))

keys = jax.random.split(jax.random.PRNGKey(0), 2)

jax.jit(jax.vmap(reset_one_seed)).lower(keys).compile()
```
Versions:
```
python 3.10
jax==0.5.3
jaxlib==0.5.3
mujoco==3.3.7
```