[jax] Improve exprel numerical stability and add gradient tests (#57)

* Bump version to 0.1.4 and improve numerical stability in exprel function

* Fix formatting in pyproject.toml for version and keywords

* Implement numerically stable exprel function and add comprehensive tests

* Update lax array creation to use jax.numpy for zero initialization and update CI jax version

* Remove redundant lax array creation tests for improved clarity

* Update changelog for version 0.1.4: add exprel function, improve lax array creation, and enhance CI compatibility

Author: chaoming0625

.github/workflows/CI-daily.yml

@@ -38,7 +38,7 @@ jobs:
       fail-fast: false
       matrix:
         python-version: ["3.13"]
-        jax-version: ["0.4.38", "0.5.2", "0.6.0", "0.7.2", ""]
+        jax-version: ["0.4.38", "0.5.2", "0.6.0", "0.7.2", "0.8.0", ""]
 
     steps:
       - name: Cancel Previous Runs

changelog.md

@@ -1,5 +1,12 @@
 # Release Notes
 
+## Version 0.1.4
+
+- Added numerically stable `exprel` function with comprehensive test coverage
+- Updated lax array creation to use `jax.numpy` for zero initialization
+- Updated CI JAX version for improved compatibility
+- Improved code quality and removed redundant tests
+
 ## Version 0.1.3
 
 - Compatible with `jax>=0.8.2`

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 
 [tool.setuptools.dynamic]
-version = {attr = "saiunit.__version__"}
+version = { attr = "saiunit.__version__" }
 
 [tool.setuptools.packages.find]
 exclude = [
@@ -44,7 +44,7 @@ classifiers = [
     'Topic :: Software Development :: Libraries',
 ]
 
-keywords = ['physical unit', 'physical quantity', 'brain modeling', 'scientific computing', 'AI for science',]
+keywords = ['physical unit', 'physical quantity', 'brain modeling', 'scientific computing', 'AI for science', ]
 
 dependencies = [
     'jax',
@@ -64,9 +64,7 @@ repository = 'https://github.com/chaobrain/saiunit'
 "Documentation" = "https://saiunit.readthedocs.io/"
 
 [project.optional-dependencies]
-testing = [
-    'pytest',
-]
+testing = ['pytest']
 cpu = ["jax[cpu]"]
 cuda12 = ["jax[cuda12]"]
 cuda13 = ["jax[cuda13]"]

saiunit/__init__.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.1.3"
+__version__ = "0.1.4"
 __version_info__ = tuple(map(int, __version__.split(".")))
 
 from . import _matplotlib_compat

saiunit/_compatible_import.py

@@ -24,13 +24,19 @@
     'safe_map',
     'unzip2',
     'wrap_init',
+    'Primitive',
 ]
 
 T = TypeVar("T")
 T1 = TypeVar("T1")
 T2 = TypeVar("T2")
 T3 = TypeVar("T3")
 
+if jax.__version_info__ < (0, 4, 38):
+    from jax.core import Primitive
+else:
+    from jax.extend.core import Primitive
+
 
 def wrap_init(fun: Callable, args: tuple, kwargs: dict, name: str):
     if jax.__version_info__ < (0, 6, 0):

saiunit/lax/_lax_array_creation.py

@@ -17,6 +17,7 @@
 
 import jax
 from jax import lax
+import jax.numpy as jnp
 
 from .._base import Unit, Quantity
 from .._misc import set_module_as, maybe_custom_array
@@ -28,7 +29,8 @@
     'zeros_like_array',
 
     # array creation(misc)
-    'iota', 'broadcasted_iota',
+    'iota',
+    'broadcasted_iota',
 ]
 
 
@@ -43,13 +45,13 @@ def zeros_like_array(
         if unit is not None:
             assert isinstance(unit, Unit), 'unit must be an instance of Unit.'
             x = x.in_unit(unit)
-        return Quantity(lax.zeros_like_array(x.mantissa), unit=x.unit)
+        return Quantity(jnp.zeros_like(x.mantissa), unit=x.unit)
     else:
         if unit is not None:
             assert isinstance(unit, Unit), 'unit must be an instance of Unit.'
-            return lax.zeros_like_array(x) * unit
+            return jnp.zeros_like(x) * unit
         else:
-            return lax.zeros_like_array(x)
+            return jnp.zeros_like(x)
 
 
 # array creation (misc)

saiunit/lax/_lax_array_creation_test.py

@@ -37,25 +37,6 @@ def __init__(self, *args, **kwargs):
 
         print()
 
-    @parameterized.product(
-        array=[jnp.array([1.0, 2.0]), jnp.array([[1.0, 2.0], [3.0, 4.0]])],
-        unit=[second, meter]
-    )
-    def test_lax_array_creation_given_array(self, array, unit):
-        bulax_fun_list = [getattr(bulax, fun) for fun in lax_array_creation_given_array]
-        lax_fun_list = [getattr(lax, fun) for fun in lax_array_creation_given_array]
-
-        for bulax_fun, lax_fun in zip(bulax_fun_list, lax_fun_list):
-            print(f'fun: {bulax_fun.__name__}')
-
-            result = bulax_fun(array)
-            expected = lax_fun(array)
-            assert_quantity(result, expected)
-
-            result = bulax_fun(array, unit=unit)
-            expected = lax_fun(array)
-            assert_quantity(result, expected, unit=unit)
-
     @parameterized.product(
         value=[1, 10, 100],
         unit=[second, meter]