Repository: google/jax
Pull request #21394 (2024-05-23 12:18:33)

Implement LRU cache eviction for persistent compilation cache
This PR is part of the implementation of LRU cache eviction using the `mtime` attribute provided by the filesystem. The current PR does not support GCS, but this problem will be solved in a subsequent PR.

More details in the design doc.

hawkinsp: > Note, I saw some NFS server being configured to not update mtime to speed up the server. Maybe document that this can happen and in that case, this will revert to creation time? The first time I saw the behavior without knowing the reason, it took times to understand what was going on.

Yes, I suspect there's a chance you might see stale `mtime` values if you stick the cache on NFS and you're accessing it concurrently from multiple clients (see `lookupcache` in the NFS docs). I'm not sure there's a lot we can do about that, though.
skye: This is a first cut at the LRU eviction implementation, so it isn't expected to work well with network file systems yet (notably GCS, which many Cloud TPU users use for their cache storage). We'll iterate from here. I don't think we should publicly document this until it works well across filesystems, but absolutely agree this should eventually be in https://jax.readthedocs.io/en/latest/persistent_compilation_cache.html.
ayaka14732: Test fails because the test utilises `filelock`, which is not installed.

```
=================================== FAILURES ===================================
_______________________ LRUCacheTest.test_cache_eviction _______________________
[gw5] linux -- Python 3.12.3 /opt/hostedtoolcache/Python/3.12.3/x64/bin/python
tests/lru_cache_test.py:39: in test_cache_eviction
    cache = Impl(path, max_size=884700)
jax/_src/lru_cache.py:57: in __init__
    raise RuntimeError("Please install filelock package to set `jax_compilation_cache_max_size`")
E   RuntimeError: Please install filelock package to set `jax_compilation_cache_max_size`
```
superbobry: Add filelock to build/test-requirements.txt?
ayaka14732: Linter complains about [import-not-found]:

```
jax/_src/lru_cache.py:27: error: Cannot find implementation or library stub for module named "filelock"  [import-not-found]
jax/_src/lru_cache.py:175: error: Cannot find implementation or library stub for module named "fsspec.utils"  [import-not-found]
jax/_src/lru_cache.py:175: error: Cannot find implementation or library stub for module named "fsspec"  [import-not-found]
jax/_src/compilation_cache.py:80: error: Cannot find implementation or library stub for module named "fsspec.utils"  [import-not-found]
jax/_src/compilation_cache.py:80: note: See https://mypy.readthedocs.io/en/stable/running_mypy.html#missing-imports
jax/_src/compilation_cache.py:80: error: Cannot find implementation or library stub for module named "fsspec"  [import-not-found]
jax/_src/compilation_cache.py:215: error: Argument 1 to "extract_executable_and_time" has incompatible type "bytes | None"; expected "bytes"  [arg-type]
Found 6 errors in 2 files (checked 501 source files)
```
ayaka14732: Just realised that JAX had a `FileSystemCache` that supports LRU cache eviction introduced in https://github.com/google/jax/pull/6869, but was subsequently removed in https://github.com/google/jax/pull/10771 to support GCS. This is exactly one of the challenges that I faced in this PR. Fortunately, I've devised potential solutions to simultaneously support LRU cache eviction and GCS compatibility. This is going to be completely solved in a subsequent PR.
ayaka14732: All comments resolved
ayaka14732: All review comments done
skye: Please fix the failing lint_and_typecheck (looks like you need to fix a type annotation?)

Patch:
diff --git a/jax/BUILD b/jax/BUILD
index 70a9eda86..81050790b 100644
--- a/jax/BUILD
+++ b/jax/BUILD
@@ -390,6 +390,7 @@ pytype_strict_library(
         ":compilation_cache_interface",
         ":config",
         ":gfile_cache",
+        ":lru_cache",
         ":monitoring",
         ":path",
         "//jax/_src/lib",
@@ -415,6 +416,15 @@ pytype_strict_library(
     ],
 )
 
