src/transformers/models/wav2vec2/modeling_flax_wav2vec2.py/FlaxWav2Vec2PreTrainedModel/init_weights
class FlaxWav2Vec2PreTrainedModel:    def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple) -> FrozenDict:
        # init input tensors
        input_values = jnp.zeros(input_shape, dtype="i4")
        attention_mask = jnp.ones_like(input_values)
        params_rng, dropout_rng = jax.random.split(rng, 2)
        rngs = {"params": params_rng, "dropout": dropout_rng}

        return self.module.init(rngs, input_values, attention_mask, return_dict=False)["params"]