update jax grid

alanlujan91 · Oct 13, 2023 · d9f15d4 · d9f15d4
1 parent c5fe27e
commit d9f15d4
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/multinterp/regular.py b/src/multinterp/regular.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import cupy as cp
+import jax.numpy as jnp
 import numpy as np
 
 from multinterp.backend._numba import numba_get_coordinates, numba_map_coordinates
@@ -41,6 +42,7 @@ def get_methods():
 
         get_coords["jax"] = jax_get_coordinates
         map_coords["jax"] = jax_map_coordinates
+        get_grad["jax"] = jnp.gradient
     except ImportError:
         pass