Method test_zipinput

axlearn/common/input_composite_test.py:96–151

(self, is_training)

    @parameterized.parameters(False, True)
    def test_zipinput(self, is_training):
        cfg = ZipInput.default_config().set(
            name="input",
            is_training=is_training,
            inputs={
                "0": self._input_config(
                    [1, 2, 3, 4, 5, 6, 7, 8, 9], batch_size=3, out_signature="number_1", repeat=None
                ),
                "1": self._input_config(
                    [11, 12, 13, 14, 15, 16, 17, 18, 19],
                    batch_size=2,
                    out_signature="number_2",
                    repeat=None,
                ),
            },
        )
        dataset = cfg.instantiate(parent=None)
        expected_train_batches = [
            {
                "0": {"number_1": jnp.asarray([1, 2, 3]), "index": jnp.asarray([0, 1, 2])},
                "1": {"number_2": jnp.asarray([11, 12]), "index": jnp.asarray([0, 1])},
            },
            {
                "0": {"number_1": jnp.asarray([4, 5, 6]), "index": jnp.asarray([3, 4, 5])},
                "1": {"number_2": jnp.asarray([13, 14]), "index": jnp.asarray([2, 3])},
            },
        ]
        expected_eval_batches = [
            {
                "0": {"number_1": jnp.asarray([1, 2, 3]), "index": jnp.asarray([0, 1, 2])},
                "1": {"number_2": jnp.asarray([11, 12]), "index": jnp.asarray([0, 1])},
            },
            {
                "0": {"number_1": jnp.asarray([4, 5, 6]), "index": jnp.asarray([3, 4, 5])},
                "1": {"number_2": jnp.asarray([13, 14]), "index": jnp.asarray([2, 3])},
            },
        ]
        expected_batches = expected_train_batches if is_training else expected_eval_batches
        for batch_index, batch in enumerate(dataset):
            if batch_index >= len(expected_batches):
                break
            self.assertNestedAllClose(as_tensor(expected_batches[batch_index]), batch)
        self.assertEqual(
            batch_index,
            len(expected_batches),  # pylint: disable=undefined-loop-variable
        )

        dataset = cfg.instantiate(parent=None)
        for batch_index, batch in enumerate(dataset.dataset()):
            if batch_index >= len(expected_batches):
                break
            self.assertNestedAllClose(as_tensor(expected_batches[batch_index]), batch)
        self.assertEqual(
            batch_index,
            len(expected_batches),  # pylint: disable=undefined-loop-variable
        )