code: def stochastic_depth(
      input: torch.tensor, p: float, mode: str, training: bool = True
  ):
      """Apply stochastic depth.
  
      Implements the Stochastic Depth from `"Deep Networks with Stochastic Depth"
      <https://arxiv.org/abs/1603.09382>`_ used for randomly dropping residual
      branches of residual architectures.
  
      Args:
          input (Tensor[N, ...]): The input tensor or arbitrary dimensions with the first
                      one being its batch i.e. a batch with ``N`` rows.
          p (float): probability of the input to be zeroed.
          mode (str): ``"batch"`` or ``"row"``.
                      ``"batch"`` randomly zeroes the entire input, ``"row"`` zeroes
                      randomly selected rows from the batch.
          training: apply stochastic depth if is ``True``. Default: ``True``
  
      Returns:
          Tensor[N, ...]: The randomly zeroed tensor.
      """
      if p < 0.0 or p > 1.0:
          raise ValueError(
              "drop probability has to be between 0 and 1, but got {}".format(
                  p
              )
          )
      if mode not in ["batch", "row"]:
          raise ValueError(
              "mode has to be either 'batch' or 'row', but got {}".format(
                  mode
              )
          )
      if not training or p == 0.0:
          return input
  
      survival_rate = 1.0 - p
      if mode == "row":
          size = [input.shape[0]] + [1] * (input.ndim - 1)
      else:
          size = [1] * input.ndim
      noise = torch.empty(size, dtype=input.dtype, device=input.device)
      noise = noise.bernoulli_(survival_rate).div_(survival_rate)
      return input * noise
  
imports: import math
  from functools import partial
  import torch
  import torch.nn
  import torch.nn.functional
  from einops import rearrange
  from opt_einsum import contract
label_desc: Use image node_id 1 for calling a global function with example usage: stochastic_depth(input, p, mode, training) and returns: unknown, input