Method _runner_config

axlearn/common/inference_test.py:229–257

(
        self,
        *,
        mesh_shape: tuple[int, int],
        mesh_axis_names: tuple[str, str],
        param_dtype: jnp.dtype,
        inference_dtype: Optional[jnp.dtype],
        data_partition: DataPartitionType,
        ckpt_dir: str,
        use_ema: bool = False,
    )

    # pylint: disable-next=no-self-use
    def _runner_config(
        self,
        *,
        mesh_shape: tuple[int, int],
        mesh_axis_names: tuple[str, str],
        param_dtype: jnp.dtype,
        inference_dtype: Optional[jnp.dtype],
        data_partition: DataPartitionType,
        ckpt_dir: str,
        use_ema: bool = False,
    ):
        inference_runner_cfg = InferenceRunner.default_config().set(
            mesh_shape=mesh_shape,
            mesh_axis_names=mesh_axis_names,
            model=DummyModel.default_config().set(dtype=param_dtype),
            inference_dtype=inference_dtype,
            input_batch_partition_spec=data_partition,
        )
        if use_ema:
            inference_runner_cfg.init_state_builder = RestoreAndConvertBuilder.default_config().set(
                name="builder",
                builder=TensorStoreStateStorageBuilder.default_config().set(
                    dir=ckpt_dir, validation=CheckpointValidationType.CONTAINS_STATE_UP_TO_DTYPE
                ),
                converter=EmaParamsConverter.default_config(),
            )
        else:
            inference_runner_cfg.init_state_builder.set(dir=ckpt_dir)
        return inference_runner_cfg

    # pylint: disable-next=no-self-use
    def _build_ckpt(