# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import torch
from xautodl import spaces
from xautodl.xlayers import super_core

def _create_stel(input_dim, output_dim, order):
    return super_core.SuperSequential(
        super_core.SuperLinear(input_dim, output_dim),
        super_core.SuperTransformerEncoderLayer(
            output_dim,
            num_heads=spaces.Categorical(2, 4, 6),
            mlp_hidden_multiplier=spaces.Categorical(1, 2, 4),
            order=order,
        ),
    )


# -

batch, seq_dim, input_dim = 1, 4, 6
order = super_core.LayerOrder.PreNorm

out1_dim = spaces.Categorical(12, 24, 36)
out2_dim = spaces.Categorical(24, 36, 48)
out3_dim = spaces.Categorical(36, 72, 100)
layer1 = _create_stel(input_dim, out1_dim, order)
layer2 = _create_stel(out1_dim, out2_dim, order)
layer3 = _create_stel(out2_dim, out3_dim, order)
model = super_core.SuperSequential(layer1, layer2, layer3)
print(model)

inputs = torch.rand(batch, seq_dim, input_dim)
outputs = model(inputs)

abstract_space = model.abstract_search_space
abstract_space.clean_last()
abstract_child = abstract_space.random(reuse_last=True)
# print("The abstract child program is:\n{:}".format(abstract_child))
model.enable_candidate()
model.set_super_run_type(super_core.SuperRunMode.Candidate)
model.apply_candidate(abstract_child)
outputs = model(inputs)

print(outputs.shape)



StarcoderdataJupyterScriptsDedupFiltered