code: def register_augmentation_op(
      name: str, features: Union[str, List[str]]
  ):
      if isinstance(features, str):
          features = [features]
  
      def wrap(cls):
          for feature in features:
              augmentation_op_registry = (
                  get_augmentation_op_registry().get(feature, {})
              )
              augmentation_op_registry[name] = cls
              get_augmentation_op_registry()[
                  feature
              ] = augmentation_op_registry
          return cls
  
      return wrap
  
file_path: ludwig/ludwig/utils/augmentation_utils.py
imports: from typing import Dict, List, Union
  from ludwig.api_annotations import DeveloperAPI
  from ludwig.utils.registry import Registry
label_desc: Use image node_id 2 for calling a global function with example usage: register_augmentation_op(name, features) and returns: wrap, cls