Repository: ecmwf/earthkit-data
Pull request #533 (2024-11-20 12:00:02)

Add jax backend
Add support for jax as a backend

codecov-commenter: ## [Codecov](https://app.codecov.io/gh/ecmwf/earthkit-data/pull/533?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf) Report
Attention: Patch coverage is `25.00000%` with `15 lines` in your changes missing coverage. Please review.
> Project coverage is 90.35%. Comparing base [(`837f240`)](https://app.codecov.io/gh/ecmwf/earthkit-data/commit/837f240f2ff5e3dc2dc985d24f3be6b52038f1c0?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf) to head [(`5519a15`)](https://app.codecov.io/gh/ecmwf/earthkit-data/commit/5519a15afcbe42b8aedc55e6db9754d827030688?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf).

| [Files with missing lines](https://app.codecov.io/gh/ecmwf/earthkit-data/pull/533?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf) | Patch % | Lines |
|---|---|---|
| [tests/utils/test\_array.py](https://app.codecov.io/gh/ecmwf/earthkit-data/pull/533?src=pr&el=tree&filepath=tests%2Futils%2Ftest_array.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf#diff-dGVzdHMvdXRpbHMvdGVzdF9hcnJheS5weQ==) | 25.00% | [15 Missing :warning: ](https://app.codecov.io/gh/ecmwf/earthkit-data/pull/533?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf) |

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##           develop     #533      +/-   ##
===========================================
- Coverage    90.47%   90.35%   -0.13%     
===========================================
  Files          148      148              
  Lines        10838    10857      +19     
  Branches       511      511              
===========================================
+ Hits          9806     9810       +4     
- Misses         868      883      +15     
  Partials       164      164              
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/ecmwf/earthkit-data/pull/533?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf).

----
🚨 Try these New Features:

- [Flaky Tests Detection](https://docs.codecov.com/docs/test-result-ingestion-beta) - Detect and resolve failed and flaky tests

Patch:
diff --git a/src/earthkit/data/utils/array.py b/src/earthkit/data/utils/array.py
index caf55e4..30765ff 100644
--- a/src/earthkit/data/utils/array.py
+++ b/src/earthkit/data/utils/array.py
@@ -167,6 +167,40 @@ class PytorchBackend(ArrayBackend):
         return {"float64": torch.float64, "float32": torch.float32}
 
 
+class JaxBackend(ArrayBackend):
+    _name = "jax"
+    _module_name = "jax"
+
+    def _make_sample(self):
+        import jax.numpy as jarray
+
+        return jarray.ones(2)
+
+    @cached_property
+    def module(self):
+        import jax.numpy as jarray
+
+        return jarray
+
+    def to_numpy(self, v):
+        return v.numpy()
+
+    def from_numpy(self, v):
+
+        return self.from_other(v)
+
+    def from_other(self, v, **kwargs):
+        import jax.numpy as jarray
+
+        return jarray.array(v, **kwargs)
+
+    @cached_property
+    def dtypes(self):
+        import jax.numpy as jarray
+
+        return {"float64": jarray.float64, "float32": jarray.float32}
+
+
 class CupyBackend(ArrayBackend):
     _name = "cupy"
     _module_name = "cupy"
@@ -203,8 +237,9 @@ class CupyBackend(ArrayBackend):
 _NUMPY = NumpyBackend()
 _PYTORCH = PytorchBackend()
 _CUPY = CupyBackend()
+_JAX = JaxBackend()
 
-_BACKENDS = [_NUMPY, _PYTORCH, _CUPY]
+_BACKENDS = [_NUMPY, _PYTORCH, _CUPY, _JAX]
 _BACKENDS_BY_NAME = {v._name: v for v in _BACKENDS}
 _BACKENDS_BY_MODULE = {v._module_name: v for v in _BACKENDS}
 
@@ -226,6 +261,8 @@ def backend_from_array(array, raise_exception=True):
             return _PYTORCH
         elif _NAMESPACE.api.is_cupy_array(array):
             return _CUPY
+        elif _NAMESPACE.api.is_jax_array(array):
+            return _JAX
 
     if raise_exception:
         raise ValueError(f"Can't find namespace for array type={type(array)}")


Test patch:
diff --git a/src/earthkit/data/testing.py b/src/earthkit/data/testing.py
index 76d3562..01c9afc 100644
--- a/src/earthkit/data/testing.py
+++ b/src/earthkit/data/testing.py
@@ -123,6 +123,7 @@ NO_COVJSONKIT = not modules_installed("covjsonkit")
 NO_PYTORCH = not modules_installed("torch")
 NO_RIOXARRAY = not modules_installed("rioxarray")
 NO_CUPY = not modules_installed("cupy")
+NO_JAX = not modules_installed("jax")
 if not NO_CUPY:
     try:
         import cupy as cp
diff --git a/tests/utils/test_array.py b/tests/utils/test_array.py
index b5d67d8..91a8c27 100644
--- a/tests/utils/test_array.py
+++ b/tests/utils/test_array.py
@@ -12,8 +12,10 @@
 import pytest
 
 from earthkit.data.testing import NO_CUPY
+from earthkit.data.testing import NO_JAX
 from earthkit.data.testing import NO_PYTORCH
 from earthkit.data.utils.array import _CUPY
+from earthkit.data.utils.array import _JAX
 from earthkit.data.utils.array import _NUMPY
 from earthkit.data.utils.array import _PYTORCH
 from earthkit.data.utils.array import get_backend
@@ -75,7 +77,7 @@ def test_utils_array_backend_pytorch():
     assert np.isclose(b.namespace.mean(v), 1.0)
 
 
-@pytest.mark.skipif(NO_CUPY, reason="No pytorch installed")
+@pytest.mark.skipif(NO_CUPY, reason="No cupy installed")
 def test_utils_array_backend_cupy():
     b = get_backend("cupy")
     assert b.name == "cupy"
@@ -102,6 +104,33 @@ def test_utils_array_backend_cupy():
     assert np.isclose(b.namespace.mean(v), 1.0)
 
 
+@pytest.mark.skipif(NO_JAX, reason="No jax installed")
+def test_utils_array_backend_jax():
+    b = get_backend("jax")
+    assert b.name == "jax"
+    assert b is _JAX
+
+    import jax.numpy as ja
+    import numpy as np
+
+    v = ja.ones(10, dtype=ja.float64)
+    v_np = np.ones(10, dtype=np.float64)
+    v_lst = [1.0] * 10
+
+    # assert b.is_native_array(v)
+    # assert id(b.from_backend(v, b)) == id(v)
+    # assert id(b.from_backend(v, None)) == id(v)
+    assert ja.allclose(b.from_numpy(v_np, None), v)
+    assert ja.allclose(b.from_other(v_lst, dtype=ja.float64), v)
+    assert get_backend(v) is b
+
+    r = b.to_numpy(v)
+    assert isinstance(r, np.ndarray)
+    assert np.allclose(r, v_np)
+
+    assert np.isclose(b.namespace.mean(v), 1.0)
+
+
 if __name__ == "__main__":
     from earthkit.data.testing import main