Disable typeguard externally

PiperOrigin-RevId: 698706356

.github/workflows/pytest_and_autopublish.yml

@@ -40,6 +40,7 @@ jobs:
     # * sweep_utils_test: Depends on kxm
     # * lpips_test: Missing VGG weights
     # * partial_loader_test: Orbax partial checkpoint loader not yet open-sourced  (TODO(epot): Restore)
+    # * typing tests: Not yet supported due to typeguard version issues.
     - name: Run core tests
       run: |
         pytest -vv -n auto \
@@ -48,7 +49,9 @@ jobs:
           --ignore=kauldron/xm/ \
           --ignore=kauldron/metrics/lpips_test.py \
           --ignore=kauldron/checkpoints/partial_loader_test.py \
-          --ignore=kauldron/utils/sweep_utils_test.py
+          --ignore=kauldron/utils/sweep_utils_test.py \
+          --ignore=kauldron/typing/shape_spec_test.py \
+          --ignore=kauldron/typing/type_check_test.py
 
   # Auto-publish when version is increased
   publish-job:

kauldron/typing/__init__.py

@@ -16,11 +16,13 @@
 
 from __future__ import annotations
 
+# pylint: disable=g-multiple-import,g-importing-member
+
 from typing import Any, Callable, Hashable, Sequence, Union
 
-from clu.data.dataset_iterator import ArraySpec, ElementSpec, PyTree  # pylint: disable=g-multiple-import,g-importing-member
+from clu.data.dataset_iterator import ArraySpec, ElementSpec, PyTree
 import jax
-from kauldron.typing.array_types import (  # pylint: disable=g-multiple-import
+from kauldron.typing.array_types import (
     Array,
     Bool,
     Complex,
@@ -42,12 +44,11 @@
     UInt8,
     XArray,
 )
-from kauldron.typing.shape_spec import Dim, Memo, Shape  # pylint: disable=g-multiple-import,g-importing-member
-from kauldron.typing.type_check import TypeCheckError, typechecked  # pylint: disable=g-multiple-import,g-importing-member
+from kauldron.typing.shape_spec import Dim, Memo, Shape
+from kauldron.typing.type_check import check_type
+from kauldron.typing.type_check import typechecked
+from kauldron.typing.type_check import TypeCheckError
 import numpy as np
-import typeguard as _typeguard
-# make typeguard.check_type accessible in this namespace
-check_type = _typeguard.check_type
 
 PRNGKey = UInt32["2"]
 PRNGKeyLike = Union[int, Sequence[int], np.ndarray, PRNGKey]

kauldron/typing/shape_spec.py

@@ -83,8 +83,8 @@ def _assert_caller_is_typechecked_func() -> None:
   if stack[i + 1].function != "_reraise_with_shape_info":
     caller_name = stack[i].function
     raise AssertionError(
-        "Dim and Shape only work inside of @typechecked functions. But"
-        f" {caller_name!r} lacks @typechecked."
+        "Dim and Shape not yet supported due to `typeguard` issue."
+        f" Raised in {caller_name!r}"
     )
 
 

kauldron/typing/type_check.py

@@ -29,7 +29,9 @@
 from etils import epy
 import jaxtyping
 from kauldron.typing import shape_spec
-import typeguard
+
+with epy.lazy_imports():
+  import typeguard  # pylint: disable=g-import-not-at-top
 
 
 # a global switch to disable typechecking
@@ -39,7 +41,18 @@
 _undef = object()
 
 
-class TypeCheckError(typeguard.TypeCheckError):
+def check_type(
+    value: Any,
+    expected_type: Any,
+) -> None:
+  """Ensure that value matches expected_type, alias for typeguard.check_type."""
+  if True:  # Typeguard not yet supported
+    return
+  return typeguard.check_type(value, expected_type)
+
+
+exc_cls = Exception
+class TypeCheckError(exc_cls):
   """Indicates a runtime typechecking error from the @typechecked decorator."""
 
   def __init__(
@@ -99,6 +112,9 @@ def _annotation_repr(ann: Any) -> str:
 
 def typechecked(fn):
   """Decorator to enable runtime type-checking and shape-checking."""
+  if True:  # Typeguard not yet supported
+    return fn
+
   if hasattr(fn, "__wrapped__"):
     raise AssertionError("@typechecked should be the innermost decorator")
 
@@ -417,7 +433,3 @@ def add_custom_checker_lookup_fn(lookup_fn):
       break
   else:  # prepend
     checker_lookup_fns[:0] = [lookup_fn]
-
-
-add_custom_checker_lookup_fn(_array_spec_checker_lookup)
-add_custom_checker_lookup_fn(_dataclass_checker_lookup)

pyproject.toml

@@ -43,8 +43,9 @@ dependencies = [
     "tfds-nightly",  # TODO(klausg): switch back to tensorflow_datasets>=4.9.7
     # once released: https://github.com/tensorflow/datasets/commit/d4bfd59863c6cb5b64d043b7cb6ab566e7d92440
     "tqdm",
+    # TODO(klausg): Restore typeguard or switch to something else
     # closest match to the internal typeguard
-    "typeguard@git+https://github.com/agronholm/typeguard@0dd7f7510b7c694e66a0d17d1d58d185125bad5d",
+    # "typeguard@git+https://github.com/agronholm/typeguard@0dd7f7510b7c694e66a0d17d1d58d185125bad5d",
     "typing_extensions",
     "xmanager",
     # lazy deps (should ideally remove those)