diff --git a/recml/core/training/batch_folding.py b/recml/core/training/batch_folding.py
new file mode 100644
index 0000000..ae4115c
--- /dev/null
+++ b/recml/core/training/batch_folding.py
@@ -0,0 +1,239 @@
+# Copyright 2024 RecML authors <recommendations-ml@google.com>.
+#
+# 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
+#
+#     http://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.
+"""Folds the dense leaves of training batches into a few arrays.
+
+The Keras JAX training loop converts and copies every leaf of a batch to the
+device one by one, in Python. With many leaves per batch, as in packed sequence
+training, this holds the GIL for tens of milliseconds per step, and the thread
+that starts the training steps needs the GIL too.
+
+`fold_dataset` concatenates the dense leaves of each batch in `tf.data`, one
+array per dtype and leading dimension, so that only a few arrays cross to the
+device. `model_unfold_scope` makes the jitted train step of a model split them
+back, so the model sees the original batch.
+
+Terms:
+  leaf: one tensor of a batch, in the order of `keras.tree.flatten`.
+  leaf group: the leaves that fold into one array. They share dtype and leading
+    dimension, and there are 2 or more of them.
+  folded array: the concatenation of a leaf group, of shape
+    `[leading_dim, total_width]`, where the width of a leaf is the product of
+    its dimensions after the first.
+"""
+
+import collections
+from collections.abc import Callable, Iterator, Sequence
+import contextlib
+import dataclasses
+import functools
+import math
+from typing import Any
+
+from absl import logging
+import keras
+import tensorflow as tf
+
+
+@dataclasses.dataclass(frozen=True)
+class _LeafGroups:
+  """Which leaves of a batch fold into which array, and how to undo it.
+
+  Attributes:
+    structure: The element structure of the dataset before folding.
+    specs: The `tf.TensorSpec` (or other spec) of each leaf, in flatten order.
+    folded_leaf_groups: For each folded array, the indices of its leaves in the
+      order of their columns.
+    passthrough_leaf_indices: The indices of the leaves that pass through.
+  """
+
+  structure: Any
+  specs: Sequence[Any]
+  folded_leaf_groups: Sequence[Sequence[int]]
+  passthrough_leaf_indices: Sequence[int]
+
+  def leaf_width(self, leaf_index: int) -> int:
+    """Returns the number of columns of a leaf in its folded array."""
+    return math.prod(self.specs[leaf_index].shape[1:])
+
+  # All control flow here is over static Python specs, not tensors.
+  @tf.autograph.experimental.do_not_convert
+  def fold(self, *element: Any) -> tuple[tuple[Any, ...], tuple[Any, ...]]:
+    """Returns `(folded_arrays, passthrough_leaves)` of a dataset element."""
+    # `Dataset.map` passes the components of a tuple element as arguments.
+    if not isinstance(self.structure, tuple):
+      [element] = element
+    leaves = keras.tree.flatten(element)
+    folded_arrays = []
+    for leaf_group in self.folded_leaf_groups:
+      columns = [  # Each entry has shape [leading_dim, leaf_width].
+          tf.reshape(leaves[i], [-1, self.leaf_width(i)]) for i in leaf_group
+      ]
+      folded_arrays.append(tf.concat(columns, axis=1))
+    passthrough_leaves = [leaves[i] for i in self.passthrough_leaf_indices]
+    return tuple(folded_arrays), tuple(passthrough_leaves)
+
+  def unfold(self, element: Any) -> Any:
+    """Returns the original element of a `(folded_arrays, ...)` element.
+
+    The leading dimension of the arrays may differ from the one in `specs`, as
+    when the arrays hold the batches of several hosts.
+
+    Args:
+      element: `(folded_arrays, passthrough_leaves)` as returned by `fold`.
+
+    Returns:
+      The element in `structure`.
+    """
+    folded_arrays, passthrough_leaves = element
+    leaves = [None] * len(self.specs)
+    for folded_array, leaf_group in zip(
+        folded_arrays, self.folded_leaf_groups, strict=True
+    ):
+      leading_dim = folded_array.shape[0]
+      first_column = 0
+      for i in leaf_group:
+        last_column = first_column + self.leaf_width(i)
+        columns = folded_array[:, first_column:last_column]
+        leaf_shape = (leading_dim, *self.specs[i].shape[1:])
+        leaves[i] = keras.ops.reshape(columns, leaf_shape)
+        first_column = last_column
+    for i, leaf in zip(
+        self.passthrough_leaf_indices, passthrough_leaves, strict=True
+    ):
+      leaves[i] = leaf
+    return keras.tree.pack_sequence_as(self.structure, leaves)
+
+
+def fold_dataset(
+    dataset: tf.data.Dataset,
+) -> tuple[tf.data.Dataset, Callable[[Any], Any] | None]:
+  """Folds the dense leaves of the elements of `dataset`.
+
+  Leaves of rank 1 or more whose dimensions after the first are all known are
+  grouped by dtype and static leading dimension. Each group of 2 or more leaves
+  becomes one array of shape `[leading_dim, total_width]`. All other leaves pass
+  through unchanged. Leaves with an unknown leading dimension form their own
+  group; `tf.concat` fails if their leading dimensions differ at run time.
+
+  For example, a batch of 2 rows (keys flatten in sorted order)
+
+    {"ids": int32 [[1, 2], [3, 4]],            # shape [2, 2], width 2
+     "mask": int32 [[[1, 0]], [[0, 1]]],       # shape [2, 1, 2], width 2
+     "present": bool [True, False],            # only bool leaf
+     "weight": float32 [0.5, 2.0]}             # only float32 leaf
+
+  folds into one int32 array, with "ids" in columns 0-1 and "mask" in columns
+  2-3, while "present" and "weight" have no leaf to fold with:
+
+    ((int32 [[1, 2, 1, 0], [3, 4, 0, 1]],), (present, weight))
+
+  Args:
+    dataset: A batched dataset.
+
+  Returns:
+    `(folded_dataset, unfold)`. The elements of `folded_dataset` are
+    `(folded_arrays, passthrough_leaves)` tuples. `unfold` maps such an element
+    back to the structure of the elements of `dataset`; it accepts arrays with
+    a different leading dimension, as on multiple hosts, and runs under `jit`.
+    If no leaves can be folded, returns `(dataset, None)`.
+
+  Raises:
+    TypeError: If `dataset` is not a `tf.data.Dataset`.
+  """
+  if not isinstance(dataset, tf.data.Dataset):
+    raise TypeError(
+        f"Folding batches needs a `tf.data.Dataset`, got {type(dataset)}."
+    )
+  structure = dataset.element_spec
+  specs = keras.tree.flatten(structure)
+  leaf_indices_by_dtype_and_leading_dim = collections.defaultdict(list)
+  for i, spec in enumerate(specs):
+    if (
+        isinstance(spec, tf.TensorSpec)
+        and spec.shape.rank
+        and spec.shape[1:].is_fully_defined()
+    ):
+      key = (spec.dtype, spec.shape[0])
+      leaf_indices_by_dtype_and_leading_dim[key].append(i)
+  folded_leaf_groups = [
+      leaf_indices
+      for leaf_indices in leaf_indices_by_dtype_and_leading_dim.values()
+      if len(leaf_indices) > 1
+  ]
+  if not folded_leaf_groups:
+    return dataset, None
+  folded_leaf_indices = set().union(*folded_leaf_groups)
+  leaf_groups = _LeafGroups(
+      structure=structure,
+      specs=specs,
+      folded_leaf_groups=folded_leaf_groups,
+      passthrough_leaf_indices=[
+          i for i in range(len(specs)) if i not in folded_leaf_indices
+      ],
+  )
+  logging.info(
+      "Folding %d of %d batch leaves into %d arrays.",
+      len(folded_leaf_indices),
+      len(specs),
+      len(folded_leaf_groups),
+  )
+  return dataset.map(leaf_groups.fold), leaf_groups.unfold
+
+
+@contextlib.contextmanager
+def model_unfold_scope(
+    model: keras.Model, unfold: Callable[[Any], Any] | None
+) -> Iterator[None]:
+  """Makes the train step of `model` unfold its batches.
+
+  `model.fit` must run on the folded dataset inside this context. The model,
+  its compiled loss and metrics, and its optimizer must be built (for example
+  with `model._symbolic_build(data_batch=unfolded_dataset.element_spec)`):
+  Keras builds unbuilt parts from the first batch, which is folded.
+
+  Args:
+    model: A built Keras model on the JAX backend.
+    unfold: The function returned by `fold_dataset`. `None` leaves `model`
+      untouched.
+
+  Yields:
+    None.
+
+  Raises:
+    ValueError: If `unfold` is set and `model` is not built.
+  """
+  if unfold is None:
+    yield
+    return
+  if not model.built:
+    raise ValueError("The model must be built to train on folded batches.")
+  had_instance_train_step = "train_step" in vars(model)
+  train_step = model.train_step
+
+  # A 2-line closure over `train_step` and `unfold` reads more simply than a
+  # class with the same two fields.
+  @functools.wraps(train_step)
+  def unfolded_train_step(state, data):
+    return train_step(state, unfold(data))
+
+  model.train_step = unfolded_train_step
+  model.train_function = None
+  try:
+    yield
+  finally:
+    if had_instance_train_step:
+      model.train_step = train_step
+    else:
+      del model.train_step
+    model.train_function = None
diff --git a/recml/core/training/batch_folding_test.py b/recml/core/training/batch_folding_test.py
new file mode 100644
index 0000000..f9e4c56
--- /dev/null
+++ b/recml/core/training/batch_folding_test.py
@@ -0,0 +1,306 @@
+# Copyright 2024 RecML authors <recommendations-ml@google.com>.
+#
+# 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
+#
+#     http://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 absl.testing import absltest
+import keras
+import numpy as np
+from recml.core.training import batch_folding
+import tensorflow as tf
+
+_BATCH_SIZE = 4
+_NUM_BATCHES = 3
+
+
+def _features() -> dict[str, np.ndarray]:
+  rng = np.random.default_rng(0)
+  n = _BATCH_SIZE * _NUM_BATCHES
+  return {
+      "ids": rng.integers(0, 10, [n, 3]).astype(np.int32),
+      "segments": rng.integers(0, 10, [n, 2, 2]).astype(np.int32),
+      "values": rng.random([n, 5]).astype(np.float32),
+      "weight": rng.random([n]).astype(np.float32),
+      # The only bool leaf: it has no group to fold with.
+      "present": rng.random([n]) > 0.5,
+  }
+
+
+def _labels() -> np.ndarray:
+  n = _BATCH_SIZE * _NUM_BATCHES
+  return np.random.default_rng(1).random([n, 1]).astype(np.float32)
+
+
+def _dataset() -> tf.data.Dataset:
+  return tf.data.Dataset.from_tensor_slices((_features(), _labels())).batch(
+      _BATCH_SIZE, drop_remainder=True
+  )
+
+
+def _assert_trees_equal(actual, expected):
+  keras.tree.assert_same_structure(actual, expected)
+  for a, e in zip(keras.tree.flatten(actual), keras.tree.flatten(expected)):
+    a, e = np.asarray(a), np.asarray(e)
+    np.testing.assert_array_equal(a, e)
+    np.testing.assert_equal(a.dtype, e.dtype)
+
+
+class _Model(keras.Model):
+  """Regression on the features; its loss needs every leaf."""
+
+  def __init__(self):
+    super().__init__()
+    self._dense = keras.layers.Dense(1)
+
+  def call(self, inputs):
+    ids = keras.ops.cast(inputs["ids"], "float32")
+    segments = keras.ops.cast(inputs["segments"], "float32")
+    x = keras.ops.concatenate(
+        [ids, keras.ops.reshape(segments, (-1, 4)), inputs["values"]], axis=1
+    )
+    present = keras.ops.cast(inputs["present"], "float32")
+    return self._dense(x) * (inputs["weight"] * present)[:, None]
+
+  def compute_loss(
+      self, x=None, y=None, y_pred=None, sample_weight=None, training=True
+  ):
+    del x, sample_weight, training
+    return keras.ops.mean(keras.ops.square(y_pred - y))
+
+
+def _built_model() -> _Model:
+  keras.utils.set_random_seed(0)
+  model = _Model()
+  model.compile(optimizer=keras.optimizers.SGD(0.1))
+  model(next(iter(_dataset()))[0])
+  return model
+
+
+class FoldDatasetTest(absltest.TestCase):
+
+  def test_docstring_example(self):
+    dataset = tf.data.Dataset.from_tensors({
+        "ids": np.array([[1, 2], [3, 4]], np.int32),
+        "mask": np.array([[[1, 0]], [[0, 1]]], np.int32),
+        "present": np.array([True, False]),
+        "weight": np.array([0.5, 2.0], np.float32),
+    })
+
+    folded_dataset, unfold = batch_folding.fold_dataset(dataset)
+
+    folded = next(folded_dataset.as_numpy_iterator())
+    expected_folded = (
+        (np.array([[1, 2, 1, 0], [3, 4, 0, 1]], np.int32),),
+        (np.array([True, False]), np.array([0.5, 2.0], np.float32)),
+    )
+    _assert_trees_equal(folded, expected_folded)
+    _assert_trees_equal(unfold(folded), next(dataset.as_numpy_iterator()))
+
+  def test_folds_leaves_by_dtype_and_leading_dimension(self):
+    folded_dataset, _ = batch_folding.fold_dataset(_dataset())
+
+    folded_arrays, other_leaves = folded_dataset.element_spec
+    # int32: ids (3) + segments (4). float32: values (5) + weight (1) +
+    # labels (1).
+    self.assertCountEqual(
+        [(spec.dtype, tuple(spec.shape)) for spec in folded_arrays],
+        [(tf.int32, (_BATCH_SIZE, 7)), (tf.float32, (_BATCH_SIZE, 7))],
+    )
+    self.assertEqual(
+        [(spec.dtype, tuple(spec.shape)) for spec in other_leaves],
+        [(tf.bool, (_BATCH_SIZE,))],
+    )
+
+  def test_unfold_restores_the_elements(self):
+    dataset = _dataset()
+    folded_dataset, unfold = batch_folding.fold_dataset(dataset)
+
+    num_elements = 0
+    for element, folded in zip(
+        dataset.as_numpy_iterator(), folded_dataset.as_numpy_iterator()
+    ):
+      _assert_trees_equal(unfold(folded), element)
+      num_elements += 1
+    self.assertEqual(num_elements, _NUM_BATCHES)
+
+  def test_unfold_accepts_another_leading_dimension(self):
+    dataset = _dataset()
+    folded_dataset, unfold = batch_folding.fold_dataset(dataset)
+    elements = list(dataset.as_numpy_iterator())
+    folded = list(folded_dataset.as_numpy_iterator())
+
+    # On several hosts, the arrays of a step hold the batches of all hosts.
+    joined = keras.tree.map_structure(
+        lambda *leaves: np.concatenate(leaves, axis=0), *folded
+    )
+
+    _assert_trees_equal(
+        unfold(joined),
+        keras.tree.map_structure(
+            lambda *leaves: np.concatenate(leaves, axis=0), *elements
+        ),
+    )
+
+  def test_folds_a_dict_element(self):
+    dataset = tf.data.Dataset.from_tensor_slices(_features()).batch(
+        _BATCH_SIZE, drop_remainder=True
+    )
+    folded_dataset, unfold = batch_folding.fold_dataset(dataset)
+
+    for element, folded in zip(
+        dataset.as_numpy_iterator(),
+        folded_dataset.as_numpy_iterator(),
+        strict=True,
+    ):
+      _assert_trees_equal(unfold(folded), element)
+
+  def test_ragged_leaves_pass_through(self):
+    features = _features()
+    dataset = tf.data.Dataset.from_tensor_slices({
+        "ids": features["ids"],
+        "segments": features["segments"],
+        "ragged": tf.ragged.constant(
+            [[1, 2], [3]] * (_BATCH_SIZE * _NUM_BATCHES // 2), dtype=tf.int32
+        ),
+    }).batch(_BATCH_SIZE, drop_remainder=True)
+
+    folded_dataset, unfold = batch_folding.fold_dataset(dataset)
+
+    for element, folded in zip(dataset, folded_dataset, strict=True):
+      folded_arrays, other_leaves = folded
+      self.assertLen(folded_arrays, 1)
+      [ragged] = other_leaves
+      restored = unfold(folded)
+      self.assertIs(restored["ragged"], ragged)
+      self.assertEqual(ragged.to_list(), element["ragged"].to_list())
+      np.testing.assert_array_equal(restored["ids"], element["ids"])
+      np.testing.assert_array_equal(restored["segments"], element["segments"])
+
+  def test_one_leaf_per_dtype_is_not_folded(self):
+    features = _features()
+    dataset = tf.data.Dataset.from_tensor_slices(
+        {"ids": features["ids"], "values": features["values"]}
+    ).batch(_BATCH_SIZE, drop_remainder=True)
+
+    folded_dataset, unfold = batch_folding.fold_dataset(dataset)
+
+    self.assertIs(folded_dataset, dataset)
+    self.assertIsNone(unfold)
+
+  def test_folds_leaves_with_an_unknown_leading_dimension(self):
+    # 10 rows batched by 4 without `drop_remainder` yield sizes [4, 4, 2].
+    features = {k: v[:10] for k, v in _features().items()}
+    labels = _labels()[:10]
+    dataset = tf.data.Dataset.from_tensor_slices((features, labels)).batch(
+        _BATCH_SIZE
+    )
+
+    folded_dataset, unfold = batch_folding.fold_dataset(dataset)
+
+    folded_arrays, other_leaves = folded_dataset.element_spec
+    self.assertCountEqual(
+        [(spec.dtype, tuple(spec.shape)) for spec in folded_arrays],
+        [(tf.int32, (None, 7)), (tf.float32, (None, 7))],
+    )
+    self.assertEqual(
+        [(spec.dtype, tuple(spec.shape)) for spec in other_leaves],
+        [(tf.bool, (None,))],
+    )
+    num_elements = 0
+    for element, folded in zip(
+        dataset.as_numpy_iterator(), folded_dataset.as_numpy_iterator()
+    ):
+      _assert_trees_equal(unfold(folded), element)
+      num_elements += 1
+    self.assertEqual(num_elements, 3)
+
+  def test_unknown_trailing_dimension_passes_through(self):
+    dataset = tf.data.Dataset.from_generator(
+        lambda: [
+            {"a": np.zeros([2, 3], np.int32), "b": np.zeros([2, 4], np.int32)}
+        ],
+        output_signature={
+            "a": tf.TensorSpec(shape=(2, None), dtype=tf.int32),
+            "b": tf.TensorSpec(shape=(2, None), dtype=tf.int32),
+        },
+    )
+
+    folded_dataset, unfold = batch_folding.fold_dataset(dataset)
+
+    self.assertIs(folded_dataset, dataset)
+    self.assertIsNone(unfold)
+
+  def test_rejects_other_dataset_types(self):
+    with self.assertRaisesRegex(TypeError, "tf.data.Dataset"):
+      batch_folding.fold_dataset([1, 2, 3])
+
+
+class ModelUnfoldScopeTest(absltest.TestCase):
+
+  def test_training_on_folded_batches_gives_the_same_weights(self):
+    dataset = _dataset()
+    expected = _built_model()
+    expected.fit(dataset, epochs=2, verbose=0)
+    model = _built_model()
+    folded_dataset, unfold = batch_folding.fold_dataset(dataset)
+
+    with batch_folding.model_unfold_scope(model, unfold):
+      model.fit(folded_dataset, epochs=2, verbose=0)
+
+    self.assertNotEmpty(model.trainable_weights)
+    for actual, want in zip(model.get_weights(), expected.get_weights()):
+      np.testing.assert_allclose(actual, want, rtol=1e-6)
+    initial = _built_model()
+    self.assertFalse(
+        np.allclose(model.get_weights()[0], initial.get_weights()[0])
+    )
+
+  def test_restores_the_train_step_on_exit(self):
+    model = _built_model()
+    _, unfold = batch_folding.fold_dataset(_dataset())
+
+    with batch_folding.model_unfold_scope(model, unfold):
+      self.assertIn("train_step", vars(model))
+      self.assertEqual(model.train_step.__name__, "train_step")
+
+    self.assertNotIn("train_step", vars(model))
+    self.assertIsNone(model.train_function)
+    # The model trains on plain batches again.
+    model.fit(_dataset(), epochs=1, verbose=0)
+
+  def test_restores_an_instance_train_step_on_exit(self):
+    model = _built_model()
+    custom_step = lambda state, data: _Model.train_step(model, state, data)
+    model.train_step = custom_step
+    _, unfold = batch_folding.fold_dataset(_dataset())
+
+    with batch_folding.model_unfold_scope(model, unfold):
+      self.assertIsNot(model.train_step, custom_step)
+
+    self.assertIs(model.train_step, custom_step)
+
+  def test_no_unfold_leaves_the_model_untouched(self):
+    model = _built_model()
+
+    with batch_folding.model_unfold_scope(model, None):
+      self.assertNotIn("train_step", vars(model))
+
+  def test_rejects_an_unbuilt_model(self):
+    model = _Model()
+    _, unfold = batch_folding.fold_dataset(_dataset())
+
+    with self.assertRaisesRegex(ValueError, "must be built"):
+      with batch_folding.model_unfold_scope(model, unfold):
+        pass
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/recml/core/training/keras_trainer.py b/recml/core/training/keras_trainer.py
index be865fc..3dcfacc 100644
--- a/recml/core/training/keras_trainer.py
+++ b/recml/core/training/keras_trainer.py
@@ -31,6 +31,7 @@
 import keras
 import orbax.checkpoint as ocp
 import psutil
+from recml.core.training import batch_folding
 from recml.core.training import core
 from recml.core.training import device_prefetch
 from recml.core.utils import py_utils
@@ -130,6 +131,7 @@ def __init__(
       eval_callbacks: Sequence[keras.callbacks.Callback] = (),
       device_prefetch_depth: int = 0,
       skip_batch_logs: bool = False,
+      fold_batch_leaves: bool = False,
   ):
     """Initializes the instance.
 
@@ -165,13 +167,20 @@ def __init__(
         `train_and_evaluate`.
       skip_batch_logs: During `train`, always drop `ProgbarLogger` (`verbose=0`,
         which suppresses `Epoch N/M` and epoch-end stdout progress lines while
-        keeping TensorBoard summaries and checkpoint logs) and disable
-        the per-step hooks of exact `keras_utils.EpochSummaryCallback`
-        instances (custom subclasses are kept as-is), and also skip the
-        per-step conversion of `logs` to Python when no callback needs per-step
-        logs (see `batch_callbacks.ZeroSyncCallbackList`). Epoch-end summaries
-        and checkpoints are unchanged. Not applied in `train_and_evaluate`.
-        Defaults to False.
+        keeping TensorBoard summaries and checkpoint logs) and disable the
+        per-step hooks of exact `keras_utils.EpochSummaryCallback` instances
+        (custom subclasses are kept as-is), and also skip the per-step
+        conversion of `logs` to Python when no callback needs per-step logs (see
+        `batch_callbacks.ZeroSyncCallbackList`). Epoch-end summaries and
+        checkpoints are unchanged. Not applied in `train_and_evaluate`. Defaults
+        to False.
+      fold_batch_leaves: During `train`, concatenate the dense leaves of each
+        batch into one array per dtype and leading dimension in `tf.data`, and
+        split them back inside the jitted train step (see `batch_folding`).
+        Fewer arrays then cross to the device on each step, which reduces the
+        host time per step for batches with many leaves. The model sees the same
+        batches. Requires the JAX backend and a `tf.data.Dataset`. Not applied
+        in `train_and_evaluate`. Defaults to False.
       resume_training_launcher: Launcher to resume training.
       enable_xmanager_measurements: Whether to enable XManager measurements.
       enable_autoxprof: Whether to enable AutoXprof.
@@ -179,12 +188,18 @@ def __init__(
       s2_logging_settings: Settings for S2 logging.
 
     Raises:
-      ValueError: If `device_prefetch_depth` is negative.
+      ValueError: If `device_prefetch_depth` is negative, or if
+        `fold_batch_leaves` is set on a backend other than JAX.
     """
     if device_prefetch_depth < 0:
       raise ValueError(
           f"device_prefetch_depth must be >= 0, got {device_prefetch_depth}."
       )
+    if fold_batch_leaves and keras.backend.backend() != "jax":
+      raise ValueError(
+          "fold_batch_leaves needs the JAX backend, got"
+          f" {keras.backend.backend()!r}."
+      )
 
     keras.utils.set_random_seed(rng_seed)
 
@@ -209,6 +224,7 @@ def __init__(
     self._steps_per_loop = steps_per_loop
     self._device_prefetch_depth = device_prefetch_depth
     self._skip_batch_logs = skip_batch_logs
+    self._fold_batch_leaves = fold_batch_leaves
     self._marker_path = os.path.join(
         model_dir, core.TRAINING_COMPLETE_MARKER_FILE
     )
@@ -377,8 +393,17 @@ def train(self, task: KerasTask) -> core.Logs:
           steps=self._steps_per_loop,
           model=model,
       )
-    with device_prefetch.model_prefetch_scope(
-        model, depth=self._device_prefetch_depth
+    unfold = None
+    if self._fold_batch_leaves:
+      # `fit` builds the model, the compiled loss and metrics, and the
+      # optimizer from its first batch, which would be folded.
+      model._symbolic_build(data_batch=dataset.element_spec)  # pylint: disable=protected-access
+      dataset, unfold = batch_folding.fold_dataset(dataset)
+    with (
+        batch_folding.model_unfold_scope(model, unfold),
+        device_prefetch.model_prefetch_scope(
+            model, depth=self._device_prefetch_depth
+        ),
     ):
       history = model.fit(
           dataset,
@@ -431,10 +456,14 @@ def evaluate(self, task: KerasTask) -> core.Logs:
 
   def train_and_evaluate(self, task: KerasTask) -> core.Logs:
     """Trains and evaluates a Keras model."""
-    if self._device_prefetch_depth or self._skip_batch_logs:
+    if (
+        self._device_prefetch_depth
+        or self._skip_batch_logs
+        or self._fold_batch_leaves
+    ):
       logging.warning(
-          "device_prefetch_depth and skip_batch_logs only apply to `train`;"
-          " ignoring them in `train_and_evaluate`."
+          "device_prefetch_depth, skip_batch_logs and fold_batch_leaves only"
+          " apply to `train`; ignoring them in `train_and_evaluate`."
       )
     train_dataset = task.create_dataset(
         training=True, **self._maybe_get_dataset_kws(task)
diff --git a/recml/core/training/keras_trainer_options_test.py b/recml/core/training/keras_trainer_options_test.py
index f140a41..e4fb807 100644
--- a/recml/core/training/keras_trainer_options_test.py
+++ b/recml/core/training/keras_trainer_options_test.py
@@ -26,6 +26,7 @@
 from recml.core.training import device_prefetch
 from recml.core.training import keras_trainer
 from recml.core.utils import keras_utils
+import tensorflow as tf
 
 
 def _trainer(tmpdir, **overrides):
@@ -68,6 +69,32 @@ def summary(self, print_fn=None):
     del print_fn
 
 
+class _StepModel(_DummyModel):
+  """A `_DummyModel` whose train function runs `train_step`, like Keras."""
+
+  built = True
+
+  def __init__(self):
+    super().__init__()
+    self.datasets_in_fit = []
+
+  def train_step(self, state, data):
+    return data, state
+
+  def _symbolic_build(self, iterator=None, data_batch=None):
+    del iterator, data_batch
+
+  def make_train_function(self, force=False):
+    if self.train_function is not None and not force:
+      return
+    train_step = self.train_step
+    self.train_function = lambda state, it: train_step(state, next(it))
+
+  def fit(self, dataset, **kwargs):
+    self.datasets_in_fit.append(dataset)
+    return super().fit(dataset, **kwargs)
+
+
 def _task(model, exported):
   return types.SimpleNamespace(
       create_dataset=lambda training: list(range(100)),
@@ -147,6 +174,97 @@ def test_train_and_evaluate_warns_that_prefetch_is_not_applied(self):
         trainer.train_and_evaluate(task)
     self.assertIn("device_prefetch_depth", logs.output[0])
 
+  def test_fold_batch_leaves_folds_the_dataset_and_unfolds_in_the_step(self):
+    features = {
+        "a": np.arange(40, dtype=np.float32).reshape([20, 2]),
+        "b": -np.arange(60, dtype=np.float32).reshape([20, 3]),
+    }
+    dataset = tf.data.Dataset.from_tensor_slices(features).batch(
+        2, drop_remainder=True
+    )
+    model = _StepModel()
+    task = types.SimpleNamespace(
+        create_dataset=lambda training: dataset,
+        create_model=lambda: model,
+        export_model=lambda m, model_dir: None,
+    )
+
+    _trainer(self._tmpdir, fold_batch_leaves=True).train(task)
+
+    # `fit` got 1 folded array per batch instead of the 2 leaves.
+    self.assertLen(model.datasets_in_fit, 1)
+    dataset_in_fit = model.datasets_in_fit[0]
+    folded_specs, other_specs = dataset_in_fit.element_spec
+    self.assertLen(folded_specs, 1)
+    self.assertEmpty(other_specs)
+    # The train step got the original batches.
+    self.assertLen(model.batches_in_fit, 5)
+    for step, batch in enumerate(model.batches_in_fit):
+      self.assertCountEqual(batch, ["a", "b"])
+      for key, value in batch.items():
+        np.testing.assert_array_equal(
+            value, features[key][2 * step : 2 * step + 2]
+        )
+    self.assertNotIn("train_step", vars(model))
+
+  def test_rejects_fold_batch_leaves_on_other_backends(self):
+    with mock.patch.object(keras.backend, "backend", return_value="tensorflow"):
+      with self.assertRaisesRegex(ValueError, "needs the JAX backend"):
+        _trainer(self._tmpdir, fold_batch_leaves=True)
+
+  def test_fold_batch_leaves_with_compiled_loss_and_metrics(self):
+    rng = np.random.default_rng(0)
+    x = rng.random((20, 4)).astype(np.float32)
+    y = rng.random((20, 1)).astype(np.float32)
+
+    def build_model():
+      keras.utils.set_random_seed(0)
+      inputs = keras.Input(shape=(4,), dtype="float32")
+      model = keras.Model(inputs=inputs, outputs=keras.layers.Dense(1)(inputs))
+      # The compiled loss and metrics stay unbuilt until the first batch.
+      model.compile(optimizer="sgd", loss="mse", metrics=["mae"])
+      return model
+
+    def train(model_dir, **options):
+      model = build_model()
+      os.makedirs(model_dir)
+      history = _trainer(
+          model_dir, train_steps=10, steps_per_loop=5, **options
+      ).train(
+          types.SimpleNamespace(
+              create_dataset=lambda training: tf.data.Dataset.from_tensor_slices(
+                  (x, y)
+              )
+              .batch(2, drop_remainder=True)
+              .repeat(),
+              create_model=lambda: model,
+              export_model=lambda m, model_dir: None,
+          )
+      )
+      return model, history.history
+
+    want_model, want_history = train(os.path.join(self._tmpdir, "plain"))
+    model, history = train(
+        os.path.join(self._tmpdir, "folded"), fold_batch_leaves=True
+    )
+
+    for actual, want in zip(
+        model.get_weights(), want_model.get_weights(), strict=True
+    ):
+      np.testing.assert_allclose(actual, want, rtol=1e-6)
+    self.assertCountEqual(history, want_history)
+    for key, values in want_history.items():
+      np.testing.assert_allclose(history[key], values, rtol=1e-6, err_msg=key)
+
+  def test_train_and_evaluate_warns_that_folding_is_not_applied(self):
+    trainer = _trainer(self._tmpdir, fold_batch_leaves=True)
+    task = mock.Mock()
+    task.create_dataset.side_effect = RuntimeError("stop")
+    with self.assertLogs(level="WARNING") as logs:
+      with self.assertRaisesRegex(RuntimeError, "stop"):
+        trainer.train_and_evaluate(task)
+    self.assertIn("fold_batch_leaves", logs.output[0])
+
   def test_train_and_evaluate_warns_that_skip_batch_logs_is_not_applied(self):
     trainer = _trainer(self._tmpdir, skip_batch_logs=True)
     task = mock.Mock()