Repository: jax-ml/jax
Pull request #24081 (2024-10-02 20:40:26)

Implement jax.numpy.spacing
Somehow we've missed this numpy API up until now.

jakevdp: Hmm, this approach seems to be failing for inputs very close to zero, because `nextafter` returns a denormal, and `lax.sub` flushes denormals.
jakevdp: > How close are we talking?

Anything where the next higher number is a denormal. For example:
```python
>>> x = np.float32(0)

>>> np.spacing(x)
np.float32(1e-45)

jnp.spacing(x)
Array(0., dtype=float32)
```
This is because of the issue noted here: https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.nextafter.html
dfm: Interesting! When I run `spacing(np.float32(0))` after copying your new implementation into colab (where I was originally experimenting), it returns `Array(1.e-45, dtype=float32)`, just like the `np` version, but on my mac, I see `Array(0., dtype=float32)`, like you do. I guess this is related to the "in _some_ environments" part of that docstring?

It doesn't seem awesome to have environment-dependent behavior here, but perhaps it's not worth digging into it too deeply. Maybe it's ok to keep this implementation with the appropriate caveat in the docstring, unless you have an idea for an easy fix!
jakevdp: I think I managed to work around the denormal issue – we'll see if the tests pass. I also added special logic for `float16` so as to match NumPy's outputs (not sure whether we should do that, or just note that JAX's convention differs).

Agreed it's probably worth filing an upstream bug.

Patch:
diff --git a/.github/workflows/asan.yaml b/.github/workflows/asan.yaml
index ca2842edc..fd3054aab 100644
--- a/.github/workflows/asan.yaml
+++ b/.github/workflows/asan.yaml
@@ -8,48 +8,42 @@ on:
   schedule:
     - cron: "0 12 * * *" # Daily at 12:00 UTC
   workflow_dispatch: # allows triggering the workflow run manually
-  pull_request: # Automatically trigger on pull requests affecting this file
+  pull_request:  # Automatically trigger on pull requests affecting this file
     branches:
       - main
     paths:
-      - '**/workflows/asan.yml'
+      - '**workflows/asan.yml'
 
 jobs:
-  asan:
-    runs-on: linux-x86-n2-64
-    container:
-      image: index.docker.io/library/ubuntu@sha256:b359f1067efa76f37863778f7b6d0e8d911e3ee8efa807ad01fbf5dc1ef9006b # ratchet:ubuntu:24.04
+  upstream-dev:
+    runs-on: ubuntu-20.04-16core
     strategy:
       fail-fast: false
-    defaults:
-      run:
-        shell: bash -l {0}
+      matrix:
+        python-version: ["3.12"]
     steps:
-      - uses: actions/checkout@d632683dd7b4114ad314bca15554477dd762a938 # v4.2.0
+      - uses: actions/checkout@d632683dd7b4114ad314bca15554477dd762a938  # v4.2.0
         with:
           path: jax
-      - uses: actions/checkout@d632683dd7b4114ad314bca15554477dd762a938 # v4.2.0
+      - uses: actions/checkout@d632683dd7b4114ad314bca15554477dd762a938  # v4.2.0
         with:
           repository: python/cpython
           path: cpython
           ref: v3.12.6
       - name: Install clang 18
-        env:
-          DEBIAN_FRONTEND: noninteractive
         run: |
-          apt update
-          apt install -y clang-18 libstdc++-14-dev build-essential libssl-dev \
-            zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev curl git \
-            libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev \
-            libffi-dev liblzma-dev
+          wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
+          echo deb http://apt.llvm.org/focal/ llvm-toolchain-focal-18 main | sudo tee -a /etc/apt/sources.list
+          sudo apt update
+          sudo apt install clang-18 libstdc++-10-dev
       - name: Build CPython with ASAN enabled
         env:
           ASAN_OPTIONS: detect_leaks=0
         run: |
           cd cpython
           mkdir ${GITHUB_WORKSPACE}/cpythonasan
-          CC=clang-18 CXX=clang++-18 ./configure --prefix ${GITHUB_WORKSPACE}/cpythonasan --with-address-sanitizer --without-pymalloc
-          make -j64
+          CC=clang-18 ./configure --prefix ${GITHUB_WORKSPACE}/cpythonasan --with-address-sanitizer --without-pymalloc
+          make -j16
           make install
           ${GITHUB_WORKSPACE}/cpythonasan/bin/python3 -m venv ${GITHUB_WORKSPACE}/venv
       - name: Install JAX test requirements
@@ -65,10 +59,7 @@ jobs:
         run: |
           source ${GITHUB_WORKSPACE}/venv/bin/activate
           cd jax
-          python build/build.py \
-            --bazel_options=--color=yes
-            --bazel_options=--copt=-fsanitize=address \
-            --clang_path=/usr/bin/clang-18 \
+          python build/build.py --bazel_options=--copt=-fsanitize=address --clang_path=/usr/bin/clang-18
           pip install dist/jaxlib-*.whl
           pip install -e .
       - name: Run tests
@@ -86,3 +77,4 @@ jobs:
           echo "JAX_SKIP_SLOW_TESTS=$JAX_SKIP_SLOW_TESTS"
           # The LD_PRELOAD works around https://github.com/google/sanitizers/issues/934#issuecomment-649516500
           LD_PRELOAD=/lib/x86_64-linux-gnu/libstdc++.so.6 python -m pytest -n auto --tb=short --maxfail=20 tests
+ 
\ No newline at end of file
diff --git a/docs/jax.numpy.rst b/docs/jax.numpy.rst
index d6b7d74bd..9eb518464 100644
--- a/docs/jax.numpy.rst
+++ b/docs/jax.numpy.rst
@@ -376,6 +376,7 @@ namespace; they are listed below.
     size
     sort
     sort_complex
+    spacing
     split
     sqrt
     square
diff --git a/jax/_src/numpy/ufuncs.py b/jax/_src/numpy/ufuncs.py
index abeb7b775..d7afab1e7 100644
--- a/jax/_src/numpy/ufuncs.py
+++ b/jax/_src/numpy/ufuncs.py
@@ -1451,6 +1451,50 @@ def nextafter(x: ArrayLike, y: ArrayLike, /) -> Array:
   """
   return lax.nextafter(*promote_args_inexact("nextafter", x, y))
 
+
+@partial(jit, inline=True)
+def spacing(x: ArrayLike, /) -> Array:
+  """Return the spacing between ``x`` and the next adjacent number.
+
+  JAX implementation of :func:`numpy.spacing`.
+
+  Args:
+    x: real-valued array. Integer or boolean types will be cast to float.
+
+  Returns:
+    Array of same shape as ``x`` containing spacing between each entry of
+    ``x`` and its closest adjacent value.
+
+  See also:
+    - :func:`jax.numpy.nextafter`: find the next representable value.
+
+  Examples:
+    >>> x = jnp.array([0.0, 0.25, 0.5, 0.75, 1.0], dtype='float32')
+    >>> jnp.spacing(x)
+    Array([1.4012985e-45, 2.9802322e-08, 5.9604645e-08, 5.9604645e-08,
+          1.1920929e-07], dtype=float32)
+
+    For ``x = 1``, the spacing is equal to the ``eps`` value given by
+    :class:`jax.numpy.finfo`:
+
+    >>> x = jnp.float32(1)
+    >>> jnp.spacing(x) == jnp.finfo(x.dtype).eps
+    Array(True, dtype=bool)
+  """
+  arr, = promote_args_inexact("spacing", x)
+  if dtypes.isdtype(arr.dtype, "complex floating"):
+    raise ValueError("jnp.spacing is not defined for complex inputs.")
+  inf = _lax_const(arr, np.inf)
+  smallest_subnormal = dtypes.finfo(arr.dtype).smallest_subnormal
+
+  # Numpy's behavior seems to depend on dtype
+  if arr.dtype == 'float16':
+    return lax.nextafter(arr, inf) - arr
+  else:
+    result = lax.nextafter(arr, copysign(inf, arr)) - arr
+    return _where(result == 0, copysign(smallest_subnormal, arr), result)
+
+
 # Logical ops
 @partial(jit, inline=True)
 def _logical_and(x: ArrayLike, y: ArrayLike, /) -> Array:
diff --git a/jax/numpy/__init__.py b/jax/numpy/__init__.py
index 20c37c559..bd8068729 100644
--- a/jax/numpy/__init__.py
+++ b/jax/numpy/__init__.py
@@ -443,6 +443,7 @@ from jax._src.numpy.ufuncs import (
     sin as sin,
     sinc as sinc,
     sinh as sinh,
+    spacing as spacing,
     sqrt as sqrt,
     square as square,
     subtract as subtract,
diff --git a/jax/numpy/__init__.pyi b/jax/numpy/__init__.pyi
index fd2524bc3..0ea9b5ee7 100644
--- a/jax/numpy/__init__.pyi
+++ b/jax/numpy/__init__.pyi
@@ -808,6 +808,7 @@ def sort(
     order: None = ...,
 ) -> Array: ...
 def sort_complex(a: ArrayLike) -> Array: ...
+def spacing(x: ArrayLike, /) -> Array: ...
 def split(
     ary: ArrayLike,
     indices_or_sections: int | Sequence[int] | ArrayLike,
diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl
index e1d03f5f6..37861d98e 100644
--- a/third_party/xla/workspace.bzl
+++ b/third_party/xla/workspace.bzl
@@ -21,8 +21,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 #    curl -L https://github.com/openxla/xla/archive/<git hash>.tar.gz | sha256sum
 #    and update XLA_SHA256 with the result.
 
-XLA_COMMIT = "0b058d5ddeae050c9dc3ed93ac179f07037f840b"
-XLA_SHA256 = "29d8a871b9920b47c3659479fb1577a08cce19ac93dfe1dceab86b6fbeb25b55"
+XLA_COMMIT = "cd6e808c59f53b40a99df1f1b860db9a3e598bff"
+XLA_SHA256 = "65cb6d63ef4083b35775052636cb9c629f86db6947c8b91711923ba31dbdcde8"
 
 def repo():
     tf_http_archive(


Test patch:
diff --git a/tests/lax_numpy_operators_test.py b/tests/lax_numpy_operators_test.py
index 45a780c9f..744a99fb7 100644
--- a/tests/lax_numpy_operators_test.py
+++ b/tests/lax_numpy_operators_test.py
@@ -126,6 +126,8 @@ JAX_ONE_TO_ONE_OP_RECORDS = [
     op_record("negative", 1, number_dtypes, all_shapes, jtu.rand_default, ["rev"]),
     op_record("nextafter", 2, [f for f in float_dtypes if f != jnp.bfloat16],
               all_shapes, jtu.rand_default, ["rev"], inexact=True, tolerance=0),
+    op_record("spacing", 1, float_dtypes, all_shapes, jtu.rand_default, ["rev"],
+              inexact=True, tolerance=0),
     op_record("not_equal", 2, all_dtypes, all_shapes, jtu.rand_some_equal, ["rev"]),
     op_record("array_equal", 2, number_dtypes, all_shapes, jtu.rand_some_equal, ["rev"]),
     op_record("array_equiv", 2, number_dtypes, all_shapes, jtu.rand_some_equal, ["rev"]),
@@ -701,6 +703,34 @@ class JaxNumpyOperatorTests(jtu.JaxTestCase):
     dx = jax.grad(jax.numpy.i0)(0.0)
     self.assertArraysEqual(dx, 0.0)
 
+  @jtu.sample_product(
+      shape=all_shapes,
+      dtype=default_dtypes,
+  )
+  def testSpacingIntegerInputs(self, shape, dtype):
+    rng = jtu.rand_int(self.rng(), low=-64, high=64)
+    args_maker = lambda: [rng(shape, dtype)]
+    computation_dtype = jnp.spacing(rng(shape, dtype)).dtype
+    np_func = lambda x: np.spacing(np.array(x).astype(computation_dtype))
+    self._CheckAgainstNumpy(np_func, jnp.spacing, args_maker, check_dtypes=True, tol=0)
+    self._CompileAndCheck(jnp.spacing, args_maker, tol=0)
+
+  @jtu.sample_product(dtype = float_dtypes)
+  @jtu.skip_on_devices("tpu")
+  def testSpacingSubnormals(self, dtype):
+    zero = np.array(0, dtype=dtype)
+    inf = np.array(np.inf, dtype=dtype)
+    x = [zero]
+    for i in range(5):
+      x.append(np.nextafter(x[-1], -inf))  # negative denormals
+    x = x[::-1]
+    for i in range(5):
+      x.append(np.nextafter(x[-1], inf))  # positive denormals
+    x = np.array(x, dtype=dtype)
+    args_maker = lambda: [x]
+    self._CheckAgainstNumpy(np.spacing, jnp.spacing, args_maker, check_dtypes=True, tol=0)
+    self._CompileAndCheck(jnp.spacing, args_maker, tol=0)
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())
diff --git a/tests/logging_test.py b/tests/logging_test.py
index a1d6695a1..70f619de5 100644
--- a/tests/logging_test.py
+++ b/tests/logging_test.py
@@ -92,14 +92,10 @@ class LoggingTest(jtu.JaxTestCase):
       python = sys.executable
       assert "python" in python
       env_variables = {"TF_CPP_MIN_LOG_LEVEL": "1"}
-      if os.getenv("ASAN_OPTIONS"):
-        env_variables["ASAN_OPTIONS"] = os.getenv("ASAN_OPTIONS")
       if os.getenv("PYTHONPATH"):
         env_variables["PYTHONPATH"] = os.getenv("PYTHONPATH")
       if os.getenv("LD_LIBRARY_PATH"):
         env_variables["LD_LIBRARY_PATH"] = os.getenv("LD_LIBRARY_PATH")
-      if os.getenv("LD_PRELOAD"):
-        env_variables["LD_PRELOAD"] = os.getenv("LD_PRELOAD")
       # Make sure C++ logging is at default level for the test process.
       proc = subprocess.run(
           [python, f.name],