Method map_source

axlearn/common/input_grain_test.py:406–411

(dispatch_config)

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

        # Configure dispatch.
        cfg = Input.default_config().set(source=map_source, input_dispatcher=input_dispatcher)