gumbel_softmax_bottleneck

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.

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)

K = int(num_embeddings)

D = int(embedding_dim)

hard = bool(hard)

input_is_logits = bool(input_is_logits)

use_bias = bool(use_bias)

kl_weight = float(kl_weight)

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

metrics

build(input_shape)

set_temperature(value: float)

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

get_config()