Skip to content

Weight Mapping Builders API

Auto-generated from source — see Weight Porting for how to write a new one.

Mapping builders

keras_climate.weights.mappings.afno_mapping.build_afno_mapper 

build_afno_mapper(depth, keras_prefix='')
Source code in keras_climate/weights/mappings/afno_mapping.py 
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
def build_afno_mapper(depth, keras_prefix=""):
    kp = f"{keras_prefix}/" if keras_prefix else ""
    rules = [
        (r"^patch_embed\.proj\.weight$", f"{kp}patch_embed_proj/kernel"),
        (r"^patch_embed\.proj\.bias$", f"{kp}patch_embed_proj/bias"),
        (r"^head\.weight$", f"{kp}head/kernel"),
        (r"^head\.bias$", f"{kp}head/bias"),
    ]
    for i in range(depth):
        prefix = f"blocks.{i}"
        keras_block = f"{kp}block{i}"
        rules += [
            (rf"^{prefix}\.norm1\.weight$", f"{keras_block}/norm1/gamma"),
            (rf"^{prefix}\.norm1\.bias$", f"{keras_block}/norm1/beta"),
            (rf"^{prefix}\.filter\.w1$", f"{keras_block}/filter/w1"),
            (rf"^{prefix}\.filter\.b1$", f"{keras_block}/filter/b1"),
            (rf"^{prefix}\.filter\.w2$", f"{keras_block}/filter/w2"),
            (rf"^{prefix}\.filter\.b2$", f"{keras_block}/filter/b2"),
            (rf"^{prefix}\.norm2\.weight$", f"{keras_block}/norm2/gamma"),
            (rf"^{prefix}\.norm2\.bias$", f"{keras_block}/norm2/beta"),
            (rf"^{prefix}\.mlp\.fc1\.weight$", f"{keras_block}/mlp/fc1/kernel"),
            (rf"^{prefix}\.mlp\.fc1\.bias$", f"{keras_block}/mlp/fc1/bias"),
            (rf"^{prefix}\.mlp\.fc2\.weight$", f"{keras_block}/mlp/fc2/kernel"),
            (rf"^{prefix}\.mlp\.fc2\.bias$", f"{keras_block}/mlp/fc2/bias"),
        ]

    compiled = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, repl in compiled:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)
        return None

    return mapper

keras_climate.weights.mappings.anysat_mapping.build_anysat_mapper 

build_anysat_mapper(modalities, keras_prefix='anysat_encoder')
Source code in keras_climate/weights/mappings/anysat_mapping.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
def build_anysat_mapper(modalities, keras_prefix="anysat_encoder"):
    rules = [
        (r"^cls_token$", f"{keras_prefix}/cls_token"),
        (
            r"^pos_encoding\.coord_mlp\.fc1\.weight$",
            f"{keras_prefix}/pos_encoding/coord_mlp/fc1/kernel",
        ),
        (
            r"^pos_encoding\.coord_mlp\.fc1\.bias$",
            f"{keras_prefix}/pos_encoding/coord_mlp/fc1/bias",
        ),
        (
            r"^pos_encoding\.coord_mlp\.fc2\.weight$",
            f"{keras_prefix}/pos_encoding/coord_mlp/fc2/kernel",
        ),
        (
            r"^pos_encoding\.coord_mlp\.fc2\.bias$",
            f"{keras_prefix}/pos_encoding/coord_mlp/fc2/bias",
        ),
    ]
    for m in modalities:
        rules += [
            (rf"^modtok_{re.escape(m)}$", f"{keras_prefix}/modtok_{m}"),
            (
                rf"^embed_{re.escape(m)}\.proj\.weight$",
                f"{keras_prefix}/embed_{m}/proj/kernel",
            ),
            (
                rf"^embed_{re.escape(m)}\.proj\.bias$",
                f"{keras_prefix}/embed_{m}/proj/bias",
            ),
        ]

    return build_vit_mapper(
        keras_prefix, block_name_fn=lambda i: f"block{i}", extra_rules=rules
    )

keras_climate.weights.mappings.anysat_release_mapping.build_anysat_release_mapper 

build_anysat_release_mapper(modalities, depth, keras_prefix='anysat_release')
Source code in keras_climate/weights/mappings/anysat_release_mapping.py 
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
def build_anysat_release_mapper(modalities, depth, keras_prefix="anysat_release"):

    def mapper(torch_key):
        if torch_key == "cls_token":
            return f"{keras_prefix}/cls_token"

        m = re.match(r"^projector_(.+?)\.(.*)$", torch_key)
        if m:
            modality, rest = m.group(1), m.group(2)
            if modality not in modalities:
                return None
            mapped = _map_projector(modality, rest)
            if mapped is None:
                return None
            return f"{keras_prefix}/projector_{modality}/{mapped}"

        m = re.match(r"^spatial_encoder\.(.*)$", torch_key)
        if m:
            mapped = _map_local_encoder(m.group(1))
            if mapped is None:
                return None
            return f"{keras_prefix}/spatial_encoder/{mapped}"

        m = re.match(r"^blocks\.(\d+)\.(.*)$", torch_key)
        if m:
            idx, rest = int(m.group(1)), m.group(2)
            if idx == depth:
                mapped = _map_cross_block(rest)
                if mapped is None:
                    return None
                return f"{keras_prefix}/block_cross/{mapped}"
            mapped = _map_plain_block(rest)
            if mapped is None:
                return None
            return f"{keras_prefix}/block{idx}/{mapped}"

        return None

    return mapper

keras_climate.weights.mappings.autoformer_mapping.build_autoformer_mapper 

build_autoformer_mapper(encoder_layers, decoder_layers)
Source code in keras_climate/weights/mappings/autoformer_mapping.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def build_autoformer_mapper(encoder_layers, decoder_layers):
    rules = []
    rules += _embedding_rules("enc_embedding", "enc_embedding")
    rules += _embedding_rules("dec_embedding", "dec_embedding")

    for i in range(encoder_layers):
        rules += _encoder_layer_rules(i)
    for i in range(decoder_layers):
        rules += _decoder_layer_rules(i)

    rules += [
        (r"^projection\.weight$", "projection/kernel"),
        (r"^projection\.bias$", "projection/bias"),
    ]

    compiled = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, repl in compiled:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)
        return None

    return mapper

keras_climate.weights.mappings.clay_mapping.build_clay_identity_mapper 

build_clay_identity_mapper()
Source code in keras_climate/weights/mappings/clay_mapping.py 
64
65
def build_clay_identity_mapper():
    return lambda k: k

keras_climate.weights.mappings.clay_mapping.convert_clay_encoder_state_dict 

