class_docstr: AugMix data augmentation method based on
  `"AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty" <https://arxiv.org/abs/1912.02781>`_.
  If the image is torch Tensor, it should be of type torch.uint8, and it is expected
  to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
  If img is PIL Image, it is expected to be in mode "L" or "RGB".
  
  Args:
      severity (int): The severity of base augmentation operators. Default is ``3``.
      mixture_width (int): The number of augmentation chains. Default is ``3``.
      chain_depth (int): The depth of augmentation chains. A negative value denotes stochastic depth sampled from the interval [1, 3].
          Default is ``-1``.
      alpha (float): The hyperparameter for the probability distributions. Default is ``1.0``.
      all_ops (bool): Use all operations (including brightness, contrast, color and sharpness). Default is ``True``.
      interpolation (InterpolationMode): Desired interpolation enum defined by
          :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``.
          If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported.
      fill (sequence or number, optional): Pixel fill value for the area outside the transformed
          image. If given a number, the value is used for all bands respectively.
docstr:     img (PIL Image or Tensor): Image to be transformed.
  
  Returns:
      PIL Image or Tensor: Transformed image.
docstr_tok: ["img","(","PIL","Image","or","Tensor",")",":","Image","to","be","transformed",".","Returns",":","PIL","Image","or","Tensor",":","Transformed","image","."]
code: def forward(self, orig_img: Tensor) -> Tensor:
      """
          img (PIL Image or Tensor): Image to be transformed.
  
      Returns:
          PIL Image or Tensor: Transformed image.
      """
      fill = self.fill
      channels, height, width = F.get_dimensions(orig_img)
      if isinstance(orig_img, Tensor):
          img = orig_img
          if isinstance(fill, (int, float)):
              fill = [float(fill)] * channels
          elif fill is not None:
              fill = [float(f) for f in fill]
      else:
          img = self._pil_to_tensor(orig_img)
  
      op_meta = self._augmentation_space(
          self._PARAMETER_MAX, (height, width)
      )
  
      orig_dims = list(img.shape)
      batch = img.view([1] * max(4 - img.ndim, 0) + orig_dims)
      batch_dims = [batch.size(0)] + [1] * (batch.ndim - 1)
  
      # Sample the beta weights for combining the original and augmented image. To get Beta, we use a Dirichlet
      # with 2 parameters. The 1st column stores the weights of the original and the 2nd the ones of augmented image.
      m = self._sample_dirichlet(
          torch.tensor(
              [self.alpha, self.alpha], device=batch.device
          ).expand(batch_dims[0], -1)
      )
  
      # Sample the mixing weights and combine them with the ones sampled from Beta for the augmented images.
      combined_weights = self._sample_dirichlet(
          torch.tensor(
              [self.alpha] * self.mixture_width, device=batch.device
          ).expand(batch_dims[0], -1)
      ) * m[:, 1].view([batch_dims[0], -1])
  
      mix = m[:, 0].view(batch_dims) * batch
      for i in range(self.mixture_width):
          aug = batch
          depth = (
              self.chain_depth
              if self.chain_depth > 0
              else int(torch.randint(low=1, high=4, size=(1,)).item())
          )
          for _ in range(depth):
              op_index = int(torch.randint(len(op_meta), (1,)).item())
              op_name = list(op_meta.keys())[op_index]
              magnitudes, signed = op_meta[op_name]
              magnitude = (
                  float(
                      magnitudes[
                          torch.randint(
                              self.severity, (1,), dtype=torch.long
                          )
                      ].item()
                  )
                  if magnitudes.ndim > 0
                  else 0.0
              )
              if signed and torch.randint(2, (1,)):
                  magnitude *= -1.0
              aug = _apply_op(
                  aug,
                  op_name,
                  magnitude,
                  interpolation=self.interpolation,
                  fill=fill,
              )
          mix.add_(combined_weights[:, i].view(batch_dims) * aug)
      mix = mix.view(orig_dims).to(dtype=img.dtype)
  
      if not isinstance(orig_img, Tensor):
          return self._tensor_to_pil(mix)
      return mix
  
imports: import math
  from enum import Enum
  from typing import Dict, List, Optional, Tuple
  import torch
  from torch import Tensor
  from .None import functional, InterpolationMode
label_desc: Use image node_id 6 for calling the AugMix obj's underlying member method code with example usage: obj.forward(orig_img) and returns: mix, self