Method get_test_config

tests/data_loader_test.py:51–62

Generate config for tests

(self, reuse_example_batch, **kwargs)

    self.mock_data_iterator = MagicMock()

  def get_test_config(self, reuse_example_batch, **kwargs):
    """Generate config for tests"""
    args = {
        "run_name": "test",
        "enable_checkpointing": False,
        "reuse_example_batch": reuse_example_batch,
    }
    args.update(kwargs)
    return pyconfig.initialize(
        [None, os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml")],
        **args,
    )

  def test_load_next_batch_success(self):
    expected_shape = [jax.device_count(), self.config.max_target_length]