src/transformers/models/beit/modeling_flax_beit.py/FlaxBeitPreTrainedModel/init_weights
class FlaxBeitPreTrainedModel:    def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple) -> FrozenDict:
        # init input tensors
        pixel_values = jnp.zeros(input_shape, dtype=self.dtype)

        params_rng, dropout_rng = jax.random.split(rng)
        dropout_rng, droppath_rng = jax.random.split(dropout_rng)
        rngs = {"params": params_rng, "dropout": dropout_rng, "droppath": droppath_rng}

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