class StochasticDepth(nn.Module):
    """Create a stochastic depth layer.

        Note: When using :meth:`Module.apply() `, make
        sure to include an RNG seed named ``'stochastic_depth'``.
        StochasticDepth isn't necessary for variable initialization.

        Reference
            - Deep Networks with Stochastic Depth
            - https://github.com/tensorflow/models/blob/v2.14.2/official/vision/modeling/layers/nn_layers.py#L226-L261
            - https://pytorch.org/vision/main/_modules/torchvision/ops/stochastic_depth.html#StochasticDepth
            - https://flax.readthedocs.io/en/latest/_modules/flax/linen/stochastic.html#Dropout

    Attributes:
        stochastic_depth_drop_rate: the stochastic depth probability.
            (_not_ the keep rate!)
        deterministic: if false the inputs are scaled by ``1 / (1 - rate)`` and
            masked, whereas if true, no mask is applied and the inputs are returned as is.
        rng_collection: the rng collection name to use when requesting an rng key
    """

    stochastic_depth_drop_rate: float
    deterministic: Optional[bool] = None
    rng_collection: str = "stochastic_depth"

    @nn.compact
    def __call__(
        self,
        inputs,
        deterministic: Optional[bool] = None,
        rng: Optional[KeyArray] = None,
    ):
        """Applies a random stochastic depth mask to the input.

        Args:
            inputs: the inputs that should be randomly masked.
            deterministic: if false the inputs are scaled by ``1 / (1 - rate)``
                and masked, whereas if true, no mask is applied and the inputs
                are returnedas is.
            rng: an optional PRNGKey used as the random key, if not specified,
                one will be generated using ``make_rng`` with the
                ``rng_collection`` name.

        Returns:
            The masked inputs reweighted to preserve mean.
        """
        deterministic = merge_param("deterministic", self.deterministic, deterministic)

        if (self.stochastic_depth_drop_rate == 0.0) or deterministic:
            return inputs

        if rng is None:
            rng = self.make_rng(self.rng_collection)

        keep_prob = 1.0 - self.stochastic_depth_drop_rate
        batch_size = inputs.shape[0]

        broadcast_shape = list([batch_size] + [1] * int(inputs.ndim - 1))

        mask = random.bernoulli(rng, p=keep_prob, shape=broadcast_shape)
        mask = jnp.broadcast_to(mask, inputs.shape)
        return lax.select(mask, inputs / keep_prob, jnp.zeros_like(inputs))


def get_stochastic_depth_rate(init_rate, i, n):
    """Get drop connect rate for the ith block.

    Args:
        init_rate: A `float` of initial drop rate.
        i: An `int` of order of the current block.
        n: An `int` total number of blocks.

    Returns:
        Drop rate of the ith block.
    """
    if init_rate < 0 or init_rate > 1:
        raise ValueError("Initial drop rate must be within 0 and 1.")
    return init_rate * float(i) / n

    rng = jax.random.PRNGKey(seed=config.seed)
    params_rng, dropout_rng, stochastic_depth_rng = jax.random.split(rng, num=3)
    rngs = {
        "params": params_rng,
        "dropout": dropout_rng,
        "stochastic_depth": stochastic_depth_rng,
    }
...
    dropout_rngs = jax.random.split(dropout_rng, jax.local_device_count())
    stochastic_depth_rngs = jax.random.split(
        stochastic_depth_rng, jax.local_device_count()
    )

def train_step(
...
):
    dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)
    stochastic_depth_rng, new_stochastic_depth_rng = jax.random.split(
        stochastic_depth_rng
    )

...
    def loss_fn(params):
        """loss function used for training."""
...
            logits = state.apply_fn(
                {"params": params},
                batch["image"],
                rngs={"dropout": dropout_rng, "stochastic_depth": stochastic_depth_rng},
            )

...

    return new_state, metrics, new_dropout_rng, new_stochastic_depth_rng

class LayerScale(nn.Module):
    """Create a layer scale.

        Reference
            - [Going deeper with Image Transformers](https://arxiv.org/abs/2103.17239v2)

    Attributes:
        init_values:
        projection_dim:
    """

    projection_dim: int
    init_values: float = 1e-6
    dtype: Any = jnp.float32

    def setup(self):
        initializer = nn.initializers.constant(value=self.init_values, dtype=self.dtype)
        self.scale = self.param(
            "scale",
            initializer,
            (self.projection_dim,),
            self.dtype,
        )

    @nn.compact
    def __call__(self, inputs):
        return inputs * self.scale

        kernel_initializer = jax.nn.initializers.truncated_normal(stddev=0.02)
        conv = partial(
            nn.Conv,
            use_bias=self.use_bias,
            kernel_init=kernel_initializer,
            dtype=self.dtype,
        )