solved an unsupported deprecation, ops.index. replaced to jnp.index_e…

…xp. See jax-ml/jax#10293.
HajimeKawahara · Apr 14, 2022 · 4b350aa · 4b350aa
1 parent b561111
commit 4b350aa
Show file tree

Hide file tree

Showing 5 changed files with 5 additions and 5 deletions.
diff --git a/src/exojax/experimental/ditfnu.py b/src/exojax/experimental/ditfnu.py
@@ -15,7 +15,7 @@
 from jax.lax import scan
 import tqdm
 from exojax.spec.ditkernel import folded_voigt_kernel
-from jax.ops import index as joi
+from jax.numpy import index_exp as joi
 from exojax.spec.dit import getix
 
 

diff --git a/src/exojax/experimental/moditfnu.py b/src/exojax/experimental/moditfnu.py
@@ -15,7 +15,7 @@
 from jax.lax import scan
 import tqdm
 from exojax.spec.ditkernel import folded_voigt_kernel_logst
-from jax.ops import index as joi
+from jax.numpy import index_exp as joi
 from exojax.spec.dit import getix
 
 

diff --git a/src/exojax/experimental/redit.py b/src/exojax/experimental/redit.py
@@ -6,7 +6,7 @@
 from jax.lax import scan
 import tqdm
 from exojax.spec.ditkernel import folded_voigt_kernel_logst
-from jax.ops import index as joi
+from jax.numpy import index_exp as joi
 from exojax.spec.dit import getix
 from exojax.spec.modit import inc2D_givenx
 from exojax.spec.lpf import voigt

diff --git a/src/exojax/spec/dit.py b/src/exojax/spec/dit.py
@@ -9,7 +9,7 @@
 from jax import jit, vmap
 from jax.lax import scan
 from exojax.spec.ditkernel import fold_voigt_kernel
-from jax.ops import index as joi
+from jax.numpy import index_exp as joi
 from exojax.spec.atomll import padding_2Darray_for_each_atom
 from exojax.spec.rtransfer import dtauM
 

diff --git a/src/exojax/spec/modit.py b/src/exojax/spec/modit.py
@@ -12,7 +12,7 @@
 from exojax.spec.ditkernel import fold_voigt_kernel_logst
 from exojax.spec.ditkernel import voigt_kernel_logst
 
-from jax.ops import index as joi
+from jax.numpy import index_exp as joi
 from exojax.spec.dit import getix
 
 # exomol