Add Keras trainer.

PiperOrigin-RevId: 744903627

Author: Hilly12

mlrx/__init__.py

@@ -27,10 +27,12 @@
 from mlrx.training.core import Experiment
 from mlrx.training.core import run_experiment
 from mlrx.training.core import Trainer
-from mlrx.training.jax import JaxState
-from mlrx.training.jax import JaxTask
-from mlrx.training.jax import JaxTrainer
-from mlrx.training.jax import KerasState
+from mlrx.training.jax_trainer import JaxState
+from mlrx.training.jax_trainer import JaxTask
+from mlrx.training.jax_trainer import JaxTrainer
+from mlrx.training.jax_trainer import KerasState
+from mlrx.training.keras_trainer import KerasTask
+from mlrx.training.keras_trainer import KerasTrainer
 from mlrx.training.optax_factory import AdagradFactory
 from mlrx.training.optax_factory import AdamFactory
 from mlrx.training.optax_factory import OptimizerFactory

mlrx/training/core.py

@@ -14,7 +14,7 @@
 """Core training library for Jax."""
 
 import abc
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 import dataclasses
 import enum
 from typing import Any, Generic, TypeVar
@@ -33,6 +33,8 @@
 TRAINING_COMPLETE_MARKER_FILE = "marker.txt"
 TRAIN_LOG_DIRNAME = "train"
 EVAL_LOG_DIRNAME = "val"
+KERAS_MODEL_SAVEFILE = "model.keras"
+ORBAX_CHECKPOINT_DEFAULT_KEY = "default"
 
 DEFAULT_RNG_SEED = 0
 IN_TRAINER_CONTEXT = False  # Set to true when run from the main trainer.
@@ -171,6 +173,15 @@ def get_iterators(
   return train_dataset, eval_datasets  # pytype: disable=bad-return-type
 
 
+def get_shape(
+    x: tf.Tensor | tf.SparseTensor | tf.RaggedTensor,
+) -> Sequence[int | None]:
+  """Gets the shape of a dense / sparse / ragged tensor."""
+  if isinstance(x, tf.SparseTensor):
+    return [x.shape[0]] + [None for _ in x.shape[1:]]
+  return x.shape.as_list()
+
+
 def in_tracing_context() -> bool:
   """Returns whether the current context is a tracing context."""
   return isinstance(jnp.ones(()), jax.core.Tracer)

mlrx/training/jax.py renamed to mlrx/training/jax_trainer.py

@@ -25,13 +25,13 @@
 import jax.numpy as jnp
 import jaxtyping as jt
 import optax
-from mlrx.training import jax as jax_lib
+from mlrx.training import jax_trainer
 from mlrx.training import partitioning
 import tensorflow as tf
 import tensorflow_datasets as tfds
 
 
-class _MNISTTask(jax_lib.JaxTask):
+class _MNISTTask(jax_trainer.JaxTask):
   """Task for fitting a CNN on MNIST."""
 
   def create_datasets(self) -> tuple[tf.data.Dataset, tf.data.Dataset]:
@@ -126,7 +126,7 @@ def setUp(self):
   def test_mnist_e2e(self):
     model_dir = self.create_tempdir().full_path
     task = _MNISTTask()
-    trainer = jax_lib.JaxTrainer(
+    trainer = jax_trainer.JaxTrainer(
         partitioner=partitioning.DataParallelPartitioner(),
         train_steps=1000,
         steps_per_eval=50,

mlrx/training/jax_quality_test.py renamed to mlrx/training/jax_trainer_quality_test.py

@@ -30,7 +30,7 @@
 import optax
 import orbax.checkpoint as ocp
 from mlrx.training import core
-from mlrx.training import jax as jax_lib
+from mlrx.training import jax_trainer
 from mlrx.training import partitioning
 import tensorflow as tf
 
@@ -42,7 +42,7 @@ def __call__(self, inputs: jax.Array) -> jax.Array:
     return nn.Dense(1, kernel_init=nn.initializers.constant(-1.0))(inputs)
 
 
-class _JaxTask(jax_lib.JaxTask):
+class _JaxTask(jax_trainer.JaxTask):
 
   def create_datasets(
       self,
@@ -90,7 +90,7 @@ def eval_step(
     return {"loss": clu_metrics.Average.from_model_output(loss)}
 
 
-class _KerasJaxTask(jax_lib.JaxTask):
+class _KerasJaxTask(jax_trainer.JaxTask):
 
   def create_datasets(self) -> tf.data.Dataset:
     def _map_fn(x: int):
@@ -106,7 +106,7 @@ def _map_fn(x: int):
 
   def create_state(
       self, batch: jt.PyTree, rng: jax.Array
-  ) -> jax_lib.KerasState:
+  ) -> jax_trainer.KerasState:
     x, _ = batch
 
     model = keras.Sequential(
@@ -122,11 +122,11 @@ def create_state(
     model.build(x.shape)
 
     optimizer = optax.adagrad(0.1)
-    return jax_lib.KerasState.create(model=model, tx=optimizer)
+    return jax_trainer.KerasState.create(model=model, tx=optimizer)
 
   def train_step(
-      self, batch: jt.PyTree, state: jax_lib.KerasState, rng: jax.Array
-  ) -> tuple[jax_lib.KerasState, Mapping[str, clu_metrics.Metric]]:
+      self, batch: jt.PyTree, state: jax_trainer.KerasState, rng: jax.Array
+  ) -> tuple[jax_trainer.KerasState, Mapping[str, clu_metrics.Metric]]:
     x, y = batch
 
     def _loss_fn(tvars):
@@ -140,7 +140,7 @@ def _loss_fn(tvars):
     return state, {"loss": clu_metrics.Average.from_model_output(loss)}
 
   def eval_step(
-      self, batch: jt.PyTree, state: jax_lib.KerasState
+      self, batch: jt.PyTree, state: jax_trainer.KerasState
   ) -> Mapping[str, clu_metrics.Metric]:
     x, y = batch
     y_pred, _ = state.model.stateless_call(state.tvars, state.ntvars, x)
@@ -208,13 +208,13 @@ def setUp(self):
   )
   def test_jax_trainer(
       self,
-      task_cls: type[jax_lib.JaxTask],
+      task_cls: type[jax_trainer.JaxTask],
       mode: str,
       expected_keys: Sequence[str],
   ):
     model_dir = self.create_tempdir().full_path
     task = task_cls()
-    trainer = jax_lib.JaxTrainer(
+    trainer = jax_trainer.JaxTrainer(
         partitioner=partitioning.DataParallelPartitioner(data_axis="batch"),
         train_steps=12,
         steps_per_eval=3,
@@ -258,7 +258,7 @@ class State:
         ),
     )
     state = State(step=10, opt_state=tx.init({"a": jnp.ones((10, 10))}))
-    metrics = jax_lib._state_metrics(state)
+    metrics = jax_trainer._state_metrics(state)
     self.assertIn("optimizer/learning_rate", metrics)
     self.assertEqual(metrics["optimizer/learning_rate"].compute(), 0.1)