Method iter_source

axlearn/common/input_grain_test.py:420–426

(dispatch_config)

        # Test against iter ds.
        def iter_source(dispatch_config):
            iter_ds = fake_grain_source(list(range(10)))
            iter_ds = maybe_to_iter_dataset(iter_ds)
            iter_ds = iter_ds.map(lambda x: x + 1)
            return per_feed_batch(
                iter_ds, global_batch_size=10, dispatch_config=dispatch_config, drop_remainder=False
            )

        cfg = Input.default_config().set(source=iter_source, input_dispatcher=input_dispatcher)
        x = cfg.set(name="test").instantiate(parent=None)