migrate kauldron to ktyping

PiperOrigin-RevId: 822524116

Author: Qwlouse

kauldron/data/deprecated.py

@@ -37,7 +37,8 @@
 from kauldron.data import pipelines
 from kauldron.data import utils
 from kauldron.data.loaders import base as base_data_loader
-from kauldron.typing import PRNGKeyLike, PyTree  # pylint: disable=g-importing-member,g-multiple-import
+from kauldron.ktyping import PyTree  # pylint: disable=g-importing-member
+from kauldron.typing import PRNGKeyLike  # pylint: disable=g-importing-member
 from kauldron.utils import config_util
 import tensorflow as tf
 import tensorflow_datasets as tfds

kauldron/data/pipelines.py

@@ -26,7 +26,8 @@
 from kauldron.data import data_utils
 from kauldron.data import iterators
 from kauldron.data import utils
-from kauldron.typing import PRNGKeyLike, PyTree  # pylint: disable=g-importing-member,g-multiple-import
+from kauldron.ktyping import PyTree  # pylint: disable=g-importing-member
+from kauldron.typing import PRNGKeyLike  # pylint: disable=g-importing-member
 from kauldron.utils import config_util
 
 # Output of `tfds.as_numpy`

kauldron/evals/eval_impl.py

@@ -32,10 +32,10 @@
 from kauldron.checkpoints import partial_loader
 from kauldron.evals import evaluators as evaluators_lib
 from kauldron.evals import run_strategies
+from kauldron.ktyping import PyTree  # pylint: disable=g-importing-member
 from kauldron.train import auxiliaries
 from kauldron.train import train_step
 from kauldron.train import trainer_lib
-from kauldron.typing import PyTree  # pylint: disable=g-importing-member
 from kauldron.utils import constants
 from kauldron.utils.status_utils import status  # pylint: disable=g-importing-member
 import orbax.checkpoint as ocp

kauldron/evals/fewshot_evaluator.py

@@ -27,11 +27,11 @@
 from kauldron import data
 from kauldron import kontext
 from kauldron.evals import evaluators
+from kauldron.ktyping import Array, Float, Int, Scalar, check_type, typechecked  # pylint: disable=g-multiple-import,g-importing-member
 from kauldron.metrics import base
 from kauldron.metrics import base_state
 from kauldron.train import auxiliaries
 from kauldron.train import train_step
-from kauldron.typing import Array, Float, Int, Scalar, check_type, typechecked  # pylint: disable=g-multiple-import,g-importing-member
 from kauldron.utils import config_util
 from kauldron.utils import kdash
 from kauldron.utils import utils

kauldron/kd.py

@@ -42,6 +42,7 @@
   from kauldron import train
   from kauldron import testing
   from kauldron import typing
+  from kauldron import ktyping
   from kauldron.utils import api as utils
   from kauldron.utils import from_xid
   from kauldron.utils import kdash

kauldron/losses/base.py

@@ -25,8 +25,8 @@
 from jax import numpy as jnp
 from kauldron import kontext
 from kauldron import metrics
+from kauldron.ktyping import Array, Float, PyTree  # pylint: disable=g-multiple-import,g-importing-member
 from kauldron.metrics import base_state
-from kauldron.typing import Array, Float, PyTree  # pylint: disable=g-multiple-import,g-importing-member
 
 
 Schedule = Callable[[int], float]

kauldron/losses/simple.py

@@ -20,8 +20,8 @@
 
 import jax.numpy as jnp
 from kauldron import kontext
+from kauldron.ktyping import Array, Float, Int, typechecked  # pylint: disable=g-multiple-import,g-importing-member
 from kauldron.losses import base
-from kauldron.typing import Array, Float, Int, typechecked  # pylint: disable=g-multiple-import,g-importing-member
 import optax
 
 
@@ -80,6 +80,7 @@ def get_values(self, preds: Float["*a"], targets: Float["*a"]) -> Float["*a"]:
 @dataclasses.dataclass(eq=True, frozen=True, kw_only=True)
 class Huber(base.Loss):
   """Huber loss."""
+
   delta: float = 1.0
 
   preds: kontext.Key = kontext.REQUIRED
@@ -113,7 +114,8 @@ def get_values(
     targets = self._safe_normalize(targets)
 
     similarity = jnp.sum(preds * targets, axis=-1, keepdims=True)
-    return - similarity
+    return -similarity
+
 
 # ============================== Classification ===============================
 

kauldron/main.py

@@ -67,7 +67,6 @@
 @_catch_post_mortem
 def main(_):
   tf.config.set_visible_devices([], "GPU")
-  kd.typing.enable_kd_type_checking()  # Enable custom checks before resolve
   eval_names = _EVAL_NAMES.value
   cfg = _CONFIG.value
   trainer: kd.train.Trainer = kd.konfig.resolve(cfg)

kauldron/metrics/auto_state.py

@@ -26,9 +26,9 @@
 import jax
 import jax.numpy as jnp
 from kauldron import kontext
+from kauldron.ktyping import Array, PyTree  # pylint: disable=g-multiple-import,g-importing-member
 from kauldron.metrics import base_state
 from kauldron.metrics.base_state import EMPTY  # pylint: disable=g-importing-member
-from kauldron.typing import Array, PyTree  # pylint: disable=g-multiple-import,g-importing-member
 import numpy as np
 
 _MetricT = TypeVar("_MetricT")

kauldron/metrics/base.py

@@ -24,8 +24,8 @@
 import flax
 import jax
 from kauldron import kontext
+from kauldron.ktyping import Float, PyTree  # pylint: disable=g-multiple-import,g-importing-member
 from kauldron.metrics import base_state
-from kauldron.typing import Float, PyTree  # pylint: disable=g-multiple-import,g-importing-member
 
 _FnT = TypeVar("_FnT")