{
  "$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": "Residual Vector Quantizer with Exponential Moving Average codebook updates.\n\nReplaces the gradient-based codebook loss with EMA updates to codebook\nembeddings, following van den Oord et al. 2017 (VQ-VAE).  Only the\ncommitment loss is back-propagated; codebook vectors are updated via\nrunning averages of assigned encoder outputs.",
      "name": "ema_residual_vector_quantizer",
      "path": "helia_edge.layers.ema_residual_vector_quantizer",
      "submodules": [],
      "summary": "Residual Vector Quantizer with Exponential Moving Average codebook updates.",
      "symbols": [
        {
          "description": "Residual VQ with EMA codebook updates.\n\nInstead of learning codebook embeddings via gradient descent (which requires\na codebook loss term), this layer maintains exponential moving averages of\ncluster assignment counts and embedding sums.  Codebook vectors are derived\nfrom these running statistics with Laplace smoothing for numerical stability.\n\nOnly the *commitment loss* is back-propagated through the encoder; the\nstraight-through estimator copies gradients from the decoder to the encoder\nas in the standard VQ-VAE.\n\nInput:  ``[..., D]``  (last dim = ``embedding_dim``)\nOutput: ``[..., D]``  (sum of per-level dequantized vectors)\n\nMetrics (logged via ``metrics`` property):\n    - ``rvq_l{l}_perplexity``, ``rvq_l{l}_usage``,\n      ``rvq_l{l}_bits_per_index``\n    - ``rvq_perplexity_mean``, ``rvq_usage_mean``,\n      ``rvq_bits_per_index_sum`` (entropy lower bound)\n\n:::note[Losses added per level]\n- ``beta * ||stop(q_l) - r_l||^2``  (commitment only; no codebook\n  gradient loss)\n\n:::\n\nExample:\n\n```python\nrvq = EmaResidualVectorQuantizer(\n    num_levels=4,\n    num_embeddings=64,\n    embedding_dim=16,\n    ema_decay=0.99,\n)\ny = rvq(z, training=True)          # forward + EMA update\ny, indices = rvq(z, return_indices=True)  # also return codes\n```\n\n:::note[References]\n- van den Oord, A., Vinyals, O. & Kavukcuoglu, K. (2017).\n  *Neural Discrete Representation Learning*. NeurIPS.\n\n:::",
          "examples": [],
          "id": "helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer",
          "kind": "class",
          "language": "python",
          "members": [
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.M",
              "kind": "constant",
              "language": "python",
              "members": [],
              "name": "M",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "M = int(num_levels)",
              "source": {
                "line": 84,
                "path": "helia_edge/layers/ema_residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/ema_residual_vector_quantizer.py#L84"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.D",
              "kind": "constant",
              "language": "python",
              "members": [],
              "name": "D",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "D = int(embedding_dim)",
              "source": {
                "line": 85,
                "path": "helia_edge/layers/ema_residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/ema_residual_vector_quantizer.py#L85"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.Ks",
              "kind": "attribute",
              "language": "python",
              "members": [],
              "name": "Ks",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "Ks",
              "source": {
                "line": 89,
                "path": "helia_edge/layers/ema_residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/ema_residual_vector_quantizer.py#L89"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.beta",
              "kind": "attribute",
              "language": "python",
              "members": [],
              "name": "beta",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "beta = float(beta)",
              "source": {
                "line": 92,
                "path": "helia_edge/layers/ema_residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/ema_residual_vector_quantizer.py#L92"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.ema_decay",
              "kind": "attribute",
              "language": "python",
              "members": [],
              "name": "ema_decay",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "ema_decay = float(ema_decay)",
              "source": {
                "line": 93,
                "path": "helia_edge/layers/ema_residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/ema_residual_vector_quantizer.py#L93"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.epsilon",
              "kind": "attribute",
              "language": "python",
              "members": [],
              "name": "epsilon",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "epsilon = float(epsilon)",
              "source": {
                "line": 94,
                "path": "helia_edge/layers/ema_residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/ema_residual_vector_quantizer.py#L94"
              },
              "summary": ""
            },
            {
              "description": "Expose per-level + aggregate metrics so ``Model.fit`` logs them.",
              "examples": [],
              "id": "helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.metrics",
              "kind": "attribute",
              "language": "python",
              "members": [],
              "name": "metrics",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "metrics",
              "source": {
                "line": 264,
                "path": "helia_edge/layers/ema_residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/ema_residual_vector_quantizer.py#L264"
              },
              "summary": "Expose per-level + aggregate metrics so Model.fit logs them."
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.build",
              "kind": "method",
              "language": "python",
              "members": [],
              "name": "build",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "build(input_shape)",
              "source": {
                "line": 109,
                "path": "helia_edge/layers/ema_residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/ema_residual_vector_quantizer.py#L109"
              },
              "summary": ""
            },
            {
              "description": "Quantize *x* through all residual levels.",
              "examples": [],
              "id": "helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.call",
              "kind": "method",
              "language": "python",
              "members": [],
              "name": "call",
              "params": [
                {
                  "description": "``[..., D]`` latent to be quantized.",
                  "name": "x",
                  "type": "keras.KerasTensor"
                },
                {
                  "default": "False",
                  "description": "If ``True``, run EMA codebook updates.",
                  "name": "training",
                  "type": "bool"
                },
                {
                  "default": "False",
                  "description": "If ``True``, also return per-level flat indices.",
                  "name": "return_indices",
                  "type": "bool"
                }
              ],
              "raises": [],
              "returns": [
                {
                  "description": "``y`` or ``(y, indices_list)``: dequantized vector and optional",
                  "type": "keras.KerasTensor | tuple[keras.KerasTensor, list[keras.KerasTensor]]"
                },
                {
                  "description": "per-level index tensors.",
                  "type": "keras.KerasTensor | tuple[keras.KerasTensor, list[keras.KerasTensor]]"
                }
              ],
              "signature": "call(\n    x: keras.KerasTensor,\n    training: bool = False,\n    return_indices: bool = False,\n) -> keras.KerasTensor | tuple[keras.KerasTensor, list[keras.KerasTensor]]",
              "source": {
                "line": 195,
                "path": "helia_edge/layers/ema_residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/ema_residual_vector_quantizer.py#L195"
              },
              "summary": "Quantize x through all residual levels."
            },
            {
              "description": "Return list of per-level flat index tensors ``[N]`` (no gradients).",
              "examples": [],
              "id": "helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.encode",
              "kind": "method",
              "language": "python",
              "members": [],
              "name": "encode",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "encode(x: keras.KerasTensor) -> list[keras.KerasTensor]",
              "source": {
                "line": 272,
                "path": "helia_edge/layers/ema_residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/ema_residual_vector_quantizer.py#L272"
              },
              "summary": "Return list of per-level flat index tensors [N] (no gradients)."
            },
            {
              "description": "Sum per-level code vectors from *indices_list* and reshape.",
              "examples": [],
              "id": "helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.decode",
              "kind": "method",
              "language": "python",
              "members": [],
              "name": "decode",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "decode(indices_list: list[keras.KerasTensor], original_shape: tuple[int, ...]) -> keras.KerasTensor",
              "source": {
                "line": 284,
                "path": "helia_edge/layers/ema_residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/ema_residual_vector_quantizer.py#L284"
              },
              "summary": "Sum per-level code vectors from indiceslist and reshape."
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.get_config",
              "kind": "method",
              "language": "python",
              "members": [],
              "name": "get_config",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "get_config()",
              "source": {
                "line": 296,
                "path": "helia_edge/layers/ema_residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/ema_residual_vector_quantizer.py#L296"
              },
              "summary": ""
            }
          ],
          "name": "EmaResidualVectorQuantizer",
          "params": [
            {
              "description": "Number of residual VQ stages (``M >= 1``).",
              "name": "num_levels",
              "type": "int"
            },
            {
              "description": "Codebook size ``K`` per level (int or per-level list).",
              "name": "num_embeddings",
              "type": "int | Sequence[int]"
            },
            {
              "description": "Latent dimensionality ``D``.",
              "name": "embedding_dim",
              "type": "int"
            },
            {
              "default": "0.25",
              "description": "Commitment loss coefficient.",
              "name": "beta",
              "type": "float"
            },
            {
              "default": "0.99",
              "description": "EMA decay rate for codebook updates (``0.99``–``0.999``\ntypical).",
              "name": "ema_decay",
              "type": "float"
            },
            {
              "default": "1e-05",
              "description": "Small constant for Laplace smoothing of cluster counts.",
              "name": "epsilon",
              "type": "float"
            }
          ],
          "raises": [],
          "returns": [],
          "signature": "EmaResidualVectorQuantizer(\n    num_levels: int,\n    num_embeddings: int | Sequence[int],\n    embedding_dim: int,\n    beta: float = 0.25,\n    ema_decay: float = 0.99,\n    epsilon: float = 1e-05,\n    **kwargs={},\n)",
          "source": {
            "line": 18,
            "path": "helia_edge/layers/ema_residual_vector_quantizer.py",
            "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/ema_residual_vector_quantizer.py#L18"
          },
          "summary": "Residual VQ with EMA codebook updates."
        }
      ]
    }
  ],
  "name": "helia_edge"
}