{
  "$schema": "https://ambiqai.github.io/helia-ui/schema/reference-model-1.json",
  "generatedFrom": {
    "sourceCommit": "f341fb11f7f5d77a4974ba8273c4cd55d67c21f0",
    "tool": "pyref",
    "version": "1.7.3"
  },
  "language": "python",
  "modules": [
    {
      "description": "# MetaFormer: Meta-Learning with Transformers\n\n## Overview\n\nMetaFormer is a transformer-based model that incorporates both spatial mixing and channel mixing blocks.\nThe architecture is designed to learn from few examples and generalize to new tasks.\n\nFor more info, refer to the original paper [MetaFormer: Meta-Learning with Transformers](https://arxiv.org/abs/2110.11605).\n\nParameters are in ``helia_edge.models.metaformer_params``.\n\n**Functions**\n\n| Name | Description |\n| --- | --- |\n| `build` | MetaFormer model from ``MetaFormerParams`` |\n| `patch_embedding` | Patch embedding layer |\n| `pool_token_mixer` | Token mixer using average pooling |\n| `conv_token_mixer` | Token mixer using separable convolution |\n| `attention_token_mixer` | Token mixer using multi-head attention |\n| `mlp_channel_mixer` | Channel mixer using MLP via 1x1 convolutions |\n| `metaformer_block` | Metaformer block |\n| `metaformer_layer` | Metaformer functional layer |",
      "name": "metaformer",
      "path": "helia_edge.models.metaformer",
      "submodules": [],
      "summary": "MetaFormer: Meta-Learning with Transformers",
      "symbols": [
        {
          "description": "Patch embedding layer using 2D convolution",
          "examples": [],
          "id": "helia_edge.models.metaformer.patch_embedding",
          "kind": "function",
          "language": "python",
          "members": [],
          "name": "patch_embedding",
          "params": [
            {
              "description": "Embedding dimension",
              "name": "embed_dim",
              "type": "int"
            },
            {
              "description": "Patch shape",
              "name": "patch_shape",
              "type": "tuple[int, int]"
            },
            {
              "default": "None",
              "description": "Stride shape. Defaults to None.",
              "name": "stride_shape",
              "type": "tuple[int, int]"
            },
            {
              "default": "'same'",
              "description": "Padding. Defaults to 'same'.",
              "name": "padding",
              "type": "str"
            }
          ],
          "raises": [],
          "returns": [],
          "signature": "patch_embedding(\n    embed_dim: int,\n    patch_shape: tuple[int, int],\n    stride_shape: tuple[int, int] | None = None,\n    padding: str = 'same',\n) -> keras.layers.Layer",
          "source": {
            "line": 30,
            "path": "helia_edge/models/metaformer.py",
            "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/models/metaformer.py#L30"
          },
          "summary": "Patch embedding layer using 2D convolution"
        },
        {
          "description": "Token mixer using average pooling",
          "examples": [],
          "id": "helia_edge.models.metaformer.pool_token_mixer",
          "kind": "function",
          "language": "python",
          "members": [],
          "name": "pool_token_mixer",
          "params": [
            {
              "default": "(2, 2)",
              "description": "Pool size. Defaults to (2, 2).",
              "name": "pool_size",
              "type": "tuple[int, int]"
            }
          ],
          "raises": [],
          "returns": [
            {
              "description": "keras.layers.Layer: Token mixer layer",
              "type": "keras.layers.Layer"
            }
          ],
          "signature": "pool_token_mixer(pool_size: tuple[int, int] = (2, 2)) -> keras.layers.Layer",
          "source": {
            "line": 56,
            "path": "helia_edge/models/metaformer.py",
            "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/models/metaformer.py#L56"
          },
          "summary": "Token mixer using average pooling"
        },
        {
          "description": "Token mixer using separable convolution",
          "examples": [],
          "id": "helia_edge.models.metaformer.conv_token_mixer",
          "kind": "function",
          "language": "python",
          "members": [],
          "name": "conv_token_mixer",
          "params": [
            {
              "description": "Embedding dimension",
              "name": "embed_dim",
              "type": "int"
            },
            {
              "default": "(3, 3)",
              "description": "Kernel size. Defaults to (3, 3).",
              "name": "kernel_size",
              "type": "tuple[int, int]"
            },
            {
              "default": "(1, 1)",
              "description": "Strides. Defaults to (1, 1).",
              "name": "strides",
              "type": "tuple[int, int]"
            }
          ],
          "raises": [],
          "returns": [
            {
              "description": "keras.layers.Layer: Token mixer layer",
              "type": "keras.Layer"
            }
          ],
          "signature": "conv_token_mixer(\n    embed_dim: int,\n    kernel_size: tuple[int, int] = (3, 3),\n    strides: tuple[int, int] = (1, 1),\n) -> keras.Layer",
          "source": {
            "line": 84,
            "path": "helia_edge/models/metaformer.py",
            "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/models/metaformer.py#L84"
          },
          "summary": "Token mixer using separable convolution"
        },
        {
          "description": "Token mixer using multi-head attention",
          "examples": [],
          "id": "helia_edge.models.metaformer.attention_token_mixer",
          "kind": "function",
          "language": "python",
          "members": [],
          "name": "attention_token_mixer",
          "params": [
            {
              "description": "Embedding dimension",
              "name": "embed_dim",
              "type": "int"
            },
            {
              "description": "Number of heads",
              "name": "num_heads",
              "type": "int"
            },
            {
              "default": "0.1",
              "description": "Dropout rate. Defaults to 0.1.",
              "name": "dropout",
              "type": "float"
            }
          ],
          "raises": [],
          "returns": [
            {
              "description": "keras.layers.Layer: Token mixer layer",
              "type": "keras.layers.Layer"
            }
          ],
          "signature": "attention_token_mixer(embed_dim: int, num_heads: int, dropout: float = 0.1) -> keras.layers.Layer",
          "source": {
            "line": 115,
            "path": "helia_edge/models/metaformer.py",
            "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/models/metaformer.py#L115"
          },
          "summary": "Token mixer using multi-head attention"
        },
        {
          "description": "Channel mixer using MLP via 1x1 convolutions",
          "examples": [],
          "id": "helia_edge.models.metaformer.mlp_channel_mixer",
          "kind": "function",
          "language": "python",
          "members": [],
          "name": "mlp_channel_mixer",
          "params": [
            {
              "description": "Embedding dimension",
              "name": "embed_dim",
              "type": "int"
            },
            {
              "default": "4",
              "description": "Expansion ratio. Defaults to 4.",
              "name": "ratio",
              "type": "int"
            },
            {
              "default": "'gelu'",
              "description": "Activation function. Defaults to \"gelu\".",
              "name": "activation",
              "type": "str"
            },
            {
              "default": "0",
              "description": "Dropout rate. Defaults to 0.",
              "name": "dropout",
              "type": "float"
            }
          ],
          "raises": [],
          "returns": [
            {
              "description": "keras.layers.Layer: Channel mixer layer",
              "type": "keras.Layer"
            }
          ],
          "signature": "mlp_channel_mixer(embed_dim: int, ratio: int = 4, activation: str = 'gelu', dropout: float = 0) -> keras.Layer",
          "source": {
            "line": 151,
            "path": "helia_edge/models/metaformer.py",
            "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/models/metaformer.py#L151"
          },
          "summary": "Channel mixer using MLP via 1x1 convolutions"
        },
        {
          "description": "Metaformer block",
          "examples": [],
          "id": "helia_edge.models.metaformer.metaformer_block",
          "kind": "function",
          "language": "python",
          "members": [],
          "name": "metaformer_block",
          "params": [
            {
              "default": "None",
              "description": "Token mixer layer. Defaults to None.",
              "name": "token_mixer",
              "type": "keras.layers.Layer"
            },
            {
              "default": "None",
              "description": "Channel mixer layer. Defaults to None.",
              "name": "channel_mixer",
              "type": "keras.layers.Layer"
            },
            {
              "default": "'mf_block'",
              "description": "Block name. Defaults to 'mf_block'.",
              "name": "name",
              "type": "str"
            }
          ],
          "raises": [],
          "returns": [
            {
              "description": "keras.layers.Layer: Metaformer block",
              "type": "keras.layers.Layer"
            }
          ],
          "signature": "metaformer_block(\n    token_mixer: keras.layers.Layer | None = None,\n    channel_mixer: keras.layers.Layer | None = None,\n    name: str = 'mf_block',\n) -> keras.layers.Layer",
          "source": {
            "line": 194,
            "path": "helia_edge/models/metaformer.py",
            "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/models/metaformer.py#L194"
          },
          "summary": "Metaformer block"
        },
        {
          "description": "MetaFormer functional layer",
          "examples": [],
          "id": "helia_edge.models.metaformer.metaformer_layer",
          "kind": "function",
          "language": "python",
          "members": [],
          "name": "metaformer_layer",
          "params": [
            {
              "description": "Input tensor",
              "name": "x",
              "type": "keras.KerasTensor"
            },
            {
              "description": "Model parameters.",
              "name": "params",
              "type": "MetaFormerParams"
            }
          ],
          "raises": [],
          "returns": [
            {
              "description": "keras.KerasTensor: Output tensor",
              "type": "keras.KerasTensor"
            }
          ],
          "signature": "metaformer_layer(x: keras.KerasTensor, params: MetaFormerParams) -> keras.KerasTensor",
          "source": {
            "line": 242,
            "path": "helia_edge/models/metaformer.py",
            "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/models/metaformer.py#L242"
          },
          "summary": "MetaFormer functional layer"
        },
        {
          "description": "Build a MetaFormer model.",
          "examples": [],
          "id": "helia_edge.models.metaformer.build",
          "kind": "function",
          "language": "python",
          "members": [],
          "name": "build",
          "params": [
            {
              "description": "Model parameters.",
              "name": "params",
              "type": "MetaFormerParams"
            },
            {
              "description": "Input shape without the batch axis; None for a variable axis.",
              "name": "input_shape",
              "type": "tuple[int | None, ...]"
            },
            {
              "default": "None",
              "description": "Static batch size; None for a dynamic batch.",
              "name": "batch_size",
              "type": "int | None"
            },
            {
              "default": "None",
              "description": "Model name; the family when None.",
              "name": "name",
              "type": "str | None"
            }
          ],
          "raises": [],
          "returns": [
            {
              "description": "keras.Model: The model, named ``metaformer`` unless ``name`` is given.",
              "type": "keras.Model"
            }
          ],
          "signature": "build(\n    params: MetaFormerParams,\n    input_shape: tuple[int | None, ...],\n    *,\n    batch_size: int | None = None,\n    name: str | None = None,\n) -> keras.Model",
          "source": {
            "line": 412,
            "path": "helia_edge/models/metaformer.py",
            "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/models/metaformer.py#L412"
          },
          "summary": "Build a MetaFormer model."
        }
      ]
    }
  ],
  "name": "helia_edge"
}