code: def custom_sharding_spec_op(sharding_spec_class, func):
      """
      Decorator to allow custom registration of ops.
      Args:
          sharding_spec_class(type): The ShardingSpec for which we need to add this custom op.
          func(Callable): The op to override (ex: torch.bmm)
      """
      class_name = sharding_spec_class.__qualname__
      if class_name not in _CUSTOM_SHARDING_SPEC_OPS:
          _CUSTOM_SHARDING_SPEC_OPS[class_name] = {}
      return functools.partial(
          _decorator_func,
          op=func,
          op_table=_CUSTOM_SHARDING_SPEC_OPS[class_name],
      )
  
imports: from abc import ABC, abstractmethod
  from dataclasses import dataclass
  import functools
  from typing import Callable, Dict, List, TYPE_CHECKING
  import torch
  from ._internals import check_tensor, get_chunked_dim_size, get_split_size, validate_non_overlapping_shards_metadata
  from torch.distributed._shard.metadata import ShardMetadata
  import torch.distributed._shard.sharded_tensor.metadata
  from torch.distributed._shard.op_registry_utils import _decorator_func
label_desc: Use image node_id 3 for calling a global function with example usage: custom_sharding_spec_op(sharding_spec_class, func) and returns: functools