code: def augmentation(
      function: Optional[Callable[..., _DataNode]] = None,
      *,
      mag_range: Optional[
          Union[Tuple[float, float], np.ndarray]
      ] = None,
      randomly_negate: Optional[bool] = None,
      mag_to_param: Optional[Callable[[float], npt.ArrayLike]] = None,
      param_device: Optional[str] = None,
      name: Optional[str] = None,
      augmentation_cls: Optional[Type[Augmentation]] = None,
  ):
      """
      A decorator turning transformations implemented with DALI into augmentations that
      can be used by automatic augmentations (e.g. AutoAugment, RandAugment, TrivialAugment).
  
      The `function` must accept at least two args: a sample and a parameter.
      The decorator handles computation of the parameter. Instead of the parameter, the
      decorated augmentation accepts magnitude bin and the total number of bins.
      Then, the bin is used to compute the parameter as if by calling
      `mag_to_param(magnitudes[magnitude_bin] * ((-1) ** random_sign))`, where
      `magnitudes=linspace(mag_range[0], mag_range[1], num_magnitude_bins)`.
  
      Args
      ----
      function : callable
          A function that accepts at least two positional args: a batch
          (represented as DataNode) to be processed, and a parameter of the transformation.
          The function must use DALI operators to process the batch and return a single output
          of such processing.
      mag_range : (number, number) or np.ndarray
          Specifies the range of applicable magnitudes for the operation.
          If the tuple is provided, the magnitudes will be computed as
          `linspace(mag_range[0], mag_range[1], num_magnitude_bins)`.
          If the np.ndarray is provided, it will be used directly instead of the linspace.
          If no `mag_range` is specified, the parameter passed to the `function` will be `None`.
      randomly_negate: bool
          If `True`, the magnitude from the mag_range will be randomly negated for every sample.
      mag_to_param: callable
          A callback that transforms the magnitude into a parameter. The parameter will be passed to
          the decorated operation instead of the plain magnitude. This way, the parameters for the
          range of magnitudes can be computed once in advance and stored as a Constant node.
          Note, the callback must return numpy arrays or data directly convertible to numpy arrays
          (in particular, no pipelined DALI operators can be used in the callback).
          The output type and dimensionality must be consistent and not depend on the magnitude.
      param_device: str
          A "cpu", "gpu", or "auto"; defaults to "cpu". Describes where to store the precomputed
          parameters (i.e. the `mag_to_param` outputs). If "auto" is specified, the CPU or GPU
          backend will be selected to match the `sample`'s backend.
      name: str
          Name of the augmentation. By default, the name of the decorated function is used.
  
      Returns
      -------
      Augmentation
          The operation wrapped with the Augmentation class so that it can be used with the `auto_aug`
          transforms.
      """
  
      def decorator(function):
          cls = augmentation_cls or Augmentation
          return cls(
              function,
              mag_range=mag_range,
              mag_to_param=mag_to_param,
              randomly_negate=randomly_negate,
              param_device=param_device,
              name=name,
          )
  
      if function is None:
          return decorator
      else:
          if not callable(function):
              raise Exception(
                  f"The `@augmentation` decorator was used to decorate the object that "
                  f"is not callable: {function}."
              )
          elif isinstance(function, Augmentation):
              # it's not clear if we should go with "update the setup" or
              # "discard and create" semantics here
              raise Exception(
                  f"The `@augmentation` was applied to already decorated Augmentation. "
                  f"Please call `{function.name}.augmentation` method to modify the augmentation "
                  f"setup or apply the decorator to the underlying `{function.name}.op` directly.\n"
                  f"Error in augmentation: {function}."
              )
          return decorator(function)
  
imports: from typing import Callable, Type, Tuple, Optional, Union
  from nvidia.dali.data_node import DataNode
  from nvidia.dali.auto_aug.core._augmentation import Augmentation
  from numpy import typing
label_desc: Use image node_id 1 for calling a global function with example usage: augmentation(function) and returns: cls, decorator, decorator