convert_clay_encoder_state_dict(flat_state_dict, depth=12, num_latent_tokens=128)
Source code in keras_climate/weights/mappings/clay_mapping.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
def convert_clay_encoder_state_dict(flat_state_dict, depth=12, num_latent_tokens=128):
    out = {}

    def copy(torch_key, keras_key):
        if torch_key in flat_state_dict:
            out[keras_key] = flat_state_dict[torch_key]

    copy("cls_token", "clay_encoder/cls_token")

    pe = "patch_embedding"
    copy(f"{pe}.fclayer.l1.weight", "clay_encoder/patch_embedding/fclayer/l1/kernel")
    copy(f"{pe}.fclayer.l1.bias", "clay_encoder/patch_embedding/fclayer/l1/bias")
    copy(f"{pe}.fclayer.l2.weight", "clay_encoder/patch_embedding/fclayer/l2/kernel")
    copy(f"{pe}.fclayer.l2.bias", "clay_encoder/patch_embedding/fclayer/l2/bias")

    wg = f"{pe}.weight_generator"
    kwg = "clay_encoder/patch_embedding/weight_generator"
    copy(f"{wg}.weight_tokens", f"{kwg}/weight_tokens")
    copy(f"{wg}.bias_token", f"{kwg}/bias_token")
    copy(f"{wg}.fc_weight.weight", f"{kwg}/fc_weight/kernel")
    copy(f"{wg}.fc_weight.bias", f"{kwg}/fc_weight/bias")
    copy(f"{wg}.fc_bias.weight", f"{kwg}/fc_bias/kernel")
    copy(f"{wg}.fc_bias.bias", f"{kwg}/fc_bias/bias")

    el = f"{wg}.encoder.layers.0"
    kel = f"{kwg}/encoder_layer"
    if f"{el}.self_attn.in_proj_weight" in flat_state_dict:
        out[f"{kel}/in_proj/kernel"] = flat_state_dict[f"{el}.self_attn.in_proj_weight"]
        out[f"{kel}/in_proj/bias"] = flat_state_dict[f"{el}.self_attn.in_proj_bias"]
    copy(f"{el}.self_attn.out_proj.weight", f"{kel}/out_proj/kernel")
    copy(f"{el}.self_attn.out_proj.bias", f"{kel}/out_proj/bias")
    copy(f"{el}.linear1.weight", f"{kel}/linear1/kernel")
    copy(f"{el}.linear1.bias", f"{kel}/linear1/bias")
    copy(f"{el}.linear2.weight", f"{kel}/linear2/kernel")
    copy(f"{el}.linear2.bias", f"{kel}/linear2/bias")
    copy(f"{el}.norm1.weight", f"{kel}/norm1/gamma")
    copy(f"{el}.norm1.bias", f"{kel}/norm1/beta")
    copy(f"{el}.norm2.weight", f"{kel}/norm2/gamma")
    copy(f"{el}.norm2.bias", f"{kel}/norm2/beta")

    for i in range(depth):
        tp, kp = f"transformer.layers.{i}", f"clay_encoder/block{i}"
        copy(f"{tp}.0.norm.weight", f"{kp}/attn_norm/gamma")
        copy(f"{tp}.0.norm.bias", f"{kp}/attn_norm/beta")
        copy(f"{tp}.0.to_qkv.weight", f"{kp}/to_qkv/kernel")
        copy(f"{tp}.0.to_out.weight", f"{kp}/to_out/kernel")
        copy(f"{tp}.1.net.0.weight", f"{kp}/ff_norm/gamma")
        copy(f"{tp}.1.net.0.bias", f"{kp}/ff_norm/beta")
        copy(f"{tp}.1.net.1.weight", f"{kp}/ff1/kernel")
        copy(f"{tp}.1.net.1.bias", f"{kp}/ff1/bias")
        copy(f"{tp}.1.net.3.weight", f"{kp}/ff2/kernel")
        copy(f"{tp}.1.net.3.bias", f"{kp}/ff2/bias")
    copy("transformer.norm.weight", "clay_encoder/norm/gamma")
    copy("transformer.norm.bias", "clay_encoder/norm/beta")

    return out

keras_climate.weights.mappings.climax_mapping.build_climax_mapper 

build_climax_mapper(num_vars, depth, decoder_depth)
Source code in keras_climate/weights/mappings/climax_mapping.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def build_climax_mapper(num_vars, depth, decoder_depth):
    rules = [
        (r"^pos_embed$", "pos_embed_layer/pos_embed"),
        (r"^channel_embed$", "var_embed_layer/var_embed"),
        (r"^channel_query$", "var_agg/var_query"),
        (r"^channel_agg\.in_proj_weight$", "var_agg/in_proj_weight"),
        (r"^channel_agg\.in_proj_bias$", "var_agg/in_proj_bias"),
        (r"^channel_agg\.out_proj\.weight$", "var_agg/out_proj/kernel"),
        (r"^channel_agg\.out_proj\.bias$", "var_agg/out_proj/bias"),
        (r"^lead_time_embed\.weight$", "lead_time_embed/kernel"),
        (r"^lead_time_embed\.bias$", "lead_time_embed/bias"),
        (r"^norm\.weight$", "norm/gamma"),
        (r"^norm\.bias$", "norm/beta"),
    ]
    for v in range(num_vars):
        rules += [
            (
                rf"^token_embeds\.{v}\.proj\.weight$",
                f"token_embeds/token_embeds{v}/kernel",
            ),
            (rf"^token_embeds\.{v}\.proj\.bias$", f"token_embeds/token_embeds{v}/bias"),
        ]
    for i in range(depth):
        rules += _block_rules(i)
    for i in range(decoder_depth + 1):
        idx = 2 * i
        rules += [
            (rf"^head\.{idx}\.weight$", f"head{idx}/kernel"),
            (rf"^head\.{idx}\.bias$", f"head{idx}/bias"),
        ]

    compiled = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, repl in compiled:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)
        return None

    return mapper

keras_climate.weights.mappings.convlstm_mapping.build_convlstm_identity_mapper 

build_convlstm_identity_mapper()
Source code in keras_climate/weights/mappings/convlstm_mapping.py 
31
32
def build_convlstm_identity_mapper():
    return lambda k: k

keras_climate.weights.mappings.convlstm_mapping.convert_convlstm_stack_state_dict 

convert_convlstm_stack_state_dict(flat_state_dict, layer_specs, keras_prefix='encoder_convlstm')
Source code in keras_climate/weights/mappings/convlstm_mapping.py 
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
def convert_convlstm_stack_state_dict(
    flat_state_dict, layer_specs, keras_prefix="encoder_convlstm"
):
    out = {}
    for i, (in_ch, hidden_ch) in enumerate(layer_specs):
        w = flat_state_dict[f"cell_list.{i}.conv.weight"]
        b = flat_state_dict[f"cell_list.{i}.conv.bias"]

        gi, gf, go, gg = np.split(w, 4, axis=0)
        w_reordered = np.concatenate([gi, gf, gg, go], axis=0)
        bi, bf, bo, bg = np.split(b, 4, axis=0)
        b_reordered = np.concatenate([bi, bf, bg, bo], axis=0)

        w_input = w_reordered[:, :in_ch]
        w_recurrent = w_reordered[:, in_ch : in_ch + hidden_ch]

        kernel = np.transpose(w_input, (2, 3, 1, 0))
        recurrent_kernel = np.transpose(w_recurrent, (2, 3, 1, 0))

        prefix = f"{keras_prefix}{i}/conv_lstm_cell"
        out[f"{prefix}/kernel"] = kernel
        out[f"{prefix}/recurrent_kernel"] = recurrent_kernel
        out[f"{prefix}/bias"] = b_reordered

    return out

keras_climate.weights.mappings.croma_mapping.build_croma_identity_mapper 

build_croma_identity_mapper()
Source code in keras_climate/weights/mappings/croma_mapping.py 
107
108
def build_croma_identity_mapper():
    return lambda k: k

keras_climate.weights.mappings.croma_mapping.convert_croma_state_dict 

