Modules¶

def __init__(
    self,
    in_channels: int,
    ch_width: list[int],
    dilation: list[int],
    n_layers: list[int],
    pad: str,
    create_block: CoreBlockBuilder,
    downsampling_block: nn.Module,
    create_upsampling_block: UpsamplingBlockBuilder,
    checkpointing: "Checkpointing | None",
    drop_path_rate: float = 0.0,
):
    super().__init__()
    self.in_channels = in_channels
    self.out_channels: int = ch_width[0]

    # Create local copies of config lists that will be reversed
    ch_width = [in_channels] + ch_width.copy()
    dilation = dilation.copy()
    n_layers = n_layers.copy()
    self.pad = pad

    match checkpointing:
        case "all":
            self.checkpoint_all = True
            checkpoint_simple = False
        case "simple":
            self.checkpoint_all = False
            checkpoint_simple = True
        case None:
            self.checkpoint_all = False
            checkpoint_simple = False
        case _:
            assert_never(checkpointing)

    # going down
    layers: list[nn.Module] = []
    for i, (a, b) in enumerate(pairwise(ch_width)):
        # Core block
        layers.append(
            create_block(
                in_channels=a,
                out_channels=b,
                dilation=dilation[i],
                n_layers=n_layers[i],
                pad=pad,
                checkpoint_simple=checkpoint_simple,
            )
        )
        # Down sampling block
        layers.append(downsampling_block)

    # Middle block
    layers.append(
        create_block(
            in_channels=b,
            out_channels=b,
            dilation=dilation[i],
            n_layers=n_layers[i],
            pad=pad,
            checkpoint_simple=checkpoint_simple,
        )
    )

    # First upsampling
    layers.append(create_upsampling_block(in_channels=b, out_channels=b))

    # Reverse for upsampling path
    ch_width.reverse()
    dilation.reverse()
    n_layers.reverse()

    # going up
    for i, (a, b) in enumerate(pairwise(ch_width[:-1])):
        layers.append(
            create_block(
                in_channels=a,
                out_channels=b,
                dilation=dilation[i],
                n_layers=n_layers[i],
                pad=pad,
                checkpoint_simple=checkpoint_simple,
            )
        )
        layers.append(create_upsampling_block(in_channels=b, out_channels=b))

    # Final conv block
    layers.append(
        create_block(
            in_channels=b,
            out_channels=b,  # this is the same as self.out_channels
            dilation=dilation[i],
            n_layers=n_layers[i],
            pad=pad,
            checkpoint_simple=checkpoint_simple,
        )
    )

    first_block = layers[0]
    assert isinstance(first_block, CoreBlock)
    self.N_pad = first_block.N_pad
    self.layers = nn.ModuleList(layers)
    self.num_steps = int(len(ch_width) - 1)
    self.drop_path = DropPath(drop_path_rate)

def __init__(
    self,
    in_channels: int,
    out_channels: int,
    patch_extent: tuple[float, float],
    perceiver: nn.Module,
) -> None:
    super().__init__()
    self.in_channels = in_channels
    self.out_channels: int = out_channels  # aka, `embed_dim`.
    self.patch_extent = patch_extent
    self.perceiver = perceiver
    # TODO(#451): The input to these position and scale linear units could be a hparam.
    self.pos_embed = nn.Linear(self.out_channels, self.out_channels)
    self.scale_embed = nn.Linear(self.out_channels, self.out_channels)

def __init__(
    self,
    in_channels: int,
    out_channels: int,
    patch_extent: tuple[float, float],
    queries_dim: int,
    perceiver_io: nn.Module,
    window_patches: int | None,
    context_patches: int | None,
) -> None:
    super().__init__()
    self.in_channels = in_channels
    self.out_channels = out_channels
    self.patch_extent = patch_extent
    if window_patches is None and context_patches is not None:
        raise ValueError(
            "window_patches must be set in order for context_patches to be set."
        )
    self.window_patches = window_patches
    self.context_patches = context_patches

    # TODO(#451): The input to these position and scale linear units could be a hparam.
    # Same pos/scale linear layers as the encoder, but applied *before* the
    # perceiver (the encoder applies them after).
    self.pos_embed = nn.Linear(in_channels, in_channels)
    self.scale_embed = nn.Linear(in_channels, in_channels)

    # Embed 3D unit-sphere coordinates into queries_dim for the PerceiverIO decoder head.
    self.query_embed = nn.Linear(3, queries_dim)

    self.perceiver_io = perceiver_io

