class_docstr: Stochastic depth with two paths.
docstr: Generate the stochastic depth mask in order to apply layer-drop.
docstr_tok: ["Generate","the","stochastic","depth","mask","in","order","to","apply","layer-drop","."]
code: def __call__(
      self, x: jnp.ndarray, deterministic: bool, use_aux_dropout: bool
  ) -> jnp.ndarray:
      """Generate the stochastic depth mask in order to apply layer-drop."""
      if not deterministic and use_aux_dropout:
          shape = (x.shape[0],) + (1,) * (x.ndim - 1)
          mask = jax.random.bernoulli(
              self.make_rng("dropout"), self.aux_stochastic_depth, shape
          )
          return x * (1.0 - mask)
      if not deterministic and not use_aux_dropout:
          shape = (x.shape[0],) + (1,) * (x.ndim - 1)
          mask = jax.random.bernoulli(
              self.make_rng("dropout"), self.stochastic_depth, shape
          )
          return x * (1.0 - mask)
      else:
          return x
  
imports: from typing import Any, Callable, Optional, Sequence
  from absl import logging
  import flax
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  import scipy
label_desc: Use image node_id 1 for calling the SplitStochasticDepth obj's underlying member method code with example usage: obj.__call__(x, deterministic, use_aux_dropout) and returns: unknown, unknown, x