convert_croma_state_dict(flat_state_dict, encoder_depth=12)
Source code in keras_climate/weights/mappings/croma_mapping.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def convert_croma_state_dict(flat_state_dict, encoder_depth=12):
    key_map = {}
    key_map["s1_encoder.linear_input.weight"] = "s1_encoder/linear_input/kernel"
    key_map["s1_encoder.linear_input.bias"] = "s1_encoder/linear_input/bias"
    key_map.update(_self_block_rules("s1_encoder", "s1_encoder", encoder_depth // 2))

    key_map["s2_encoder.linear_input.weight"] = "s2_encoder/linear_input/kernel"
    key_map["s2_encoder.linear_input.bias"] = "s2_encoder/linear_input/bias"
    key_map.update(_self_block_rules("s2_encoder", "s2_encoder", encoder_depth))

    key_map.update(_gap_ffn_rules("s1_GAP_FFN", "GAP_FFN_s1"))
    key_map.update(_gap_ffn_rules("s2_GAP_FFN", "GAP_FFN_s2"))

    key_map.update(
        _cross_block_rules("joint_encoder", "cross_encoder", encoder_depth // 2)
    )

    out = {}
    for torch_key, keras_key in key_map.items():
        if torch_key in flat_state_dict:
            out[keras_key] = flat_state_dict[torch_key]
    return out

keras_climate.weights.mappings.croma_mapping.load_croma_checkpoint 

load_croma_checkpoint(path)
Source code in keras_climate/weights/mappings/croma_mapping.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
def load_croma_checkpoint(path):
    import torch

    raw = torch.load(path, map_location="cpu")
    flat = {}
    for submodule, sub_sd in raw.items():
        for k, v in sub_sd.items():
            if k.endswith("num_batches_tracked"):
                continue
            flat[f"{submodule}.{k}"] = (
                v.detach().cpu().numpy() if hasattr(v, "detach") else np.asarray(v)
            )
    return flat

keras_climate.weights.mappings.deeplabv3plus_mapping.build_deeplabv3plus_mapper 

build_deeplabv3plus_mapper(backbone_keras_prefix='backbone', layer_counts=(3, 4, 6, 3))
Source code in keras_climate/weights/mappings/deeplabv3plus_mapping.py 
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
def build_deeplabv3plus_mapper(
    backbone_keras_prefix="backbone", layer_counts=(3, 4, 6, 3)
):
    rules = build_resnet_backbone_mapper(backbone_keras_prefix, layer_counts)

    rules += _conv_bn_rules("aspp.b0", "aspp/b0")
    for r in (6, 12, 18):
        rules += _conv_bn_rules(f"aspp.b{r}", f"aspp/b{r}")
    rules += _conv_bn_rules("aspp.pool_conv", "aspp/pool_conv")
    rules += _conv_bn_rules("aspp.project", "aspp/project")

    rules += _conv_bn_rules("low_level_project", "low_level_project")
    rules += _conv_bn_rules("decoder_conv1", "decoder_conv1")
    rules += _conv_bn_rules("decoder_conv2", "decoder_conv2")

    rules += [
        (r"^logits\.weight$", "logits/kernel"),
        (r"^logits\.bias$", "logits/bias"),
    ]

    compiled = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, repl in compiled:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)
        return None

    return mapper

keras_climate.weights.mappings.deeplabv3plus_mapping.build_resnet_backbone_mapper 

build_resnet_backbone_mapper(keras_prefix='backbone', layer_counts=(3, 4, 6, 3))
Source code in keras_climate/weights/mappings/deeplabv3plus_mapping.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
def build_resnet_backbone_mapper(keras_prefix="backbone", layer_counts=(3, 4, 6, 3)):
    rules = [
        (r"^conv1\.weight$", f"{keras_prefix}_stem/conv/kernel"),
        (r"^bn1\.weight$", f"{keras_prefix}_stem/bn/gamma"),
        (r"^bn1\.bias$", f"{keras_prefix}_stem/bn/beta"),
        (r"^bn1\.running_mean$", f"{keras_prefix}_stem/bn/moving_mean"),
        (r"^bn1\.running_var$", f"{keras_prefix}_stem/bn/moving_variance"),
    ]

    for layer_idx in range(1, 5):
        num_blocks = layer_counts[layer_idx - 1]
        stage_kp = f"{keras_prefix}_stage{layer_idx}"
        for i in range(num_blocks):
            tp = f"layer{layer_idx}.{i}"
            kp = f"{stage_kp}_block{i}"
            for c in (1, 2):
                rules += [
                    (
                        rf"^{re.escape(tp)}\.conv{c}\.weight$",
                        f"{kp}_conv{c}/conv/kernel",
                    ),
                    (rf"^{re.escape(tp)}\.bn{c}\.weight$", f"{kp}_conv{c}/bn/gamma"),
                    (rf"^{re.escape(tp)}\.bn{c}\.bias$", f"{kp}_conv{c}/bn/beta"),
                    (
                        rf"^{re.escape(tp)}\.bn{c}\.running_mean$",
                        f"{kp}_conv{c}/bn/moving_mean",
                    ),
                    (
                        rf"^{re.escape(tp)}\.bn{c}\.running_var$",
                        f"{kp}_conv{c}/bn/moving_variance",
                    ),
                ]
            rules += [
                (rf"^{re.escape(tp)}\.conv3\.weight$", f"{kp}_conv3/kernel"),
                (rf"^{re.escape(tp)}\.bn3\.weight$", f"{kp}_bn3/gamma"),
                (rf"^{re.escape(tp)}\.bn3\.bias$", f"{kp}_bn3/beta"),
                (rf"^{re.escape(tp)}\.bn3\.running_mean$", f"{kp}_bn3/moving_mean"),
                (rf"^{re.escape(tp)}\.bn3\.running_var$", f"{kp}_bn3/moving_variance"),
                (
                    rf"^{re.escape(tp)}\.downsample\.0\.weight$",
                    f"{kp}_downsample_conv/kernel",
                ),
                (
                    rf"^{re.escape(tp)}\.downsample\.1\.weight$",
                    f"{kp}_downsample_bn/gamma",
                ),
                (
                    rf"^{re.escape(tp)}\.downsample\.1\.bias$",
                    f"{kp}_downsample_bn/beta",
                ),
                (
                    rf"^{re.escape(tp)}\.downsample\.1\.running_mean$",
                    f"{kp}_downsample_bn/moving_mean",
                ),
                (
                    rf"^{re.escape(tp)}\.downsample\.1\.running_var$",
                    f"{kp}_downsample_bn/moving_variance",
                ),
            ]

    return rules

keras_climate.weights.mappings.dlinear_mapping.build_dlinear_mapper 

build_dlinear_mapper(num_channels=7, individual=False)
Source code in keras_climate/weights/mappings/dlinear_mapping.py 
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def build_dlinear_mapper(num_channels=7, individual=False):
    rules = []
    if individual:
        for c in range(num_channels):
            rules += [
                (rf"^Linear_Seasonal\.{c}\.weight$", f"linear_seasonal{c}/kernel"),
                (rf"^Linear_Seasonal\.{c}\.bias$", f"linear_seasonal{c}/bias"),
                (rf"^Linear_Trend\.{c}\.weight$", f"linear_trend{c}/kernel"),
                (rf"^Linear_Trend\.{c}\.bias$", f"linear_trend{c}/bias"),
            ]
    else:
        rules += [
            (r"^Linear_Seasonal\.weight$", "linear_seasonal/kernel"),
            (r"^Linear_Seasonal\.bias$", "linear_seasonal/bias"),
            (r"^Linear_Trend\.weight$", "linear_trend/kernel"),
            (r"^Linear_Trend\.bias$", "linear_trend/bias"),
        ]

    compiled = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, repl in compiled:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)
        return None

    return mapper

keras_climate.weights.mappings.earthformer_mapping.build_earthformer_mapper 

build_earthformer_mapper(num_enc_stages, num_dec_stages, enc_depths=None, dec_depths=None)
Source code in keras_climate/weights/mappings/earthformer_mapping.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
def build_earthformer_mapper(
    num_enc_stages, num_dec_stages, enc_depths=None, dec_depths=None
):
    enc_depths = enc_depths or [1] * num_enc_stages
    dec_depths = dec_depths or [1] * num_dec_stages

    rules = [
        (r"^stem\.conv\.weight$", "stem/conv/kernel"),
        (r"^stem\.conv\.bias$", "stem/conv/bias"),
        (r"^time_channel_proj\.weight$", "time_channel_proj/kernel"),
        (r"^time_channel_proj\.bias$", "time_channel_proj/bias"),
    ]
    for i in range(num_enc_stages - 1):
        rules += [
            (rf"^downsample{i}\.conv\.weight$", f"downsample{i}/conv/kernel"),
            (rf"^downsample{i}\.conv\.bias$", f"downsample{i}/conv/bias"),
        ]
    for i in range(num_dec_stages):
        rules += [
            (rf"^upsample{i}\.conv\.weight$", f"upsample{i}/conv/kernel"),
            (rf"^upsample{i}\.conv\.bias$", f"upsample{i}/conv/bias"),
            (rf"^skip_proj{i}\.conv\.weight$", f"skip_proj{i}/conv/kernel"),
            (rf"^skip_proj{i}\.conv\.bias$", f"skip_proj{i}/conv/bias"),
        ]

    for i in range(num_enc_stages):
        rules += _block_rules(f"enc_stage{i}", f"enc_stage{i}", enc_depths[i])
    for i in range(num_dec_stages):
        rules += _block_rules(f"dec_stage{i}", f"dec_stage{i}", dec_depths[i])

    compiled = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, repl in compiled:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)
        return None

    return mapper

keras_climate.weights.mappings.fourcastnet_mapping.build_fourcastnet_mapper 

build_fourcastnet_mapper(depth)
Source code in keras_climate/weights/mappings/fourcastnet_mapping.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
def build_fourcastnet_mapper(depth):
    afno_mapper = build_afno_mapper(depth)
    rules = [
        (r"^pos_embed$", "pos_embed_layer/pos_embed"),
        (r"^norm\.weight$", "norm/gamma"),
        (r"^norm\.bias$", "norm/beta"),
    ]
    compiled = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, repl in compiled:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)
        return afno_mapper(torch_key)

    return mapper

keras_climate.weights.mappings.fourcastnet_mapping.convert_fourcastnet_state_dict 

convert_fourcastnet_state_dict(flat_state_dict, grid_h, grid_w)
Source code in keras_climate/weights/mappings/fourcastnet_mapping.py 
13
14
15
16
17
18
19
20
def convert_fourcastnet_state_dict(flat_state_dict, grid_h, grid_w):
    out = {}
    for k, v in flat_state_dict.items():
        k = k[len("module.") :] if k.startswith("module.") else k
        if k == "pos_embed":
            v = v.reshape(1, grid_h, grid_w, v.shape[-1])
        out[k] = v
    return out

