flaport / sax

S + Autograd + XLA :: S-parameter based frequency domain circuit simulations and optimizations using JAX.
https://flaport.github.io/sax
Apache License 2.0
71 stars 17 forks source link

jax update breaks sax installation #29

Closed cpascual closed 9 months ago

cpascual commented 9 months ago

With jax update to 0.4.24, the installation of sax is broken:

$ pip install sax
   ...
$ python -c 'import sax' 
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File ".../python3.11/site-packages/sax/__init__.py", line 19, in <module>
    from . import backends as backends
  File ".../python3.11/site-packages/sax/backends/__init__.py", line 40, in <module>
    from .klu import analyze_circuit_klu, analyze_instances_klu, evaluate_circuit_klu
  File ".../python3.11/site-packages/sax/backends/klu.py", line 9, in <module>
    import klujax
  File ".../python3.11/site-packages/klujax.py", line 164, in <module>
    @xla_register_cpu(coo_mul_vec_c128, klujax_cpp.coo_mul_vec_c128)
     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../python3.11/site-packages/klujax.py", line 64, in decorator
    xla.backend_specific_translations["cpu"][primitive] = partial(fun, name)
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: module 'jax.interpreters.xla' has no attribute 'backend_specific_translations'

Not sure if this should be fixed at sax level (by pinning jax and jaxlib to 0.4.23 as here) of it it should be fixed in klujax

I can confirm that using pip install sax jax==0.4.23 jaxlib==0.4.23 works fine.

flaport commented 9 months ago

I'll probably have to fix klujax. I'll get back when it's done.

flaport commented 9 months ago

Fixed in 0.11.4.