+pytype_strict_library(
+    name = "lru_cache",
+    srcs = ["_src/lru_cache.py"],
+    deps = [
+        ":config",
+        ":compilation_cache_interface",
+    ] + py_deps("filelock"),
+)
+
 pytype_strict_library(
     name = "config",
     srcs = ["_src/config.py"],
diff --git a/jax/_src/compilation_cache.py b/jax/_src/compilation_cache.py
index 7884cda7c..9c2761517 100644
--- a/jax/_src/compilation_cache.py
+++ b/jax/_src/compilation_cache.py
@@ -34,6 +34,7 @@ from jax._src import monitoring
 from jax._src.gfile_cache import GFileCache
 from jax._src.lib import xla_client
 from jax._src.lib.mlir import ir
+from jax._src.lru_cache import LRUCache
 
 
 logger = logging.getLogger(__name__)
@@ -66,7 +67,20 @@ def set_once_cache_used(f) -> None:
 
 def get_file_cache(path: str) -> tuple[CacheInterface, str] | None:
   """Returns the file cache and the path to the cache."""
-  return GFileCache(path), path
+
+  def is_local_filesystem(path: str) -> bool:
+    return path.startswith("file://") or "://" not in path
+
+  # `LRUCache` currently only supports local filesystem. Therefore, if `path`
+  # is not on a local filesystem, instead of using `LRUCache`, we
+  # fallback to the old `GFileCache`, which does not support LRU eviction.
+  # TODO(ayx): Add cloud storage support for `LRUCache`, so that all these code
+  # can be removed.
+  if not is_local_filesystem(path):
+    return GFileCache(path), path
+
+  max_size = config.compilation_cache_max_size.value
+  return LRUCache(path, max_size=max_size), path
 
 
 def set_cache_dir(path) -> None:
diff --git a/jax/_src/config.py b/jax/_src/config.py
index ba4c18fa0..949da0358 100644
--- a/jax/_src/config.py
+++ b/jax/_src/config.py
@@ -1281,6 +1281,18 @@ compilation_cache_dir = define_optional_string_state(
           '2. The value of this flag set in the command line or by default.'),
 )
 