keras_climate.weights.mappings.fourcastnet_mapping.load_fourcastnet_checkpoint 

load_fourcastnet_checkpoint(path)
Source code in keras_climate/weights/mappings/fourcastnet_mapping.py 
 6
 7
 8
 9
10
def load_fourcastnet_checkpoint(path):
    import torch

    ckpt = torch.load(path, map_location="cpu", weights_only=False)
    return {k: v.detach().numpy() for k, v in ckpt["model_state"].items()}

keras_climate.weights.mappings.informer_mapping.build_informer_mapper 

build_informer_mapper(encoder_layers, decoder_layers)
Source code in keras_climate/weights/mappings/informer_mapping.py 
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
def build_informer_mapper(encoder_layers, decoder_layers):
    rules = []
    rules += _embedding_rules("enc_embedding", "enc_embedding")
    rules += _embedding_rules("dec_embedding", "dec_embedding")

    for i in range(encoder_layers):
        rules += _encoder_layer_rules(i)
        if i < encoder_layers - 1:
            rules += [
                (rf"^distill{i}\.conv\.weight$", f"distill{i}/conv/kernel"),
                (rf"^distill{i}\.bn\.weight$", f"distill{i}/bn/gamma"),
                (rf"^distill{i}\.bn\.bias$", f"distill{i}/bn/beta"),
                (rf"^distill{i}\.bn\.running_mean$", f"distill{i}/bn/moving_mean"),
                (rf"^distill{i}\.bn\.running_var$", f"distill{i}/bn/moving_variance"),
            ]

    for i in range(decoder_layers):
        rules += _decoder_layer_rules(i)

    rules += [
        (r"^projection\.weight$", "projection/kernel"),
        (r"^projection\.bias$", "projection/bias"),
    ]

    compiled = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, repl in compiled:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)
        return None

    return mapper

keras_climate.weights.mappings.metnet_mapping.build_metnet_mapper 

build_metnet_mapper()
Source code in keras_climate/weights/mappings/metnet_mapping.py 
79
80
def build_metnet_mapper():
    return lambda k: k

keras_climate.weights.mappings.metnet_mapping.convert_metnet_state_dict 

convert_metnet_state_dict(flat_state_dict, base_filters, num_dilated_convs=6, num_att_layers=4, num_upsample=2)
Source code in keras_climate/weights/mappings/metnet_mapping.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
def convert_metnet_state_dict(
    flat_state_dict, base_filters, num_dilated_convs=6, num_att_layers=4, num_upsample=2
):
    out = {}

    def copy(torch_key, keras_key):
        if torch_key in flat_state_dict:
            out[keras_key] = flat_state_dict[torch_key]

    copy("stem.conv.weight", "stem/conv/kernel")
    copy("stem.conv.bias", "stem/conv/bias")
    copy("stem_bn.bn.weight", "stem_bn/bn/gamma")
    copy("stem_bn.bn.bias", "stem_bn/bn/beta")
    copy("stem_bn.bn.running_mean", "stem_bn/bn/moving_mean")
    copy("stem_bn.bn.running_var", "stem_bn/bn/moving_variance")

    temporal_flat = {
        k[len("temporal_encoder.") :]: v
        for k, v in flat_state_dict.items()
        if k.startswith("temporal_encoder.")
    }
    if temporal_flat:
        temporal_converted = convert_convlstm_stack_state_dict(
            temporal_flat,
            [(base_filters, base_filters)],
            keras_prefix="temporal_encoder",
        )
        for k, v in temporal_converted.items():
            out[k.replace("temporal_encoder0", "temporal_encoder", 1)] = v

    for i in range(num_dilated_convs):
        copy(f"context_tower.{i}.conv.weight", f"context_tower_dconv{i}/kernel")
        copy(f"context_tower.{i}.conv.bias", f"context_tower_dconv{i}/bias")
        copy(f"context_tower.{i}.bn.weight", f"context_tower_bn{i}/gamma")
        copy(f"context_tower.{i}.bn.bias", f"context_tower_bn{i}/beta")
        copy(f"context_tower.{i}.bn.running_mean", f"context_tower_bn{i}/moving_mean")
        copy(
            f"context_tower.{i}.bn.running_var", f"context_tower_bn{i}/moving_variance"
        )

    copy("lead_time_embed.weight", "lead_time_embed/embeddings")
    copy("attn_proj_in.weight", "attn_proj_in/kernel")
    copy("attn_proj_in.bias", "attn_proj_in/bias")

    for i in range(num_att_layers):
        p, kp = f"axial_attn{i}", f"axial_attn{i}"
        for stream in ("row", "col"):
            copy(f"{p}.{stream}_qkv.weight", f"{kp}/{stream}_qkv/kernel")
            copy(f"{p}.{stream}_qkv.bias", f"{kp}/{stream}_qkv/bias")
            copy(f"{p}.{stream}_proj.weight", f"{kp}/{stream}_proj/kernel")
            copy(f"{p}.{stream}_proj.bias", f"{kp}/{stream}_proj/bias")
        for n in ("norm1", "norm2"):
            copy(f"{p}.{n}.weight", f"{kp}/{n}/gamma")
            copy(f"{p}.{n}.bias", f"{kp}/{n}/beta")

    for i in range(num_upsample):
        copy(f"upsample{i}.weight", f"upsample{i}/kernel")
        copy(f"upsample{i}.bias", f"upsample{i}/bias")
        copy(f"upsample_bn{i}.weight", f"upsample_bn{i}/gamma")
        copy(f"upsample_bn{i}.bias", f"upsample_bn{i}/beta")
        copy(f"upsample_bn{i}.running_mean", f"upsample_bn{i}/moving_mean")
        copy(f"upsample_bn{i}.running_var", f"upsample_bn{i}/moving_variance")

    copy("precip_head.weight", "precip_head/kernel")
    copy("precip_head.bias", "precip_head/bias")

    return out

keras_climate.weights.mappings.nbeats_mapping.build_nbeats_mapper 

build_nbeats_mapper(num_blocks, num_fc_layers=4)
Source code in keras_climate/weights/mappings/nbeats_mapping.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
def build_nbeats_mapper(num_blocks, num_fc_layers=4):
    rules = []
    for i in range(num_blocks):
        rules += _block_rules(i, num_fc_layers)

    compiled = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, repl in compiled:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)
        return None

    return mapper

keras_climate.weights.mappings.pangu_weather_mapping.build_pangu_weather_mapper 

build_pangu_weather_mapper(depths=(2, 6, 6, 2))
Source code in keras_climate/weights/mappings/pangu_weather_mapping.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def build_pangu_weather_mapper(depths=(2, 6, 6, 2)):
    rules = [
        (r"^_input_layer\.conv\.weight$", "_input_layer/conv/kernel"),
        (r"^_input_layer\.conv\.bias$", "_input_layer/conv/bias"),
        (r"^_input_layer\.conv_surface\.weight$", "_input_layer/conv_surface/kernel"),
        (r"^_input_layer\.conv_surface\.bias$", "_input_layer/conv_surface/bias"),
        (r"^downsample\.linear\.weight$", "downsample/linear/kernel"),
        (r"^downsample\.norm\.weight$", "downsample/norm/gamma"),
        (r"^downsample\.norm\.bias$", "downsample/norm/beta"),
        (r"^upsample\.linear1\.weight$", "upsample/linear1/kernel"),
        (r"^upsample\.linear2\.weight$", "upsample/linear2/kernel"),
        (r"^upsample\.norm\.weight$", "upsample/norm/gamma"),
        (r"^upsample\.norm\.bias$", "upsample/norm/beta"),
        (r"^_output_layer\.conv\.weight$", "_output_layer/conv/kernel"),
        (r"^_output_layer\.conv\.bias$", "_output_layer/conv/bias"),
        (r"^_output_layer\.conv_surface\.weight$", "_output_layer/conv_surface/kernel"),
        (r"^_output_layer\.conv_surface\.bias$", "_output_layer/conv_surface/bias"),
    ]
    for i, depth in enumerate(depths):
        for j in range(depth):
            rules += _block_rules(i, j)

    compiled = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, repl in compiled:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)
        return None

    return mapper

keras_climate.weights.mappings.pangu_weather_mapping.convert_pangu_weather_state_dict 

convert_pangu_weather_state_dict(flat_state_dict)
Source code in keras_climate/weights/mappings/pangu_weather_mapping.py 
 4
 5
 6
 7
 8
 9
