{
  "$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": "",
      "name": "residual_vector_quantizer",
      "path": "helia_edge.layers.residual_vector_quantizer",
      "submodules": [],
      "summary": "",
      "symbols": [
        {
          "description": "Residual Vector Quantizer (RVQ) with straight-through estimator.\n\nInput:  [..., D]  (last dim = embedding_dim)\nOutput: [..., D]  (sum of per-level dequantized vectors; gradients pass through x)\n\nMetrics (logged via `metrics` property):\n  - rvq_l{l}_perplexity, rvq_l{l}_usage, rvq_l{l}_bits_per_index\n  - rvq_perplexity_mean, rvq_usage_mean, rvq_bits_per_index_sum (entropy lower bound)\n\n:::note[Losses added per level]\n- beta * ||stop(q_l) - r_l||^2  +  ||q_l - stop(r_l)||^2,\n  where r_l is the current residual and q_l the level-l code vector.\n\n:::",
          "examples": [],
          "id": "helia_edge.layers.residual_vector_quantizer.ResidualVectorQuantizer",
          "kind": "class",
          "language": "python",
          "members": [
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.residual_vector_quantizer.ResidualVectorQuantizer.M",
              "kind": "constant",
              "language": "python",
              "members": [],
              "name": "M",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "M = int(num_levels)",
              "source": {
                "line": 42,
                "path": "helia_edge/layers/residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/residual_vector_quantizer.py#L42"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.residual_vector_quantizer.ResidualVectorQuantizer.D",
              "kind": "constant",
              "language": "python",
              "members": [],
              "name": "D",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "D = int(embedding_dim)",
              "source": {
                "line": 43,
                "path": "helia_edge/layers/residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/residual_vector_quantizer.py#L43"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.residual_vector_quantizer.ResidualVectorQuantizer.Ks",
              "kind": "attribute",
              "language": "python",
              "members": [],
              "name": "Ks",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "Ks",
              "source": {
                "line": 48,
                "path": "helia_edge/layers/residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/residual_vector_quantizer.py#L48"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.residual_vector_quantizer.ResidualVectorQuantizer.beta",
              "kind": "attribute",
              "language": "python",
              "members": [],
              "name": "beta",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "beta = float(beta)",
              "source": {
                "line": 51,
                "path": "helia_edge/layers/residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/residual_vector_quantizer.py#L51"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.residual_vector_quantizer.ResidualVectorQuantizer.metrics",
              "kind": "attribute",
              "language": "python",
              "members": [],
              "name": "metrics",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "metrics",
              "source": {
                "line": 159,
                "path": "helia_edge/layers/residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/residual_vector_quantizer.py#L159"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.residual_vector_quantizer.ResidualVectorQuantizer.build",
              "kind": "method",
              "language": "python",
              "members": [],
              "name": "build",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "build(input_shape)",
              "source": {
                "line": 64,
                "path": "helia_edge/layers/residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/residual_vector_quantizer.py#L64"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.residual_vector_quantizer.ResidualVectorQuantizer.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, also returns list of flat indices (one tensor per level).",
                  "name": "return_indices",
                  "type": "bool"
                }
              ],
              "raises": [],
              "returns": [
                {
                  "description": "y or (y, indices_list): dequantized vector and optional per-level indices.",
                  "type": "keras.KerasTensor | tuple[keras.KerasTensor, list[keras.KerasTensor]]"
                }
              ],
              "signature": "call(\n    x: keras.KerasTensor,\n    return_indices: bool = False,\n) -> keras.KerasTensor | tuple[keras.KerasTensor, list[keras.KerasTensor]]",
              "source": {
                "line": 93,
                "path": "helia_edge/layers/residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/residual_vector_quantizer.py#L93"
              },
              "summary": ""
            },
            {
              "description": "Return list of per-level flat index tensors [N] (no gradients).",
              "examples": [],
              "id": "helia_edge.layers.residual_vector_quantizer.ResidualVectorQuantizer.encode",
              "kind": "method",
              "language": "python",
              "members": [],
              "name": "encode",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "encode(x)",
              "source": {
                "line": 164,
                "path": "helia_edge/layers/residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/residual_vector_quantizer.py#L164"
              },
              "summary": "Return list of per-level flat index tensors [N] (no gradients)."
            },
            {
              "description": "Sum per-level code vectors from indices_list and reshape to original_shape.",
              "examples": [],
              "id": "helia_edge.layers.residual_vector_quantizer.ResidualVectorQuantizer.decode",
              "kind": "method",
              "language": "python",
              "members": [],
              "name": "decode",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "decode(indices_list, original_shape)",
              "source": {
                "line": 176,
                "path": "helia_edge/layers/residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/residual_vector_quantizer.py#L176"
              },
              "summary": "Sum per-level code vectors from indiceslist and reshape to originalshape."
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.residual_vector_quantizer.ResidualVectorQuantizer.get_config",
              "kind": "method",
              "language": "python",
              "members": [],
              "name": "get_config",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "get_config()",
              "source": {
                "line": 185,
                "path": "helia_edge/layers/residual_vector_quantizer.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/residual_vector_quantizer.py#L185"
              },
              "summary": ""
            }
          ],
          "name": "ResidualVectorQuantizer",
          "params": [
            {
              "description": "int, number of residual VQ stages (M >= 1)",
              "name": "num_levels",
              "type": "int"
            },
            {
              "description": "int OR sequence[int], codebook size K for each level",
              "name": "num_embeddings",
              "type": "int | Sequence[int]"
            },
            {
              "description": "int, latent dimensionality D",
              "name": "embedding_dim",
              "type": "int"
            },
            {
              "default": "0.25",
              "description": "float, commitment coefficient per level",
              "name": "beta",
              "type": "float"
            }
          ],
          "raises": [],
          "returns": [],
          "signature": "ResidualVectorQuantizer(\n    num_levels: int,\n    num_embeddings: int | Sequence[int],\n    embedding_dim: int,\n    beta: float = 0.25,\n    **kwargs={},\n)",
          "source": {
            "line": 8,
            "path": "helia_edge/layers/residual_vector_quantizer.py",
            "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/residual_vector_quantizer.py#L8"
          },
          "summary": "Residual Vector Quantizer (RVQ) with straight-through estimator."
        }
      ]
    }
  ],
  "name": "helia_edge"
}