def __init__(self, in_channels: int, out_channels: int):
    super().__init__()
    self.linear = torch.nn.Linear(in_channels, out_channels)

def __init__(self, upsampling: int | tuple[int, int] = 2):
    super().__init__()
    if isinstance(upsampling, int):
        upsampling = (upsampling, upsampling)
    if tuple(upsampling) != (2, 2):
        raise ValueError(
            "ZonallyPeriodicBilinearUpsample only supports 2x upsampling"
        )
    self.scale_h, self.scale_w = upsampling

def __init__(self, drop_prob: float = 0.0):
    super().__init__()
    self.dropout = torch.nn.Dropout(p=drop_prob)

def __init__(
    self,
    in_channels: int = 300,
    out_channels: int = 1,
    kernel_size: int = 3,
    dilation: int = 1,
    n_layers: int = 1,
    activation: Callable[[], torch.nn.Module] = CappedGELU,
    pad="circular",
    upscale_factor: int = 4,
    norm="batch",
    checkpoint_simple: bool = False,
    pointwise_linear: bool = False,
):
    super().__init__(in_channels, out_channels, kernel_size, dilation, pad)
    assert n_layers == 1, "Can only use a single layer here!"

    # Instantiate pointwise linear to increase/decrease channel depth if necessary
    if in_channels == out_channels:
        self.skip_module = lambda x: x  # Identity-function required in forward pass
    else:
        self.skip_module = _pointwise(pointwise_linear, in_channels, out_channels)

    # Convolution block
    convblock: list[torch.nn.Module] = []
    convblock.append(
        torch.nn.Conv2d(
            in_channels=in_channels,
            out_channels=int(in_channels * upscale_factor),
            kernel_size=kernel_size,
            dilation=dilation,
        )
    )
    # BatchNorm
    if norm == "batch":
        convblock.append(torch.nn.BatchNorm2d(in_channels * upscale_factor))
    # Instance Norm
    elif norm == "instance":
        convblock.append(torch.nn.InstanceNorm2d(in_channels * upscale_factor))
    elif norm == "nonorm":
        pass
    else:
        raise NotImplementedError
    if activation is not None:
        convblock.append(activation())
    convblock.append(
        torch.nn.Conv2d(
            in_channels=int(in_channels * upscale_factor),
            out_channels=int(in_channels * upscale_factor),
            kernel_size=kernel_size,
            dilation=dilation,
        )
    )
    # BatchNorm
    if norm == "batch":
        convblock.append(torch.nn.BatchNorm2d(in_channels * upscale_factor))
    # Instance Norm
    elif norm == "instance":
        convblock.append(torch.nn.InstanceNorm2d(in_channels * upscale_factor))
    elif norm == "nonorm":
        pass
    else:
        raise NotImplementedError
    if activation is not None:
        convblock.append(activation())
    # Linear postprocessing
    convblock.append(
        _pointwise(
            pointwise_linear, int(in_channels * upscale_factor), out_channels
        )
    )
    self.convblock = torch.nn.Sequential(*convblock)
    self.checkpoint_simple = checkpoint_simple

def __init__(self, **kwargs):
    """
    :param kwargs: passed to torch.nn.ReLU
    """
    super().__init__()
    self.relu = torch.nn.ReLU(**kwargs)

def __init__(self, cap_value=10.0, **kwargs):
    """
    :param cap_value: float: value at which to clip activation
    :param kwargs: passed to torch.nn.LeadyReLU
    """
    super().__init__()
    self.relu = torch.nn.LeakyReLU(**kwargs)
    self.cap = torch.nn.Buffer(torch.tensor(cap_value, dtype=torch.float32))

def __init__(self, cap_value=10.0, **kwargs):
    """
    :param cap_value: float: value at which to clip activation
    :param kwargs: passed to torch.nn.LeadyReLU
    """
    super().__init__()
    self.gelu = torch.nn.GELU(**kwargs)
    self.cap = torch.nn.Buffer(torch.tensor(cap_value, dtype=torch.float32))