10
11
12
13
14
def convert_pangu_weather_state_dict(flat_state_dict):
    out = {}
    for k, v in flat_state_dict.items():
        if (
            k.endswith((".conv.weight", ".conv_surface.weight"))
            and v.ndim == 3
            and v.shape[-1] == 1
        ):
            v = v[:, :, 0]
        out[k] = v
    return out

keras_climate.weights.mappings.patchtst_mapping.build_patchtst_mapper 

build_patchtst_mapper(depth, use_revin=True)
Source code in keras_climate/weights/mappings/patchtst_mapping.py 
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def build_patchtst_mapper(depth, use_revin=True):
    rules = [
        (r"^patch_proj\.weight$", "patch_proj/kernel"),
        (r"^patch_proj\.bias$", "patch_proj/bias"),
        (r"^pos_embed$", "pos_embed/embeddings"),
        (r"^encoder_norm\.weight$", "encoder_norm/gamma"),
        (r"^encoder_norm\.bias$", "encoder_norm/beta"),
        (r"^forecast_head\.weight$", "forecast_head/kernel"),
        (r"^forecast_head\.bias$", "forecast_head/bias"),
    ]
    if use_revin:
        rules += [
            (r"^revin\.affine_weight$", "revin/gamma"),
            (r"^revin\.affine_bias$", "revin/beta"),
        ]

    for i in range(depth):
        tp, kp = f"blocks.{i}", f"encoder_block{i}"
        rules += [
            (rf"^{tp}\.norm1\.weight$", f"{kp}/norm1/gamma"),
            (rf"^{tp}\.norm1\.bias$", f"{kp}/norm1/beta"),
            (rf"^{tp}\.attn\.qkv\.weight$", f"{kp}/attn/qkv/kernel"),
            (rf"^{tp}\.attn\.qkv\.bias$", f"{kp}/attn/qkv/bias"),
            (rf"^{tp}\.attn\.proj\.weight$", f"{kp}/attn/proj/kernel"),
            (rf"^{tp}\.attn\.proj\.bias$", f"{kp}/attn/proj/bias"),
            (rf"^{tp}\.norm2\.weight$", f"{kp}/norm2/gamma"),
            (rf"^{tp}\.norm2\.bias$", f"{kp}/norm2/beta"),
            (rf"^{tp}\.mlp\.fc1\.weight$", f"{kp}/mlp/fc1/kernel"),
            (rf"^{tp}\.mlp\.fc1\.bias$", f"{kp}/mlp/fc1/bias"),
            (rf"^{tp}\.mlp\.fc2\.weight$", f"{kp}/mlp/fc2/kernel"),
            (rf"^{tp}\.mlp\.fc2\.bias$", f"{kp}/mlp/fc2/bias"),
        ]

    compiled = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, repl in compiled:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)
        return None

    return mapper

keras_climate.weights.mappings.ringmo_mapping.build_ringmo_decoder_mapper 

build_ringmo_decoder_mapper(keras_prefix='ringmo_decoder')
Source code in keras_climate/weights/mappings/ringmo_mapping.py 
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def build_ringmo_decoder_mapper(keras_prefix="ringmo_decoder"):
    rules = [
        (r"^decoder\.weight$", f"{keras_prefix}/conv/kernel"),
        (r"^decoder\.bias$", f"{keras_prefix}/conv/bias"),
    ]
    compiled = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, repl in compiled:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)
        return None

    return mapper

keras_climate.weights.mappings.ringmo_mapping.build_ringmo_encoder_mapper 

build_ringmo_encoder_mapper(depths, keras_prefix='ringmo_encoder')
Source code in keras_climate/weights/mappings/ringmo_mapping.py 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
def build_ringmo_encoder_mapper(depths, keras_prefix="ringmo_encoder"):
    rules = []
    rules += _conv_bn_act_rules(
        "patch_embed.conv1", f"{keras_prefix}/patch_embed/conv1"
    )
    rules += _conv_bn_act_rules(
        "patch_embed.conv2", f"{keras_prefix}/patch_embed/conv2"
    )
    rules += [
        (r"^patch_embed\.conv3\.weight$", f"{keras_prefix}/patch_embed/conv3/kernel"),
        (r"^patch_embed\.conv3\.bias$", f"{keras_prefix}/patch_embed/conv3/bias"),
        (r"^patch_embed\.norm\.weight$", f"{keras_prefix}/patch_embed/norm/gamma"),
        (r"^patch_embed\.norm\.bias$", f"{keras_prefix}/patch_embed/norm/beta"),
        (r"^norm\.weight$", f"{keras_prefix}/norm/gamma"),
        (r"^norm\.bias$", f"{keras_prefix}/norm/beta"),
    ]

    for i, depth in enumerate(depths):
        for j in range(depth):
            rules += _swin_block_rules(
                f"stages.{i}.blocks.{j}", f"{keras_prefix}/stage{i}/block{j}"
            )
        if i < len(depths) - 1:
            rules += [
                (
                    rf"^stages\.{i}\.downsample\.reduction\.weight$",
                    f"{keras_prefix}/stage{i}/downsample/reduction/kernel",
                ),
                (
                    rf"^stages\.{i}\.downsample\.norm\.weight$",
                    f"{keras_prefix}/stage{i}/downsample/norm/gamma",
                ),
                (
                    rf"^stages\.{i}\.downsample\.norm\.bias$",
                    f"{keras_prefix}/stage{i}/downsample/norm/beta",
                ),
            ]

    compiled = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, repl in compiled:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)
        return None

    return mapper

keras_climate.weights.mappings.ringmo_mapping.build_ringmo_mapper 

build_ringmo_mapper(depths, encoder_prefix='ringmo_encoder', decoder_prefix='ringmo_decoder')
Source code in keras_climate/weights/mappings/ringmo_mapping.py 
102
103
104
105
106
107
108
109
110
111
112
113
114
def build_ringmo_mapper(
    depths, encoder_prefix="ringmo_encoder", decoder_prefix="ringmo_decoder"
):
    encoder_mapper = build_ringmo_encoder_mapper(depths, encoder_prefix)
    decoder_mapper = build_ringmo_decoder_mapper(decoder_prefix)

    def mapper(torch_key):
        result = encoder_mapper(torch_key)
        if result is not None:
            return result
        return decoder_mapper(torch_key)

    return mapper

keras_climate.weights.mappings.satclip_mapping.SATCLIP_SKIP_PATTERNS module-attribute 

SATCLIP_SKIP_PATTERNS = ['^model\\.visual\\.', '^model\\.logit_scale$', '^model\\.posenc', '^model\\.location\\.posenc', '^model\\.nnet\\.']

keras_climate.weights.mappings.satclip_mapping.build_satclip_location_mapper 

build_satclip_location_mapper(num_hidden_layers=2)
Source code in keras_climate/weights/mappings/satclip_mapping.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def build_satclip_location_mapper(num_hidden_layers=2):
    rules = []
    for i in range(num_hidden_layers):
        rules += [
            (rf"^model\.location\.nnet\.layers\.{i}\.weight$", f"layers{i}/kernel"),
            (rf"^model\.location\.nnet\.layers\.{i}\.bias$", f"layers{i}/bias"),
        ]
    rules += [
        (r"^model\.location\.nnet\.last_layer\.weight$", "last_layer/kernel"),
        (r"^model\.location\.nnet\.last_layer\.bias$", "last_layer/bias"),
    ]
    compiled = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, repl in compiled:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)
        return None

    return mapper

keras_climate.weights.mappings.satclip_mapping.load_satclip_checkpoint 

load_satclip_checkpoint(path)
Source code in keras_climate/weights/mappings/satclip_mapping.py 
4
5
6
7
8
def load_satclip_checkpoint(path):
    import torch

    ckpt = torch.load(path, map_location="cpu", weights_only=False)
    return {k: v.detach().numpy() for k, v in ckpt["state_dict"].items()}

keras_climate.weights.mappings.satmae_mapping.build_satmae_decoder_mapper 

