Skip to content

Neural Operators API

Auto-generated from source — see Neural Operators for narrative descriptions.

AFNO

keras_climate.operators.afno.AFNOOperator 

AFNOOperator(input_shape=(64, 64, 3), out_channels=3, patch_size=8, embed_dim=256, depth=8, num_blocks=8, mlp_ratio=4.0, name='afno_operator')
Source code in keras_climate/operators/afno.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
def AFNOOperator(
    input_shape=(64, 64, 3),
    out_channels=3,
    patch_size=8,
    embed_dim=256,
    depth=8,
    num_blocks=8,
    mlp_ratio=4.0,
    name="afno_operator",
):
    grid = input_shape[0] // patch_size
    inputs = keras.Input(shape=input_shape, name="field")
    x = layers.Conv2D(
        embed_dim, kernel_size=patch_size, strides=patch_size, name="patch_embed_proj"
    )(inputs)

    for i in range(depth):
        x = AFNOBlock(embed_dim, num_blocks, mlp_ratio, name=f"block{i}")(x)

    x = layers.Dense(patch_size * patch_size * out_channels, name="head")(x)
    x = unpatchify_2d(x, grid, grid, patch_size, out_channels)
    return keras.Model(inputs, x, name=name)

keras_climate.operators.afno.AFNO2D 

AFNO2D(hidden_size, num_blocks=8, sparsity_threshold=0.01, hidden_size_factor=1, **kwargs)

Bases: Layer

Source code in keras_climate/operators/afno.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
def __init__(
    self,
    hidden_size,
    num_blocks=8,
    sparsity_threshold=0.01,
    hidden_size_factor=1,
    **kwargs,
):
    super().__init__(**kwargs)
    self.hidden_size = hidden_size
    self.num_blocks = num_blocks
    self.block_size = hidden_size // num_blocks
    self.hidden_size_factor = hidden_size_factor
    self.sparsity_threshold = sparsity_threshold

keras_climate.operators.afno.AFNOBlock 

AFNOBlock(hidden_size, num_blocks=8, mlp_ratio=4.0, sparsity_threshold=0.01, hidden_size_factor=1, **kwargs)

Bases: Layer

Source code in keras_climate/operators/afno.py 
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def __init__(
    self,
    hidden_size,
    num_blocks=8,
    mlp_ratio=4.0,
    sparsity_threshold=0.01,
    hidden_size_factor=1,
    **kwargs,
):
    super().__init__(**kwargs)
    self.norm1 = layers.LayerNormalization(epsilon=1e-6, name="norm1")
    self.filter = AFNO2D(
        hidden_size,
        num_blocks,
        sparsity_threshold,
        hidden_size_factor,
        name="filter",
    )
    self.norm2 = layers.LayerNormalization(epsilon=1e-6, name="norm2")
    self.mlp = MLP(int(hidden_size * mlp_ratio), hidden_size, name="mlp")