Shared Layers API

Auto-generated from source — building blocks shared across model families.

keras_climate.utils.layers

keras_climate.utils.layers.DropPath

DropPath(drop_prob=0.0, **kwargs)

Bases: Layer

Source code in keras_climate/utils/layers.py

def __init__(self, drop_prob=0.0, **kwargs):
    super().__init__(**kwargs)
    self.drop_prob = drop_prob

keras_climate.utils.layers.MLP

MLP(hidden_dim, out_dim=None, act='gelu', drop=0.0, **kwargs)

Bases: Layer

Source code in keras_climate/utils/layers.py

def __init__(self, hidden_dim, out_dim=None, act="gelu", drop=0.0, **kwargs):
    super().__init__(**kwargs)
    self.hidden_dim = hidden_dim
    self.out_dim = out_dim
    self.act = act
    self.drop = drop

keras_climate.utils.layers.GatedResidualNetwork

GatedResidualNetwork(units, dropout=0.1, **kwargs)

Bases: Layer

Source code in keras_climate/utils/layers.py

def __init__(self, units, dropout=0.1, **kwargs):
    super().__init__(**kwargs)
    self.units = units
    self.dropout = dropout

keras_climate.utils.layers.MultiHeadSelfAttention

MultiHeadSelfAttention(dim, num_heads=8, qkv_bias=True, attn_drop=0.0, proj_drop=0.0, **kwargs)

Bases: Layer

Source code in keras_climate/utils/layers.py

def __init__(
    self, dim, num_heads=8, qkv_bias=True, attn_drop=0.0, proj_drop=0.0, **kwargs
):
    super().__init__(**kwargs)
    self.dim = dim
    self.num_heads = num_heads
    self.head_dim = dim // num_heads
    self.scale = self.head_dim**-0.5
    self.qkv_bias = qkv_bias
    self.attn_drop_rate = attn_drop
    self.proj_drop_rate = proj_drop

keras_climate.utils.layers.SpatialReductionAttention

SpatialReductionAttention(dim, num_heads=8, sr_ratio=1, qkv_bias=True, attn_drop=0.0, proj_drop=0.0, **kwargs)

Bases: Layer

Source code in keras_climate/utils/layers.py

def __init__(
    self,
    dim,
    num_heads=8,
    sr_ratio=1,
    qkv_bias=True,
    attn_drop=0.0,
    proj_drop=0.0,
    **kwargs,
):
    super().__init__(**kwargs)
    self.dim = dim
    self.num_heads = num_heads
    self.head_dim = dim // num_heads
    self.scale = self.head_dim**-0.5
    self.sr_ratio = sr_ratio
    self.qkv_bias = qkv_bias
    self.attn_drop_rate = attn_drop
    self.proj_drop_rate = proj_drop

keras_climate.utils.layers.TransformerEncoderBlock

TransformerEncoderBlock(dim, num_heads, mlp_ratio=4.0, qkv_bias=True, drop=0.0, attn_drop=0.0, drop_path=0.0, **kwargs)

Bases: Layer

Source code in keras_climate/utils/layers.py

def __init__(
    self,
    dim,
    num_heads,
    mlp_ratio=4.0,
    qkv_bias=True,
    drop=0.0,
    attn_drop=0.0,
    drop_path=0.0,
    **kwargs,
):
    super().__init__(**kwargs)
    self.norm1 = layers.LayerNormalization(epsilon=1e-6, name="norm1")
    self.attn = MultiHeadSelfAttention(
        dim, num_heads, qkv_bias, attn_drop, drop, name="attn"
    )
    self.drop_path = DropPath(drop_path)
    self.norm2 = layers.LayerNormalization(epsilon=1e-6, name="norm2")
    self.mlp = MLP(int(dim * mlp_ratio), dim, drop=drop, name="mlp")

keras_climate.utils.layers.SegformerBlock

SegformerBlock(dim, num_heads, mlp_ratio=4.0, sr_ratio=1, drop=0.0, attn_drop=0.0, drop_path=0.0, **kwargs)

Bases: Layer

Source code in keras_climate/utils/layers.py