build_satmae_decoder_mapper(keras_prefix='satmae_decoder', torch_prefix='')
Source code in keras_climate/weights/mappings/satmae_mapping.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
def build_satmae_decoder_mapper(keras_prefix="satmae_decoder", torch_prefix=""):
    tp = re.escape(torch_prefix)
    rules = [
        (rf"^{tp}decoder_embed\.weight$", f"{keras_prefix}/decoder_embed/kernel"),
        (rf"^{tp}decoder_embed\.bias$", f"{keras_prefix}/decoder_embed/bias"),
        (rf"^{tp}mask_token$", f"{keras_prefix}/mask_token"),
        (rf"^{tp}decoder_pos_embed$", f"{keras_prefix}/decoder_pos_embed"),
        (rf"^{tp}decoder_norm\.weight$", f"{keras_prefix}/norm/gamma"),
        (rf"^{tp}decoder_norm\.bias$", f"{keras_prefix}/norm/beta"),
        (rf"^{tp}decoder_pred\.weight$", f"{keras_prefix}/pred/kernel"),
        (rf"^{tp}decoder_pred\.bias$", f"{keras_prefix}/pred/bias"),
    ]

    block_pattern = re.compile(
        rf"^{tp}decoder_blocks\.(\d+)\.(norm1|norm2)\.(weight|bias)$|"
        rf"^{tp}decoder_blocks\.(\d+)\.attn\.(qkv|proj)\.(weight|bias)$|"
        rf"^{tp}decoder_blocks\.(\d+)\.mlp\.(fc1|fc2)\.(weight|bias)$"
    )
    compiled_static = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, repl in compiled_static:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)

        m = block_pattern.match(torch_key)
        if not m:
            return None

        if m.group(1) is not None:
            idx, sub, param = m.group(1), m.group(2), m.group(3)
            param_name = "gamma" if param == "weight" else "beta"
            return f"{keras_prefix}/block{int(idx)}/{sub}/{param_name}"

        if m.group(4) is not None:
            idx, sub, param = m.group(4), m.group(5), m.group(6)
            keras_param = "kernel" if param == "weight" else "bias"
            return f"{keras_prefix}/block{int(idx)}/attn/{sub}/{keras_param}"

        idx, sub, param = m.group(7), m.group(8), m.group(9)
        keras_param = "kernel" if param == "weight" else "bias"
        return f"{keras_prefix}/block{int(idx)}/mlp/{sub}/{keras_param}"

    return mapper

keras_climate.weights.mappings.satmae_mapping.build_satmae_encoder_mapper 

build_satmae_encoder_mapper(keras_prefix='satmae_encoder', mode='single')
Source code in keras_climate/weights/mappings/satmae_mapping.py 
11
12
13
14
15
def build_satmae_encoder_mapper(keras_prefix="satmae_encoder", mode="single"):
    extra_rules = [
        (p, r.format(prefix=keras_prefix)) for p, r in EXTRA_EMBED_KEYS.get(mode, [])
    ]
    return build_vit_mapper(keras_prefix, extra_rules=extra_rules)

keras_climate.weights.mappings.satmae_mapping.build_satmae_mapper 

build_satmae_mapper(encoder_prefix='satmae_encoder', decoder_prefix='satmae_decoder', mode='single')
Source code in keras_climate/weights/mappings/satmae_mapping.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
def build_satmae_mapper(
    encoder_prefix="satmae_encoder", decoder_prefix="satmae_decoder", mode="single"
):
    encoder_mapper = build_satmae_encoder_mapper(encoder_prefix, mode)
    decoder_mapper = build_satmae_decoder_mapper(decoder_prefix)

    def mapper(torch_key):
        result = encoder_mapper(torch_key)
        if result is not None:
            return result
        return decoder_mapper(torch_key)

    return mapper

keras_climate.weights.mappings.scalemae_mapping.SCALEMAE_SKIP_PATTERNS module-attribute 

SCALEMAE_SKIP_PATTERNS = ['^pos_embed$', '^decoder_pos_embed$']

keras_climate.weights.mappings.scalemae_mapping.build_scalemae_mapper 

build_scalemae_mapper(keras_prefix='scalemae_encoder')
Source code in keras_climate/weights/mappings/scalemae_mapping.py 
4
5
def build_scalemae_mapper(keras_prefix="scalemae_encoder"):
    return build_vit_mapper(keras_prefix)

keras_climate.weights.mappings.segformer_mapping.build_segformer_mapper 

build_segformer_mapper(cfg, keras_name='segformer', decode_head_prefix='decode_head')
Source code in keras_climate/weights/mappings/segformer_mapping.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def build_segformer_mapper(
    cfg, keras_name="segformer", decode_head_prefix="decode_head"
):
    rules = []

    for s in range(1, 5):
        stage_kp = f"{keras_name}_backbone_stage{s}"
        rules += [
            (rf"^patch_embed{s}\.proj\.weight$", f"{stage_kp}/patch_embed/proj/kernel"),
            (rf"^patch_embed{s}\.proj\.bias$", f"{stage_kp}/patch_embed/proj/bias"),
            (rf"^patch_embed{s}\.norm\.weight$", f"{stage_kp}/patch_embed/norm/gamma"),
            (rf"^patch_embed{s}\.norm\.bias$", f"{stage_kp}/patch_embed/norm/beta"),
            (rf"^norm{s}\.weight$", f"{stage_kp}/norm/gamma"),
            (rf"^norm{s}\.bias$", f"{stage_kp}/norm/beta"),
        ]

        depth = cfg["depths"][s - 1]
        for i in range(depth):
            tp = f"block{s}.{i}"
            kp = f"{stage_kp}/block{i}"
            rules += [
                (rf"^{re.escape(tp)}\.norm1\.weight$", f"{kp}/norm1/gamma"),
                (rf"^{re.escape(tp)}\.norm1\.bias$", f"{kp}/norm1/beta"),
                (rf"^{re.escape(tp)}\.attn\.q\.weight$", f"{kp}/attn/q/kernel"),
                (rf"^{re.escape(tp)}\.attn\.q\.bias$", f"{kp}/attn/q/bias"),
                (rf"^{re.escape(tp)}\.attn\.kv\.weight$", f"{kp}/attn/kv/kernel"),
                (rf"^{re.escape(tp)}\.attn\.kv\.bias$", f"{kp}/attn/kv/bias"),
                (rf"^{re.escape(tp)}\.attn\.sr\.weight$", f"{kp}/attn/sr/kernel"),
                (rf"^{re.escape(tp)}\.attn\.sr\.bias$", f"{kp}/attn/sr/bias"),
                (rf"^{re.escape(tp)}\.attn\.norm\.weight$", f"{kp}/attn/sr_norm/gamma"),
                (rf"^{re.escape(tp)}\.attn\.norm\.bias$", f"{kp}/attn/sr_norm/beta"),
                (rf"^{re.escape(tp)}\.attn\.proj\.weight$", f"{kp}/attn/proj/kernel"),
                (rf"^{re.escape(tp)}\.attn\.proj\.bias$", f"{kp}/attn/proj/bias"),
                (rf"^{re.escape(tp)}\.norm2\.weight$", f"{kp}/norm2/gamma"),
                (rf"^{re.escape(tp)}\.norm2\.bias$", f"{kp}/norm2/beta"),
                (rf"^{re.escape(tp)}\.mlp\.fc1\.weight$", f"{kp}/mixffn_fc1/kernel"),
                (rf"^{re.escape(tp)}\.mlp\.fc1\.bias$", f"{kp}/mixffn_fc1/bias"),
                (
                    rf"^{re.escape(tp)}\.mlp\.dwconv\.dwconv\.weight$",
                    f"{kp}/mixffn_dwconv/kernel",
                ),
                (
                    rf"^{re.escape(tp)}\.mlp\.dwconv\.dwconv\.bias$",
                    f"{kp}/mixffn_dwconv/bias",
                ),
                (rf"^{re.escape(tp)}\.mlp\.fc2\.weight$", f"{kp}/mixffn_fc2/kernel"),
                (rf"^{re.escape(tp)}\.mlp\.fc2\.bias$", f"{kp}/mixffn_fc2/bias"),
            ]

    dh = decode_head_prefix
    dh_kp = f"{keras_name}_decode_head"
    for i in (4, 3, 2, 1):
        rules += [
            (
                rf"^{re.escape(dh)}\.linear_c{i}\.proj\.weight$",
                f"{dh_kp}_linear_c{i}/kernel",
            ),
            (
                rf"^{re.escape(dh)}\.linear_c{i}\.proj\.bias$",
                f"{dh_kp}_linear_c{i}/bias",
            ),
        ]
    rules += [
        (rf"^{re.escape(dh)}\.linear_fuse\.conv\.weight$", f"{dh_kp}_fuse/conv/kernel"),
        (rf"^{re.escape(dh)}\.linear_fuse\.bn\.weight$", f"{dh_kp}_fuse/bn/gamma"),
        (rf"^{re.escape(dh)}\.linear_fuse\.bn\.bias$", f"{dh_kp}_fuse/bn/beta"),
        (
            rf"^{re.escape(dh)}\.linear_fuse\.bn\.running_mean$",
            f"{dh_kp}_fuse/bn/moving_mean",
        ),
        (
            rf"^{re.escape(dh)}\.linear_fuse\.bn\.running_var$",
            f"{dh_kp}_fuse/bn/moving_variance",
        ),
        (rf"^{re.escape(dh)}\.linear_pred\.weight$", f"{dh_kp}_classifier/kernel"),
        (rf"^{re.escape(dh)}\.linear_pred\.bias$", f"{dh_kp}_classifier/bias"),
    ]

    compiled = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, repl in compiled:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)
        return None

    return mapper