+compilation_cache_max_size = define_int_state(
+    name='jax_compilation_cache_max_size',
+    default=-1,
+    help=('The maximum size (in bytes) allowed for the persistent compilation '
+          'cache. When set, the least recently accessed cache entry(s) '
+          'will be deleted once the total cache directory size '
+          'exceeds the specified limit. '
+          'Caching will be disabled if this value is set to 0. A '
+          'special value of -1 indicates no limit, allowing the cache '
+          'size to grow indefinitely.'),
+)
+
 default_dtype_bits = define_enum_state(
     name='jax_default_dtype_bits',
     enum_values=['32', '64'],
diff --git a/jax/_src/gfile_cache.py b/jax/_src/gfile_cache.py
index 301f61cc6..989844b10 100644
--- a/jax/_src/gfile_cache.py
+++ b/jax/_src/gfile_cache.py
@@ -17,6 +17,10 @@ import os
 from jax._src import path as pathlib
 from jax._src.compilation_cache_interface import CacheInterface
 
+
+# TODO (ayx): This class will be ultimately removed after `lru_cache.py` is
+# finished. It exists because the current `lru_cache.py` does not support
+# `gs://`.
 class GFileCache(CacheInterface):
 
   def __init__(self, path: str):
diff --git a/jax/_src/lru_cache.py b/jax/_src/lru_cache.py
new file mode 100644
index 000000000..480d4915e
--- /dev/null
+++ b/jax/_src/lru_cache.py
@@ -0,0 +1,184 @@
+# Copyright 2024 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import heapq
+import logging
+import pathlib
+import warnings
+
+from jax._src.compilation_cache_interface import CacheInterface
+
+
+try:
+  import filelock
+except ImportError:
+  filelock = None
+
+
+logger = logging.getLogger(__name__)
+
+
+class LRUCache(CacheInterface):
+  """Bounded cache with least-recently-used (LRU) eviction policy.
+
+  This implementation includes cache reading, writing and eviction
+  based on the LRU policy.
+
+  Notably, when ``max_size`` is set to -1, the cache eviction
+  is disabled, and the LRU cache functions as a normal cache
+  without any size limitations.
+  """
+
+  def __init__(self, path: str, *, max_size: int, lock_timeout_secs: float | None = 10):
+    """Args:
+
+      path: The path to the cache directory.
+      max_size: The maximum size of the cache in bytes. Caching will be
+        disabled if this value is set to ``0``. A special value of ``-1``
+        indicates no limit, allowing the cache size to grow indefinitely.
+      lock_timeout_secs: (optional) The timeout for acquiring a file lock.
+    """
+    # TODO(ayx): add support for cloud other filesystems such as GCS
+    if not self._is_local_filesystem(path):
+      raise NotImplementedError("LRUCache only supports local filesystem at this time.")
+
+    self.path = pathlib.Path(path)
+    self.path.mkdir(parents=True, exist_ok=True)
+
+    # TODO(ayx): having a `self._path` is required by the base class
+    # `CacheInterface`, but the base class can be removed after `LRUCache`
+    # and the original `GFileCache` are unified
+    self._path = self.path
+
+    self.eviction_enabled = max_size != -1  # no eviction if `max_size` is set to -1
+
+    if self.eviction_enabled:
+      if filelock is None:
+        raise RuntimeError("Please install filelock package to set `jax_compilation_cache_max_size`")
+
+      self.max_size = max_size
+      self.lock_timeout_secs = lock_timeout_secs
+
+      lock_path = self.path / ".lockfile"
+      self.lock = filelock.FileLock(lock_path)
+
+  def get(self, key: str) -> bytes | None:
+    """Retrieves the cached value for the given key.
+
+    Args:
+      key: The key for which the cache value is retrieved.
+
+    Returns:
+      The cached data as bytes if available; ``None`` otherwise.
+    """
+    if not key:
+      raise ValueError("key cannot be empty")
+
+    file = self.path / key
+
+    if self.eviction_enabled:
+      self.lock.acquire(timeout=self.lock_timeout_secs)
+
+    try:
+      if not file.exists():
+        logger.debug(f"Cache miss for key: {key!r}")
+        return None
+
+      logger.debug(f"Cache hit for key: {key!r}")
+      file.touch()  # update mtime
+      return file.read_bytes()
+
+    finally:
+      if self.eviction_enabled:
+        self.lock.release()
+
+  def put(self, key: str, val: bytes) -> None:
+    """Adds a new entry to the cache.
+
+    If a cache item with the same key already exists, no action
+    will be taken, even if the value is different.
+
+    Args:
+      key: The key under which the data will be stored.
+      val: The data to be stored.
+    """
+    if not key:
+      raise ValueError("key cannot be empty")
+
+    # prevent adding entries that exceed the maximum size limit of the cache
+    if self.eviction_enabled and len(val) > self.max_size:
+      msg = (f"Cache value for key {key!r} of size {len(val)} bytes exceeds "
+             f"the maximum cache size of {self.max_size} bytes")
+      warnings.warn(msg)
+      return
+
+    file = self.path / key
+
+    if self.eviction_enabled:
+      self.lock.acquire(timeout=self.lock_timeout_secs)
+
+    try:
+      if file.exists():
+        return
+
+      self._evict_if_needed(additional_size=len(val))
+      file.write_bytes(val)
+
+    finally:
+      if self.eviction_enabled:
+        self.lock.release()
+
+  def _evict_if_needed(self, *, additional_size: int = 0) -> None:
+    """Evicts the least recently used items from the cache if necessary
+    to ensure the cache does not exceed its maximum size.
+
+    Args:
+      additional_size: The size of the new entry being added to the cache.
+        This is included to account for the new entry when checking if
+        eviction is needed.
+    """
+    if not self.eviction_enabled:
+      return
+
+    # a priority queue, each element is a tuple `(file_mtime, file, file_size)`
+    h: list[tuple[int, pathlib.Path, int]] = []
+    dir_size = 0
+    for file in self.path.iterdir():
+      if file.is_file():
+        file_size = file.stat().st_size
+        file_mtime = file.stat().st_mtime_ns
+
+        dir_size += file_size
+        heapq.heappush(h, (file_mtime, file, file_size))
+
+    target_size = self.max_size - additional_size
+    # evict files until the directory size is less than or equal
+    # to `target_size`
+    while dir_size > target_size:
+      file_mtime, file, file_size = heapq.heappop(h)
+      msg = (f"Evicting cache file {file.name}: file size {file_size} bytes, "
+             f"target cache size {target_size} bytes")
+      logger.debug(msg)
+      file.unlink()
+      dir_size -= file_size
+
+  # See comments in `jax.src.compilation_cache.get_file_cache()` for details.
+  # TODO(ayx): This function has a duplicate in that place, and there is
+  # redundancy here. However, this code is temporary, and once the issue
+  # is fixed, this code can be removed.
+  @staticmethod
+  def _is_local_filesystem(path: str) -> bool:
+    return path.startswith("file://") or "://" not in path
diff --git a/jax/_src/scipy/signal.py b/jax/_src/scipy/signal.py
index 73271aad5..ae093a3ac 100644
--- a/jax/_src/scipy/signal.py
+++ b/jax/_src/scipy/signal.py
@@ -284,37 +284,6 @@ def convolve2d(in1: Array, in2: Array, mode: str = 'full', boundary: str = 'fill
     - :func:`jax.numpy.convolve`: 1D convolution
     - :func:`jax.scipy.signal.convolve`: ND convolution
     - :func:`jax.scipy.signal.correlate`: ND correlation
-
-  Examples:
-    A few 2D convolution examples:
-
-    >>> x = jnp.array([[1, 2],
-    ...                [3, 4]])
-    >>> y = jnp.array([[2, 1, 1],
-    ...                [4, 3, 4],
-    ...                [1, 3, 2]])
-
-    Full 2D convolution uses implicit zero-padding at the edges:
-
-    >>> jax.scipy.signal.convolve2d(x, y, mode='full')
-    Array([[ 2.,  5.,  3.,  2.],
-           [10., 22., 17., 12.],
-           [13., 30., 32., 20.],
-           [ 3., 13., 18.,  8.]], dtype=float32)
-
-    Specifying ``mode = 'same'`` returns a centered 2D convolution of the same size
-    as the first input:
-
-    >>> jax.scipy.signal.convolve2d(x, y, mode='same')
-    Array([[22., 17.],
-           [30., 32.]], dtype=float32)
-
-    Specifying ``mode = 'valid'`` returns only the portion of 2D convolution
-    where the two arrays fully overlap:
-
-    >>> jax.scipy.signal.convolve2d(x, y, mode='valid')
-    Array([[22., 17.],
-           [30., 32.]], dtype=float32)
   """
   if boundary != 'fill' or fillvalue != 0:
     raise NotImplementedError("convolve2d() only supports boundary='fill', fillvalue=0")
@@ -356,29 +325,6 @@ def correlate(in1: Array, in2: Array, mode: str = 'full', method: str = 'auto',
     - :func:`jax.numpy.correlate`: 1D cross-correlation
     - :func:`jax.scipy.signal.correlate2d`: 2D cross-correlation
     - :func:`jax.scipy.signal.convolve`: ND convolution
-
-  Examples:
-    A few 1D correlation examples:
-
-    >>> x = jnp.array([1, 2, 3, 2, 1])
-    >>> y = jnp.array([1, 3, 2])
-
-    Full 1D correlation uses implicit zero-padding at the edges:
-
-    >>> jax.scipy.signal.correlate(x, y, mode='full')
-    Array([ 2.,  7., 13., 15., 11.,  5.,  1.], dtype=float32)
-
-    Specifying ``mode = 'same'`` returns a centered 1D correlation of the same
-    size as the first input:
-
-    >>> jax.scipy.signal.correlate(x, y, mode='same')
-    Array([ 7., 13., 15., 11.,  5.], dtype=float32)
-
-    Specifying ``mode = 'valid'`` returns only the portion of 1D correlation
-    where the two arrays fully overlap:
-
-    >>> jax.scipy.signal.correlate(x, y, mode='valid')
-    Array([13., 15., 11.], dtype=float32)
   """
   return convolve(in1, jnp.flip(in2.conj()), mode, precision=precision, method=method)
 
@@ -418,38 +364,6 @@ def correlate2d(in1: Array, in2: Array, mode: str = 'full', boundary: str = 'fil
     - :func:`jax.numpy.correlate`: 1D cross-correlation
     - :func:`jax.scipy.signal.correlate`: ND cross-correlation
     - :func:`jax.scipy.signal.convolve`: ND convolution
-
-  Examples:
-    A few 2D correlation examples:
-
-    >>> x = jnp.array([[2, 1, 3],
-    ...                [1, 3, 1],
-    ...                [4, 1, 2]])
-    >>> y = jnp.array([[1, 3],
-    ...                [4, 2]])
-
-    Full 2D correlation uses implicit zero-padding at the edges:
-
-    >>> jax.scipy.signal.correlate2d(x, y, mode='full')
-    Array([[ 4., 10., 10., 12.],
-           [ 8., 15., 24.,  7.],
-           [11., 28., 14.,  9.],
-           [12.,  7.,  7.,  2.]], dtype=float32)
-
-    Specifying ``mode = 'same'`` returns a centered 2D correlation of the same
-    size as the first input:
-
-    >>> jax.scipy.signal.correlate2d(x, y, mode='same')
-    Array([[15., 24.,  7.],
-           [28., 14.,  9.],
-           [ 7.,  7.,  2.]], dtype=float32)
-
-    Specifying ``mode = 'valid'`` returns only the portion of 2D correlation
-    where the two arrays fully overlap:
-
-    >>> jax.scipy.signal.correlate2d(x, y, mode='valid')
-    Array([[15., 24.],
-           [28., 14.]], dtype=float32)
   """
   if boundary != 'fill' or fillvalue != 0:
     raise NotImplementedError("correlate2d() only supports boundary='fill', fillvalue=0")
diff --git a/pyproject.toml b/pyproject.toml
index 630f25835..8148308aa 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -12,6 +12,7 @@ warn_unused_ignores = true
 module = [
     "absl.*",
     "colorama.*",
+    "filelock.*",
     "importlib_metadata.*",
     "IPython.*",
     "numpy.*",


Test patch:
diff --git a/tests/BUILD b/tests/BUILD
index 69d6aff38..21fe476e3 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -1135,6 +1135,16 @@ py_test(
     ],
 )
 
+py_test(
+    name = "lru_cache_test",
+    srcs = ["lru_cache_test.py"],
+    deps = [
+        "//jax",
+        "//jax:lru_cache",
+        "//jax:test_util",
+    ] + py_deps("filelock"),
+)
+
 jax_test(
     name = "compilation_cache_test",
     srcs = ["compilation_cache_test.py"],
diff --git a/tests/lru_cache_test.py b/tests/lru_cache_test.py
new file mode 100644
index 000000000..483f65975
--- /dev/null
+++ b/tests/lru_cache_test.py
@@ -0,0 +1,151 @@
+# Copyright 2024 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import tempfile
+import time
+
+from absl.testing import absltest
+
+from jax._src import path as pathlib
+from jax._src.lru_cache import LRUCache
+import jax._src.test_util as jtu
+
+
+class LRUCacheTestCase(jtu.JaxTestCase):
+  name: str | None
+  path: pathlib.Path | None
+
+  def setUp(self):
+    super().setUp()
+    tmpdir = tempfile.TemporaryDirectory()
+    self.enter_context(tmpdir)
+    self.name = tmpdir.name
+    self.path = pathlib.Path(self.name)
+
+  def tearDown(self):
+    self.path = None
+    self.name = None
+    super().tearDown()
+
+
+class LRUCacheTest(LRUCacheTestCase):
+
+  def test_get_nonexistent_key(self):
+    cache = LRUCache(self.name, max_size=-1)
+    self.assertIsNone(cache.get("cache-a"))
+
+  def test_put_and_get_key(self):
+    cache = LRUCache(self.name, max_size=-1)
+
+    cache.put("cache-a", b"a")
+    self.assertEqual(cache.get("cache-a"), b"a")
+    self.assertEqual(set(self.path.glob("cache-*")), {self.path / "cache-a"})
+
+    cache.put("cache-b", b"b")
+    self.assertEqual(cache.get("cache-a"), b"a")
+    self.assertEqual(cache.get("cache-b"), b"b")
+    self.assertEqual(set(self.path.glob("cache-*")), {self.path / "cache-a", self.path / "cache-b"})
+
+  def test_put_empty_value(self):
+    cache = LRUCache(self.name, max_size=-1)
+
+    cache.put("cache-a", b"")
+    self.assertEqual(cache.get("cache-a"), b"")
+
+  def test_put_empty_key(self):
+    cache = LRUCache(self.name, max_size=-1)
+
+    with self.assertRaisesRegex(ValueError, r"key cannot be empty"):
+      cache.put("", b"a")
+
+  def test_eviction(self):
+    cache = LRUCache(self.name, max_size=2)
+
+    cache.put("cache-a", b"a")
+    cache.put("cache-b", b"b")
+
+    # `sleep()` is necessary to guarantee that `cache-b`"s timestamp is strictly greater than `cache-a`"s
+    time.sleep(1)
+    cache.get("cache-b")
+
+    # write `cache-c`, evict `cache-a`
+    cache.put("cache-c", b"c")
+    self.assertEqual(set(self.path.glob("cache-*")), {self.path / "cache-b", self.path / "cache-c"})
+
+    # calling `get()` on `cache-b` makes `cache-c` least recently used
+    time.sleep(1)
+    cache.get("cache-b")
+
+    # write `cache-d`, evict `cache-c`
+    cache.put("cache-d", b"d")
+    self.assertEqual(set(self.path.glob("cache-*")), {self.path / "cache-b", self.path / "cache-d"})
+
+  def test_eviction_with_empty_value(self):
+    cache = LRUCache(self.name, max_size=1)
+
+    cache.put("cache-a", b"a")
+
+    # write `cache-b` with length 0
+    # eviction should not happen even though the cache is full
+    cache.put("cache-b", b"")
+    self.assertEqual(set(self.path.glob("cache-*")), {self.path / "cache-a", self.path / "cache-b"})
+
+    # calling `get()` on `cache-a` makes `cache-b` least recently used
+    time.sleep(1)
+    cache.get("cache-a")
+
+    # writing `cache-c` should result in evicting the
+    # least recent used file (`cache-b`) first,
+    # but this is not sufficient to make room for `cache-c`,
+    # so `cache-a` should be evicted as well
+    cache.put("cache-c", b"c")
+    self.assertEqual(set(self.path.glob("cache-*")), {self.path / "cache-c"})
+
+  def test_existing_cache_dir(self):
+    cache = LRUCache(self.name, max_size=2)
+
+    cache.put("cache-a", b"a")
+
+    # simulates reinitializing the cache in another process
+    del cache
+    cache = LRUCache(self.name, max_size=2)
+
+    self.assertEqual(cache.get("cache-a"), b"a")
+
+    # ensure that the LRU policy survives cache reinitialization
+    cache.put("cache-b", b"b")
+
+    # calling `get()` on `cache-a` makes `cache-b` least recently used
+    time.sleep(1)
+    cache.get("cache-a")
+
+    # write `cache-c`, evict `cache-b`
+    cache.put("cache-c", b"c")
+    self.assertEqual(set(self.path.glob("cache-*")), {self.path / "cache-a", self.path / "cache-c"})
+
+  def test_max_size(self):
+    cache = LRUCache(self.name, max_size=1)
+
+    msg = (r"Cache value for key .+? of size \d+ bytes exceeds the maximum "
+           r"cache size of \d+ bytes")
+    with self.assertWarnsRegex(UserWarning, msg):
+      cache.put("cache-a", b"aaaa")
+    self.assertIsNone(cache.get("cache-a"))
+    self.assertEqual(set(self.path.glob("cache-*")), set())
+
+
+if __name__ == "__main__":
+  absltest.main(testLoader=jtu.JaxTestLoader())