# helia_edge.layers.gumbel_softmax_bottleneck

## helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck

`class` · `python`

```python
GumbelSoftmaxBottleneck(
    num_embeddings: int,
    embedding_dim: int,
    temperature: float = 1.0,
    hard: bool = True,
    input_is_logits: bool = False,
    use_bias: bool = True,
    kl_weight: float = 1.0,
    **kwargs={},
)
```

Discrete bottleneck via Gumbel-Softmax (Concrete) with optional straight-through hard one-hot.

:::note[Inputs]
x: [..., Din]  (features)  -- if input_is_logits=False (default), we learn a linear proj to K logits
   OR
x: [..., K]    (logits)    -- if input_is_logits=True, we treat last dim as K logits directly

:::

:::note[Outputs]
z: [..., D]    expected embedding  z = soft_one_hot @ embed   (D = embedding_dim)

:::

:::note[Adds loss]
kl_weight * mean_bits_per_index    (KL(q || Uniform(K)) in *bits*, averaged over tokens)

:::

Tracks metrics (logged via `metrics`):
  - gs_bits_per_index   (lower bound, bits/index)
  - gs_perplexity       (empirical perplexity from hard argmax histogram)
  - gs_usage            (fraction of codes used at least once in the batch)
  - gs_temperature      (current τ; useful when annealing)

Source: `helia_edge/layers/gumbel_softmax_bottleneck.py:6`

### helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.K

`constant` · `python`

```python
K = int(num_embeddings)
```

Source: `helia_edge/layers/gumbel_softmax_bottleneck.py:42`

### helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.D

`constant` · `python`

```python
D = int(embedding_dim)
```

Source: `helia_edge/layers/gumbel_softmax_bottleneck.py:43`

### helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.hard

`attribute` · `python`

```python
hard = bool(hard)
```

Source: `helia_edge/layers/gumbel_softmax_bottleneck.py:45`

### helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.input_is_logits

`attribute` · `python`

```python
input_is_logits = bool(input_is_logits)
```

Source: `helia_edge/layers/gumbel_softmax_bottleneck.py:46`

### helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.use_bias

`attribute` · `python`

```python
use_bias = bool(use_bias)
```

Source: `helia_edge/layers/gumbel_softmax_bottleneck.py:47`

### helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.kl_weight

`attribute` · `python`

```python
kl_weight = float(kl_weight)
```

Source: `helia_edge/layers/gumbel_softmax_bottleneck.py:48`

### helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.tau

`attribute` · `python`

```python
tau = self.add_weight(name='temperature', shape=(), initializer=keras.initializers.Constant(float(temperature)), trainable=False, dtype='float32')
```

Source: `helia_edge/layers/gumbel_softmax_bottleneck.py:51`

### helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.metrics

`attribute` · `python`

```python
metrics
```

Source: `helia_edge/layers/gumbel_softmax_bottleneck.py:179`

### helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.build

`method` · `python`

```python
build(input_shape)
```

Source: `helia_edge/layers/gumbel_softmax_bottleneck.py:70`

### helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.set_temperature

`method` · `python`

```python
set_temperature(value: float)
```

Source: `helia_edge/layers/gumbel_softmax_bottleneck.py:106`

### helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.call

`method` · `python`

```python
call(x, training=False, return_indices: bool = False, return_probs: bool = False)
```

Source: `helia_edge/layers/gumbel_softmax_bottleneck.py:114`

### helia_edge.layers.gumbel_softmax_bottleneck.GumbelSoftmaxBottleneck.get_config

`method` · `python`

```python
get_config()
```

Source: `helia_edge/layers/gumbel_softmax_bottleneck.py:182`