SiEPIC not working with JAX version 0.4.24

BYUCamachoLab / simphony

A simulator for photonic integrated circuits.

https://simphonyphotonics.rtfd.io

Other

119 stars 35 forks source link

SiEPIC not working with JAX version 0.4.24 #108

Closed Kricki closed 8 months ago

Kricki commented 9 months ago

Simphony version 0.7.1

In JAX version 0.4.24, the deprecated function jax.interpreters.xla: backend_specific_translations has been reomoved (see here)

This leads the import from simphony.libraries import siepic to fail. Error message: module 'jax.interpreters.xla' has no attribute "backend_specific_translations"

With JAX version 0.4.23 Simphony still works.

sequoiap commented 8 months ago

Could not reproduce the error with jax==0.4.24. Did encounter a different error with an old version of flax that was installed:

ImportError: cannot import name 'linear_util' from 'jax' (simphony/.venv/lib/python3.11/site-packages/jax/__init__.py)

This was solved by simply updating flax.

sequoiap commented 8 months ago

Also, it seems that jax.interpreters.xla was added to 0.4.25 (see here, still deprecated and will be removed again in the future, but we don't directly import that anywhere in simphony.