def __init__(
    self,
    dim,
    num_heads,
    mlp_ratio=4.0,
    sr_ratio=1,
    drop=0.0,
    attn_drop=0.0,
    drop_path=0.0,
    **kwargs,
):
    super().__init__(**kwargs)
    self.norm1 = layers.LayerNormalization(epsilon=1e-6, name="norm1")
    self.attn = SpatialReductionAttention(
        dim, num_heads, sr_ratio, attn_drop=attn_drop, proj_drop=drop, name="attn"
    )
    self.drop_path = DropPath(drop_path)
    self.norm2 = layers.LayerNormalization(epsilon=1e-6, name="norm2")
    hidden = int(dim * mlp_ratio)
    self.fc1 = layers.Dense(hidden, name="mixffn_fc1")
    self.dwconv = layers.DepthwiseConv2D(3, padding="same", name="mixffn_dwconv")
    self.act = layers.Activation("gelu")
    self.fc2 = layers.Dense(dim, name="mixffn_fc2")
    self.drop = layers.Dropout(drop)

keras_climate.utils.layers.PatchEmbed2D

PatchEmbed2D(patch_size, embed_dim, in_chans=None, norm=True, **kwargs)

Bases: Layer

Source code in keras_climate/utils/layers.py

def __init__(self, patch_size, embed_dim, in_chans=None, norm=True, **kwargs):
    super().__init__(**kwargs)
    self.patch_size = patch_size
    self.embed_dim = embed_dim
    self.use_norm = norm

keras_climate.utils.layers.OverlapPatchEmbed

OverlapPatchEmbed(patch_size, stride, embed_dim, **kwargs)

Bases: Layer

Source code in keras_climate/utils/layers.py

def __init__(self, patch_size, stride, embed_dim, **kwargs):
    super().__init__(**kwargs)
    self.patch_size = patch_size
    self.stride = stride
    self.embed_dim = embed_dim

keras_climate.utils.layers.PatchEmbed3D

PatchEmbed3D(patch_size, tubelet_size, embed_dim, **kwargs)

Bases: Layer

Source code in keras_climate/utils/layers.py

def __init__(self, patch_size, tubelet_size, embed_dim, **kwargs):
    super().__init__(**kwargs)
    self.patch_size = patch_size
    self.tubelet_size = tubelet_size
    self.embed_dim = embed_dim

keras_climate.utils.layers.DoubleConv

DoubleConv(filters, **kwargs)

Bases: Layer

Source code in keras_climate/utils/layers.py

def __init__(self, filters, **kwargs):
    super().__init__(**kwargs)
    self.conv1 = ConvBNAct(filters, 3, name="conv1")
    self.conv2 = ConvBNAct(filters, 3, name="conv2")

keras_climate.utils.layers.ASPP

ASPP(filters=256, rates=(6, 12, 18), **kwargs)

Bases: Layer

Source code in keras_climate/utils/layers.py

def __init__(self, filters=256, rates=(6, 12, 18), **kwargs):
    super().__init__(**kwargs)
    self.filters = filters
    self.rates = rates

keras_climate.utils.layers.ConvBNAct

ConvBNAct(filters, kernel_size=3, strides=1, dilation_rate=1, act='relu', use_bn=True, **kwargs)

Bases: Layer

Source code in keras_climate/utils/layers.py

def __init__(
    self,
    filters,
    kernel_size=3,
    strides=1,
    dilation_rate=1,
    act="relu",
    use_bn=True,
    **kwargs,
):
    super().__init__(**kwargs)
    # Keras's `padding="same"` pads *asymmetrically* (extra on the
    # bottom/right) whenever `strides > 1` and the input size isn't an
    # exact multiple of the stride in a way that makes the needed
    # padding even - PyTorch reference convs almost always pad
    # symmetrically via an explicit `padding=dilation*(k-1)//2`, so
    # relying on "same" here would silently misalign every ported
    # strided conv (stem/downsampling convs) against its source
    # checkpoint. Pad explicitly + "valid" whenever strides > 1 to
    # reproduce PyTorch's convention exactly; for strides == 1, "same"
    # is already symmetric, so there's nothing to fix.
    if strides > 1:
        pad = dilation_rate * (kernel_size - 1) // 2
        self.pad = layers.ZeroPadding2D(pad, name="pad") if pad > 0 else None
        conv_padding = "valid"
    else:
        self.pad = None
        conv_padding = "same"
    self.conv = layers.Conv2D(
        filters,
        kernel_size,
        strides=strides,
        padding=conv_padding,
        dilation_rate=dilation_rate,
        use_bias=not use_bn,
        name="conv",
    )
    self.bn = layers.BatchNormalization(epsilon=1e-5, name="bn") if use_bn else None
    self.act = layers.Activation(act) if act else None

