{
  "$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": "gumbel_softmax_bottleneck",
      "path": "helia_edge.layers.gumbel_softmax_bottleneck",
      "submodules": [],
      "summary": "",
      "symbols": [
        {
          "description": "Discrete bottleneck via Gumbel-Softmax (Concrete) with optional straight-through hard one-hot.\n\n:::note[Inputs]\nx: [..., Din]  (features)  -- if input_is_logits=False (default), we learn a linear proj to K logits\n   OR\nx: [..., K]    (logits)    -- if input_is_logits=True, we treat last dim as K logits directly\n\n:::\n\n:::note[Outputs]\nz: [..., D]    expected embedding  z = soft_one_hot @ embed   (D = embedding_dim)\n\n:::\n\n:::note[Adds loss]\nkl_weight * mean_bits_per_index    (KL(q || Uniform(K)) in *bits*, averaged over tokens)\n\n:::\n\nTracks metrics (logged via `metrics`):\n  - gs_bits_per_index   (lower bound, bits/index)\n  - gs_perplexity       (empirical perplexity from hard argmax histogram)\n  - gs_usage            (fraction of codes used at least once in the batch)\n  - gs_temperature      (current τ; useful when annealing)",
          "examples": [],
          "id": "helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck",
          "kind": "class",
          "language": "python",
          "members": [
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.K",
              "kind": "constant",
              "language": "python",
              "members": [],
              "name": "K",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "K = int(num_embeddings)",
              "source": {
                "line": 42,
                "path": "helia_edge/layers/gumbel_softmax_bottleneck.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/gumbel_softmax_bottleneck.py#L42"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.D",
              "kind": "constant",
              "language": "python",
              "members": [],
              "name": "D",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "D = int(embedding_dim)",
              "source": {
                "line": 43,
                "path": "helia_edge/layers/gumbel_softmax_bottleneck.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/gumbel_softmax_bottleneck.py#L43"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.hard",
              "kind": "attribute",
              "language": "python",
              "members": [],
              "name": "hard",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "hard = bool(hard)",
              "source": {
                "line": 45,
                "path": "helia_edge/layers/gumbel_softmax_bottleneck.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/gumbel_softmax_bottleneck.py#L45"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.input_is_logits",
              "kind": "attribute",
              "language": "python",
              "members": [],
              "name": "input_is_logits",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "input_is_logits = bool(input_is_logits)",
              "source": {
                "line": 46,
                "path": "helia_edge/layers/gumbel_softmax_bottleneck.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/gumbel_softmax_bottleneck.py#L46"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.use_bias",
              "kind": "attribute",
              "language": "python",
              "members": [],
              "name": "use_bias",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "use_bias = bool(use_bias)",
              "source": {
                "line": 47,
                "path": "helia_edge/layers/gumbel_softmax_bottleneck.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/gumbel_softmax_bottleneck.py#L47"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.kl_weight",
              "kind": "attribute",
              "language": "python",
              "members": [],
              "name": "kl_weight",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "kl_weight = float(kl_weight)",
              "source": {
                "line": 48,
                "path": "helia_edge/layers/gumbel_softmax_bottleneck.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/gumbel_softmax_bottleneck.py#L48"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.tau",
              "kind": "attribute",
              "language": "python",
              "members": [],
              "name": "tau",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "tau = self.add_weight(name='temperature', shape=(), initializer=keras.initializers.Constant(float(temperature)), trainable=False, dtype='float32')",
              "source": {
                "line": 51,
                "path": "helia_edge/layers/gumbel_softmax_bottleneck.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/gumbel_softmax_bottleneck.py#L51"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.metrics",
              "kind": "attribute",
              "language": "python",
              "members": [],
              "name": "metrics",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "metrics",
              "source": {
                "line": 179,
                "path": "helia_edge/layers/gumbel_softmax_bottleneck.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/gumbel_softmax_bottleneck.py#L179"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.build",
              "kind": "method",
              "language": "python",
              "members": [],
              "name": "build",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "build(input_shape)",
              "source": {
                "line": 70,
                "path": "helia_edge/layers/gumbel_softmax_bottleneck.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/gumbel_softmax_bottleneck.py#L70"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.set_temperature",
              "kind": "method",
              "language": "python",
              "members": [],
              "name": "set_temperature",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "set_temperature(value: float)",
              "source": {
                "line": 106,
                "path": "helia_edge/layers/gumbel_softmax_bottleneck.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/gumbel_softmax_bottleneck.py#L106"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.call",
              "kind": "method",
              "language": "python",
              "members": [],
              "name": "call",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "call(x, training=False, return_indices: bool = False, return_probs: bool = False)",
              "source": {
                "line": 114,
                "path": "helia_edge/layers/gumbel_softmax_bottleneck.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/gumbel_softmax_bottleneck.py#L114"
              },
              "summary": ""
            },
            {
              "description": "",
              "examples": [],
              "id": "helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.get_config",
              "kind": "method",
              "language": "python",
              "members": [],
              "name": "get_config",
              "params": [],
              "raises": [],
              "returns": [],
              "signature": "get_config()",
              "source": {
                "line": 182,
                "path": "helia_edge/layers/gumbel_softmax_bottleneck.py",
                "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/gumbel_softmax_bottleneck.py#L182"
              },
              "summary": ""
            }
          ],
          "name": "GumbelSoftmaxBottleneck",
          "params": [],
          "raises": [],
          "returns": [],
          "signature": "GumbelSoftmaxBottleneck(\n    num_embeddings: int,\n    embedding_dim: int,\n    temperature: float = 1.0,\n    hard: bool = True,\n    input_is_logits: bool = False,\n    use_bias: bool = True,\n    kl_weight: float = 1.0,\n    **kwargs={},\n)",
          "source": {
            "line": 6,
            "path": "helia_edge/layers/gumbel_softmax_bottleneck.py",
            "url": "https://github.com/AmbiqAI/helia-edge/blob/f341fb11f7f5d77a4974ba8273c4cd55d67c21f0/helia_edge/layers/gumbel_softmax_bottleneck.py#L6"
          },
          "summary": "Discrete bottleneck via Gumbel-Softmax (Concrete) with optional straight-through hard one-hot."
        }
      ]
    }
  ],
  "name": "helia_edge"
}