def _get_distributed_iterator(self, distribution):
        """Lazily compute layouts to reduce host to device transfer latency."""
        layouts = None
        for data in self.data_adapter.get_jax_iterator():
            if layouts is None:
                layouts = tree.map_structure(
                    lambda d: distribution.get_data_layout(
                        d.shape
                    ).backend_layout,
                    data,
                )
            yield _distribute_data(data, layouts)
Lazily compute layouts to reduce host to device transfer latency.
keras/src/backend/jax/trainer.py