keras_climate.weights.mappings.segformer_mapping.build_segformer_param_kind_map 

build_segformer_param_kind_map(cfg, keras_name='segformer')
Source code in keras_climate/weights/mappings/segformer_mapping.py 
197
198
199
200
201
202
203
204
205
def build_segformer_param_kind_map(cfg, keras_name="segformer"):
    out = {}
    for s in range(1, 5):
        depth = cfg["depths"][s - 1]
        for i in range(depth):
            out[f"{keras_name}_backbone_stage{s}/block{i}/mixffn_dwconv/kernel"] = (
                "depthwise_kernel"
            )
    return out

keras_climate.weights.mappings.segformer_mapping.convert_hf_segformer_state_dict 

convert_hf_segformer_state_dict(hf_state_dict, cfg)
Source code in keras_climate/weights/mappings/segformer_mapping.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
def convert_hf_segformer_state_dict(hf_state_dict, cfg):
    out = {}
    depths = cfg["depths"]

    for s in range(4):
        prefix = f"segformer.stages.{s}"
        out[f"patch_embed{s+1}.proj.weight"] = hf_state_dict[
            f"{prefix}.patch_embeddings.proj.weight"
        ]
        out[f"patch_embed{s+1}.proj.bias"] = hf_state_dict[
            f"{prefix}.patch_embeddings.proj.bias"
        ]
        out[f"patch_embed{s+1}.norm.weight"] = hf_state_dict[
            f"{prefix}.patch_embeddings.layer_norm.weight"
        ]
        out[f"patch_embed{s+1}.norm.bias"] = hf_state_dict[
            f"{prefix}.patch_embeddings.layer_norm.bias"
        ]
        out[f"norm{s+1}.weight"] = hf_state_dict[f"{prefix}.layer_norm.weight"]
        out[f"norm{s+1}.bias"] = hf_state_dict[f"{prefix}.layer_norm.bias"]

        for i in range(depths[s]):
            bp = f"{prefix}.blocks.{i}"
            tp = f"block{s+1}.{i}"
            out[f"{tp}.norm1.weight"] = hf_state_dict[f"{bp}.layernorm_before.weight"]
            out[f"{tp}.norm1.bias"] = hf_state_dict[f"{bp}.layernorm_before.bias"]
            out[f"{tp}.attn.q.weight"] = hf_state_dict[f"{bp}.attention.q_proj.weight"]
            out[f"{tp}.attn.q.bias"] = hf_state_dict[f"{bp}.attention.q_proj.bias"]
            out[f"{tp}.attn.kv.weight"] = np.concatenate(
                [
                    hf_state_dict[f"{bp}.attention.k_proj.weight"],
                    hf_state_dict[f"{bp}.attention.v_proj.weight"],
                ],
                axis=0,
            )
            out[f"{tp}.attn.kv.bias"] = np.concatenate(
                [
                    hf_state_dict[f"{bp}.attention.k_proj.bias"],
                    hf_state_dict[f"{bp}.attention.v_proj.bias"],
                ],
                axis=0,
            )
            out[f"{tp}.attn.proj.weight"] = hf_state_dict[
                f"{bp}.attention.o_proj.weight"
            ]
            out[f"{tp}.attn.proj.bias"] = hf_state_dict[f"{bp}.attention.o_proj.bias"]
            sr_key = f"{bp}.attention.sequence_reduction.sequence_reduction.weight"
            if sr_key in hf_state_dict:
                out[f"{tp}.attn.sr.weight"] = hf_state_dict[sr_key]
                out[f"{tp}.attn.sr.bias"] = hf_state_dict[
                    f"{bp}.attention.sequence_reduction.sequence_reduction.bias"
                ]
                out[f"{tp}.attn.norm.weight"] = hf_state_dict[
                    f"{bp}.attention.sequence_reduction.layer_norm.weight"
                ]
                out[f"{tp}.attn.norm.bias"] = hf_state_dict[
                    f"{bp}.attention.sequence_reduction.layer_norm.bias"
                ]
            out[f"{tp}.norm2.weight"] = hf_state_dict[f"{bp}.layernorm_after.weight"]
            out[f"{tp}.norm2.bias"] = hf_state_dict[f"{bp}.layernorm_after.bias"]
            out[f"{tp}.mlp.fc1.weight"] = hf_state_dict[f"{bp}.mlp.fc1.weight"]
            out[f"{tp}.mlp.fc1.bias"] = hf_state_dict[f"{bp}.mlp.fc1.bias"]
            out[f"{tp}.mlp.dwconv.dwconv.weight"] = hf_state_dict[
                f"{bp}.mlp.dwconv.dwconv.weight"
            ]
            out[f"{tp}.mlp.dwconv.dwconv.bias"] = hf_state_dict[
                f"{bp}.mlp.dwconv.dwconv.bias"
            ]
            out[f"{tp}.mlp.fc2.weight"] = hf_state_dict[f"{bp}.mlp.fc2.weight"]
            out[f"{tp}.mlp.fc2.bias"] = hf_state_dict[f"{bp}.mlp.fc2.bias"]

    for i in (1, 2, 3, 4):
        out[f"decode_head.linear_c{i}.proj.weight"] = hf_state_dict[
            f"decode_head.linear_projections.{i-1}.proj.weight"
        ]
        out[f"decode_head.linear_c{i}.proj.bias"] = hf_state_dict[
            f"decode_head.linear_projections.{i-1}.proj.bias"
        ]

    out["decode_head.linear_fuse.conv.weight"] = hf_state_dict[
        "decode_head.linear_fuse.weight"
    ]
    out["decode_head.linear_fuse.bn.weight"] = hf_state_dict[
        "decode_head.batch_norm.weight"
    ]
    out["decode_head.linear_fuse.bn.bias"] = hf_state_dict[
        "decode_head.batch_norm.bias"
    ]
    out["decode_head.linear_fuse.bn.running_mean"] = hf_state_dict[
        "decode_head.batch_norm.running_mean"
    ]
    out["decode_head.linear_fuse.bn.running_var"] = hf_state_dict[
        "decode_head.batch_norm.running_var"
    ]
    out["decode_head.linear_pred.weight"] = hf_state_dict[
        "decode_head.classifier.weight"
    ]
    out["decode_head.linear_pred.bias"] = hf_state_dict["decode_head.classifier.bias"]

    return out

keras_climate.weights.mappings.ssl4eo_mapping.build_ssl4eo_mapper 

build_ssl4eo_mapper(keras_prefix='backbone', layer_counts=(3, 4, 6, 3))
Source code in keras_climate/weights/mappings/ssl4eo_mapping.py 
 8
 9
10
11
12
13
14
15
16
17
18
def build_ssl4eo_mapper(keras_prefix="backbone", layer_counts=(3, 4, 6, 3)):
    rules = build_resnet_backbone_mapper(keras_prefix, layer_counts)
    compiled = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, repl in compiled:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)
        return None

    return mapper

keras_climate.weights.mappings.tft_mapping.build_tft_mapper 

