Method forward

axlearn/common/quantizer.py:362–407

Computes random projection and quantization. Args: inputs: Tensor of shape [batch_size, seq_len, input_dim]. paddings: 0/1 Tensor of shape [batch_size, seq_len]. Returns: BaseQuantizer.Output.

(self, inputs: Tensor, *, paddings: Tensor)

        return params

    def forward(self, inputs: Tensor, *, paddings: Tensor) -> BaseQuantizer.Output:
        """Computes random projection and quantization.

        Args:
            inputs: Tensor of shape [batch_size, seq_len, input_dim].
            paddings: 0/1 Tensor of shape [batch_size, seq_len].

        Returns:
            BaseQuantizer.Output.
        """
        cfg = self.config

        # [batch_size, seq_len, num_codebooks * codebook_dim].
        inputs = self.rand_proj(inputs)
        inputs_by_group = jnp.reshape(
            inputs, list(inputs.shape[:2]) + [cfg.num_codebooks, cfg.codebook_dim]
        )

        if cfg.normalize_inputs:
            # [..., num_codebooks, codebook_dim].
            inputs_by_group = l2_normalize(inputs_by_group, axis=-1, eps=1e-12)

        # When codebook is normalized, dot_product is equivalent to l2_distance.
        metric = (
            SimilarityMetric.DOT_PRODUCT if cfg.normalize_codebook else SimilarityMetric.L2_DISTANCE
        )
        q_outputs = quantize_by_nearest_neighbor(
            inputs=inputs_by_group,
            codebook=self.parameters["codebook"],
            metric=metric,
        )
        q_outputs = _apply_paddings(outputs=q_outputs, paddings=paddings)
        # Best-rq freezes the codebook.
        ids = jax.lax.stop_gradient(q_outputs.ids)
        quantized_vectors = jax.lax.stop_gradient(q_outputs.quantized_vectors)

        outputs = self.Output(
            # [batch_size, seq_len, num_codebooks].
            ids=ids,
            # [batch_size, seq_len, num_codebooks, codebook_dim].
            quantized_vectors=quantized_vectors,
        )

        onehots = _ids_to_onehots(outputs.ids, codebook_size=cfg.codebook_size, dtype=jnp.int32)
        _add_codebook_summaries(context=current_context(), onehots=onehots, paddings=paddings)
        return outputs


class KmeansVectorQuantizer(BaseQuantizer):