code: def with_cupy_rmm(func):
      """
  
      Decorator to call CuPy functions with RMM memory management. Use it
      to decorate any function that will call CuPy functions. This will ensure
      that those calls use RMM for memory allocation instead of the default
      CuPy pool. Example:
  
      .. code-block:: python
  
          @with_cupy_rmm
          def fx(...):
              a = cp.arange(10) # uses RMM for allocation
  
      """
  
      if func.__dict__.get("__cuml_rmm_wrapped", False):
          return func
  
      @wraps(func)
      def cupy_rmm_wrapper(*args, **kwargs):
          if GPU_ENABLED:
              with cupy_using_allocator(rmm_cupy_allocator):
                  return func(*args, **kwargs)
          return func(*args, **kwargs)
  
      # Mark the function as already wrapped
      cupy_rmm_wrapper.__dict__["__cuml_rmm_wrapped"] = True
  
      return cupy_rmm_wrapper
  
imports: import functools
  import operator
  import re
  from functools import wraps
  from cuml.internals.global_settings import GlobalSettings
  from cuml.internals.device_support import GPU_ENABLED
  from cuml.internals.mem_type import MemoryType
  from cuml.internals.output_type import INTERNAL_VALID_OUTPUT_TYPES, VALID_OUTPUT_TYPES
  from cuml.internals.safe_imports import cpu_only_import_from, gpu_only_import_from, UnavailableNullContext
label_desc: Use image node_id 2 for calling a global function with example usage: with_cupy_rmm(func) and returns: cupy_rmm_wrapper, func, func, func