build_tft_mapper(num_past_vars, num_future_vars, hidden_dim)
Source code in keras_climate/weights/mappings/tft_mapping.py 
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
def build_tft_mapper(num_past_vars, num_future_vars, hidden_dim):
    rules = [
        (r"^static_embed\.weight$", "static_embed/kernel"),
        (r"^static_embed\.bias$", "static_embed/bias"),
        (r"^quantile_head\.weight$", "quantile_head/kernel"),
        (r"^quantile_head\.bias$", "quantile_head/bias"),
        (r"^post_lstm_norm\.weight$", "post_lstm_norm/gamma"),
        (r"^post_lstm_norm\.bias$", "post_lstm_norm/beta"),
        (r"^post_attn_norm\.weight$", "post_attn_norm/gamma"),
        (r"^post_attn_norm\.bias$", "post_attn_norm/beta"),
        (r"^final_norm\.weight$", "final_norm/gamma"),
        (r"^final_norm\.bias$", "final_norm/beta"),
    ]
    for i in range(num_past_vars):
        rules += [
            (rf"^past_var{i}_embed\.weight$", f"past_var{i}_embed/kernel"),
            (rf"^past_var{i}_embed\.bias$", f"past_var{i}_embed/bias"),
        ]
    for i in range(num_future_vars):
        rules += [
            (rf"^future_var{i}_embed\.weight$", f"future_var{i}_embed/kernel"),
            (rf"^future_var{i}_embed\.bias$", f"future_var{i}_embed/bias"),
        ]

    for name in (
        "static_ctx_enrichment",
        "static_ctx_h",
        "static_ctx_c",
        "post_lstm_gate",
        "static_enrichment",
        "positionwise_ff",
    ):
        rules += _grn_rules(name, name, has_skip=False)

    rules += _vsn_rules(
        "past_vsn",
        "past_vsn",
        num_past_vars,
        hidden_dim,
        flatten_in_dim=num_past_vars * hidden_dim,
    )
    rules += _vsn_rules(
        "future_vsn",
        "future_vsn",
        num_future_vars,
        hidden_dim,
        flatten_in_dim=num_future_vars * hidden_dim,
    )

    num_heads_guess = 16
    rules += [
        (r"^attention\.v_layer\.weight$", "attention/v/kernel"),
        (r"^attention\.v_layer\.bias$", "attention/v/bias"),
        (r"^attention\.out_proj\.weight$", "attention/out_proj/kernel"),
        (r"^attention\.out_proj\.bias$", "attention/out_proj/bias"),
    ]

    def head_rules(num_heads):
        r = []
        for h in range(num_heads):
            r += [
                (rf"^attention\.q_layers\.{h}\.weight$", f"attention/q{h}/kernel"),
                (rf"^attention\.q_layers\.{h}\.bias$", f"attention/q{h}/bias"),
                (rf"^attention\.k_layers\.{h}\.weight$", f"attention/k{h}/kernel"),
                (rf"^attention\.k_layers\.{h}\.bias$", f"attention/k{h}/bias"),
            ]
        return r

    rules += head_rules(64)

    compiled = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, repl in compiled:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)
        return None

    return mapper

keras_climate.weights.mappings.tft_mapping.convert_tft_lstm_state_dict 

convert_tft_lstm_state_dict(flat_state_dict, torch_prefix, keras_prefix)
Source code in keras_climate/weights/mappings/tft_mapping.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
def convert_tft_lstm_state_dict(flat_state_dict, torch_prefix, keras_prefix):
    out = {}
    w_ih = flat_state_dict.get(f"{torch_prefix}.weight_ih_l0")
    w_hh = flat_state_dict.get(f"{torch_prefix}.weight_hh_l0")
    b_ih = flat_state_dict.get(f"{torch_prefix}.bias_ih_l0")
    b_hh = flat_state_dict.get(f"{torch_prefix}.bias_hh_l0")
    if w_ih is not None:
        out[f"{keras_prefix}/lstm_cell/kernel"] = w_ih
    if w_hh is not None:
        out[f"{keras_prefix}/lstm_cell/recurrent_kernel"] = w_hh
    if b_ih is not None and b_hh is not None:
        out[f"{keras_prefix}/lstm_cell/bias"] = b_ih + b_hh
    return out

keras_climate.weights.mappings.timesnet_mapping.build_timesnet_mapper 

build_timesnet_mapper(num_layers, num_kernels, use_revin=True)
Source code in keras_climate/weights/mappings/timesnet_mapping.py 
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
def build_timesnet_mapper(num_layers, num_kernels, use_revin=True):
    rules = [
        (r"^value_embed\.weight$", "value_embed/kernel"),
        (r"^value_embed\.bias$", "value_embed/bias"),
        (r"^predict_linear\.weight$", "predict_linear/kernel"),
        (r"^predict_linear\.bias$", "predict_linear/bias"),
        (r"^output_proj\.weight$", "output_proj/kernel"),
        (r"^output_proj\.bias$", "output_proj/bias"),
    ]
    if use_revin:
        rules += [
            (r"^revin\.affine_weight$", "revin/gamma"),
            (r"^revin\.affine_bias$", "revin/beta"),
        ]

    for i in range(num_layers):
        rules += [
            (rf"^norm{i}\.weight$", f"norm{i}/gamma"),
            (rf"^norm{i}\.bias$", f"norm{i}/beta"),
        ]
        for conv_name in ("conv1", "conv2"):
            for k in range(num_kernels):
                tp = f"timesblock{i}.{conv_name}.convs.{k}"
                kp = f"timesblock{i}/{conv_name}/kernel{k}"
                rules += [
                    (rf"^{re.escape(tp)}\.weight$", f"{kp}/kernel"),
                    (rf"^{re.escape(tp)}\.bias$", f"{kp}/bias"),
                ]

    compiled = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, repl in compiled:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)
        return None

    return mapper

keras_climate.weights.mappings.unet_mapping.build_mapper 

build_mapper(depth=4)
Source code in keras_climate/weights/mappings/unet_mapping.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
def build_mapper(depth=4):
    rules = []
    rules += _double_conv_rules("inc", "enc0")
    for i in range(1, depth):
        rules += _double_conv_rules(f"down{i}.maxpool_conv.1", f"enc{i}")
    rules += _double_conv_rules(f"down{depth}.maxpool_conv.1", "bottleneck")

    for k in range(1, depth + 1):
        dec_idx = depth - k
        rules += [
            (rf"^up{k}\.up\.weight$", f"upconv{dec_idx}/kernel"),
            (rf"^up{k}\.up\.bias$", f"upconv{dec_idx}/bias"),
        ]
        rules += _double_conv_rules(f"up{k}.conv", f"dec{dec_idx}")

    rules += [
        (r"^outc\.conv\.weight$", "logits/kernel"),
        (r"^outc\.conv\.bias$", "logits/bias"),
    ]

    compiled = [(re.compile(p), r) for p, r in rules]

    def mapper(torch_key):
        for pattern, repl in compiled:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)
        return None

    return mapper

keras_climate.weights.mappings.vit_mapping.build_vit_mapper 

build_vit_mapper(keras_prefix, block_name_fn=None, extra_rules=None)
Source code in keras_climate/weights/mappings/vit_mapping.py 
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
def build_vit_mapper(keras_prefix, block_name_fn=None, extra_rules=None):
    block_name_fn = block_name_fn or (lambda i: f"block{i}")

    rules = list(extra_rules or [])
    rules += [
        (r"^cls_token$", f"{keras_prefix}/cls_token"),
        (r"^pos_embed$", f"{keras_prefix}/pos_embed"),
        (r"^patch_embed\.proj\.weight$", f"{keras_prefix}/patch_embed/proj/kernel"),
        (r"^patch_embed\.proj\.bias$", f"{keras_prefix}/patch_embed/proj/bias"),
        (r"^norm\.weight$", f"{keras_prefix}/norm/gamma"),
        (r"^norm\.bias$", f"{keras_prefix}/norm/beta"),
    ]

    compiled_static = [(re.compile(p), r) for p, r in rules]

    block_pattern = re.compile(
        r"^blocks\.(\d+)\.(norm1|norm2)\.(weight|bias)$|"
        r"^blocks\.(\d+)\.attn\.(qkv|proj)\.(weight|bias)$|"
        r"^blocks\.(\d+)\.mlp\.(fc1|fc2)\.(weight|bias)$"
    )

    def mapper(torch_key):
        for pattern, repl in compiled_static:
            if pattern.match(torch_key):
                return pattern.sub(repl, torch_key)

        m = block_pattern.match(torch_key)
        if not m:
            return None

        if m.group(1) is not None:
            idx, sub, param = m.group(1), m.group(2), m.group(3)
            keras_sub = {"norm1": "norm1", "norm2": "norm2"}[sub]
            param_name = "gamma" if param == "weight" else "beta"
            return f"{keras_prefix}/{block_name_fn(int(idx))}/{keras_sub}/{param_name}"

        if m.group(4) is not None:
            idx, sub, param = m.group(4), m.group(5), m.group(6)
            keras_param = "kernel" if param == "weight" else "bias"
            return f"{keras_prefix}/{block_name_fn(int(idx))}/attn/{sub}/{keras_param}"

        idx, sub, param = m.group(7), m.group(8), m.group(9)
        keras_param = "kernel" if param == "weight" else "bias"
        return f"{keras_prefix}/{block_name_fn(int(idx))}/mlp/{sub}/{keras_param}"

    return mapper