Method batches

axlearn/common/input_base.py:213–225

Yields per-feed physical input batches (using `input_dispatcher` if configured). The caller should use `host_to_global_array` to construct a global physical batch from the per-feed physical batches returned from this method. See also `dispatch_global_batch` for constructing

(self, it: Iterator[Nested[Tensor]])

        yield from self.batches(iter(self.dataset()))

    def batches(self, it: Iterator[Nested[Tensor]]) -> Iterator[Nested[Tensor]]:
        """Yields per-feed physical input batches (using `input_dispatcher` if configured).

        The caller should use `host_to_global_array` to construct a global physical batch from the
        per-feed physical batches returned from this method.

        See also `dispatch_global_batch` for constructing a global logical batch.
        """
        for input_batch in it:
            input_batch = as_numpy_array(input_batch)
            if "input_dispatcher" in self.children:
                input_batch = self.input_dispatcher.logical_to_physical_batch(input_batch)
            yield input_batch

    def dispatch_global_batch(self, global_physical_batch: Nested[Tensor]) -> Nested[Tensor]:
        """Converts a global physical batch to a global logical batch.