src/transformers/models/wav2vec2/modeling_flax_wav2vec2.py/FlaxWav2Vec2GumbelVectorQuantizer/setup
class FlaxWav2Vec2GumbelVectorQuantizer:    def setup(self):
        self.num_groups = self.config.num_codevector_groups
        self.num_vars = self.config.num_codevectors_per_group

        if self.config.codevector_dim % self.num_groups != 0:
            raise ValueError(
                f"`config.codevector_dim {self.config.codevector_dim} must be divisible by"
                f" `config.num_codevector_groups` {self.num_groups} for concatenation"
            )

        # storage for codebook variables (codewords)
        self.codevectors = self.param(
            "codevectors",
            jax.nn.initializers.uniform(),
            (1, self.num_groups * self.num_vars, self.config.codevector_dim // self.num_groups),
        )
        self.weight_proj = nn.Dense(
            self.num_groups * self.num_vars,
            kernel_init=jax.nn.initializers.normal(1.0, self.dtype),
            dtype=self.dtype,
        )