keras_climate.utils.layers.rfft2_hw

rfft2_hw(x)

Source code in keras_climate/utils/layers.py

def rfft2_hw(x):
    x_p = ops.transpose(x, (0, 3, 1, 2))
    re, im = ops.rfft(x_p)
    re_t = ops.transpose(re, (0, 1, 3, 2))
    im_t = ops.transpose(im, (0, 1, 3, 2))
    re2, im2 = ops.fft((re_t, im_t))
    re2 = ops.transpose(re2, (0, 1, 3, 2))
    im2 = ops.transpose(im2, (0, 1, 3, 2))
    re2 = ops.transpose(re2, (0, 2, 3, 1))
    im2 = ops.transpose(im2, (0, 2, 3, 1))
    return re2, im2

keras_climate.utils.layers.irfft2_hw

irfft2_hw(re, im, H, W)

Source code in keras_climate/utils/layers.py

def irfft2_hw(re, im, H, W):
    re_p = ops.transpose(re, (0, 3, 1, 2))
    im_p = ops.transpose(im, (0, 3, 1, 2))
    re_t = ops.transpose(re_p, (0, 1, 3, 2))
    im_t = ops.transpose(im_p, (0, 1, 3, 2))
    re1, im1 = _ifft_last_axis(re_t, im_t, H)
    re1 = ops.transpose(re1, (0, 1, 3, 2))
    im1 = ops.transpose(im1, (0, 1, 3, 2))
    out = ops.irfft((re1, im1), fft_length=W)
    return ops.transpose(out, (0, 2, 3, 1))

keras_climate.utils.layers.unpatchify_2d

unpatchify_2d(x, grid_h, grid_w, patch_size, out_channels, name='unpatchify')

Source code in keras_climate/utils/layers.py

def unpatchify_2d(x, grid_h, grid_w, patch_size, out_channels, name="unpatchify"):
    out_h, out_w = grid_h * patch_size, grid_w * patch_size

    def _unpatchify(t):
        B = ops.shape(t)[0]
        t = ops.reshape(t, (B, grid_h, grid_w, patch_size, patch_size, out_channels))
        t = ops.transpose(t, (0, 1, 3, 2, 4, 5))
        return ops.reshape(t, (B, out_h, out_w, out_channels))

    return layers.Lambda(
        _unpatchify, output_shape=(out_h, out_w, out_channels), name=name
    )(x)

keras_climate.utils.layers.sincos_position_embedding

sincos_position_embedding(length, dim)

Source code in keras_climate/utils/layers.py

def sincos_position_embedding(length, dim):
    position = np.arange(length)[:, None]
    div_term = np.exp(np.arange(0, dim, 2) * -(math.log(10000.0) / dim))
    pe = np.zeros((length, dim), dtype=np.float32)
    pe[:, 0::2] = np.sin(position * div_term)
    pe[:, 1::2] = np.cos(position * div_term)
    return pe

keras_climate.utils.layers.sincos_position_embedding_2d

sincos_position_embedding_2d(h, w, dim)

Source code in keras_climate/utils/layers.py

def sincos_position_embedding_2d(h, w, dim):
    assert dim % 2 == 0
    grid_h = np.arange(h, dtype=np.float32)
    grid_w = np.arange(w, dtype=np.float32)
    grid = np.meshgrid(grid_w, grid_h)
    grid = np.stack(grid, axis=0).reshape(2, 1, h, w)

    def embed_1d(pos, d):
        omega = np.arange(d // 2, dtype=np.float32) / (d / 2.0)
        omega = 1.0 / (10000**omega)
        pos = pos.reshape(-1)
        out = np.einsum("m,d->md", pos, omega)
        return np.concatenate([np.sin(out), np.cos(out)], axis=1)

    emb_h = embed_1d(grid[0], dim // 2)
    emb_w = embed_1d(grid[1], dim // 2)
    return np.concatenate([emb_h, emb_w], axis=1)