Stephen Spencer

icon company DeepMind icon location UK

Results 1 comments of Stephen Spencer

ImportError: cannot import name 'DeviceLocalLayout' from 'jax._src.layout' (/opt/conda/lib/python3.10/site-packages/jax/_src/layout.py)

Hi there. I don't really understand this issue, it doesn't look like you're importing chex here. Could you provide a step-by-step reproduction?