Configuration¶

def build(
    self,
    in_channels: int,
    out_channels: int,
    max_patch_size: tuple[int, int],
    implementation: PerceiverImpl,
) -> nn.Module:
    """Build a regular Perceiver (used by the encoder)."""
    # This is not really a "frequency" but a maximum of the width appears to be reasonable from looking at the code.
    max_freq = max(*max_patch_size)

    num_freq_bands = 4
    fourier_dim = fourier_features_2d_dim(num_freq_bands)
    # Use the same explicit 2D Fourier features in both implementations so
    # intra-patch positions are encoded equivalently.
    if _use_flash(implementation):
        try:
            from flash_perceiver import Perceiver as FlashPerceiver  # type: ignore
        except ImportError as e:
            raise _flash_import_error() from e
        from einops.layers.torch import Rearrange

        perceiver: nn.Module = nn.Sequential(
            FourierFeatures2D(num_freq_bands=num_freq_bands, max_freq=max_freq),
            Rearrange("b ph pw v -> b (ph pw) v"),
            FlashPerceiver(
                depth=self.depth,
                input_dim=in_channels + fourier_dim,
                output_dim=out_channels,
                output_mode="average",
                latent_dim=self.latent_dim,
                num_latents=self.num_latents,
                use_flash_attn=True,
                weight_tie_layers=True,
                self_per_cross_attn=2,
            ),
        )
    elif _use_naive(implementation):
        perceiver = nn.Sequential(
            FourierFeatures2D(num_freq_bands=num_freq_bands, max_freq=max_freq),
            NaivePerceiver(
                # Required by perceiver-pytorch even when its internal
                # Fourier encoding is disabled below.
                num_freq_bands=num_freq_bands,
                max_freq=max_freq,
                depth=self.depth,
                input_axis=2,
                input_channels=in_channels + fourier_dim,
                num_classes=out_channels,
                latent_dim=self.latent_dim,
                num_latents=self.num_latents,
                weight_tie_layers=True,
                fourier_encode_data=False,
                self_per_cross_attn=2,
            ),
        )
    else:
        raise ValueError(f"Unknown perceiver implementation: {implementation}.")

    return perceiver

def build_io(
    self,
    in_channels: int,
    queries_dim: int,
    out_channels: int,
    implementation: PerceiverImpl,
) -> nn.Module:
    """Build a PerceiverIO (used by the decoder)."""
    if _use_flash(implementation):
        try:
            from flash_perceiver.perceiver import (  # type: ignore
                PerceiverIO as FlashPerceiverIO,  # type: ignore
            )
        except ImportError as e:
            raise _flash_import_error() from e
        perceiver_io: nn.Module = FlashPerceiverIO(
            depth=self.depth,
            input_dim=in_channels,
            query_dim=queries_dim,
            proj_dim=out_channels,
            num_latents=self.num_latents,
            latent_dim=self.latent_dim,
            use_flash_attn=True,
            weight_tie_layers=True,
        )
    elif _use_naive(implementation):
        from perceiver_pytorch.perceiver_io import PerceiverIO as NaivePerceiverIO

        perceiver_io = NaivePerceiverIO(
            depth=self.depth,
            dim=in_channels,
            queries_dim=queries_dim,
            logits_dim=out_channels,
            num_latents=self.num_latents,
            latent_dim=self.latent_dim,
            weight_tie_layers=True,
            decoder_ff=True,
        )
    else:
        raise ValueError(f"Unknown perceiver implementation: {implementation}.")

    return perceiver_io

@classmethod
def from_yaml_and_cli(
    cls,
    args_to_parse: list[str] | None = None,
) -> Self:
    """Load config from YAML & CLI with validation."""
    parser = argparse.ArgumentParser(
        description=cls.__doc__,
        epilog=textwrap.dedent(
            """
            CONFIG may be a path to a YAML file or, when samudra is installed,
            the name of a bundled preset such as `samudra_om4/train.yaml`.
            YAML files can include other YAML files using the !include tag,
            as in `data: !include ../data/om4.yaml` (resolved relative to the
            including file). You can also replace any JSON argument listed above
            with a YAML file by specifying it with an @ symbol,
            eg `--some_param=@../data/om4.yaml`.
            """
        ),
    )
    parser.add_argument(
        "config",
        type=str,
        help="Path to a config YAML file, or a bundled preset name",
    )

    cli_source = IncludeYamlCliSettingsSource(
        cls,
        root_parser=parser,
        # If args_to_parse is None, we parse argv, which is what `True` does
        cli_parse_args=args_to_parse if args_to_parse is not None else True,
    )

    # We do this after creating CliSettingsSource (which populates the parser)
    # so the help is complete on error.
    args = parser.parse_args(args_to_parse)

    # Then we read the YAML file specified in the CLI. A bare path resolves
    # against the filesystem; if it's missing we try a preset bundled in the
    # installed package before giving up.
    # Note that by default, YamlConfigSettingsSource will ignore missing files
    config_path = resolve_config_path(args.config)
    if not config_path.exists():
        raise FileNotFoundError(
            f"Config file `{args.config}` (full path: {config_path}) not found"
        )
    yaml_values = YamlConfigSettingsSource(cls, yaml_file=config_path)()

    return cls(
        _cli_settings_source=cli_source,
        **yaml_values,
    )

def save_yaml(self, save_path: Path) -> None:
    """Save config to YAML file."""
    with open(save_path, "w") as f:
        yaml.dump(self.model_dump(), f)