code: def basic_unit(cls: T, basic_unit_tag: bool = True) -> T:
      """
      To wrap a module as a basic unit, is to make it a primitive and stop the engine from digging deeper into it.
  
      ``basic_unit_tag`` is true by default. If set to false, it will not be explicitly mark as a basic unit, and
      graph parser will continue to parse. Currently, this is to handle a special case in ``nn.Sequential``.
  
      Although ``basic_unit`` calls ``trace`` in its implementation, it is not for serialization. Rather, it is meant
      to capture the initialization arguments for mutation. Also, graph execution engine will stop digging into the inner
      modules when it reaches a module that is decorated with ``basic_unit``.
  
      .. code-block:: python
  
          @basic_unit
          class PrimitiveOp(nn.Module):
              ...
      """
  
      # Internal flag. See nni.trace
      nni_trace_flag = os.environ.get("NNI_TRACE_FLAG", "")
      if nni_trace_flag.lower() == "disable":
          return cls
  
      if _check_wrapped(cls, "basic_unit"):
          return cls
  
      import torch.nn as nn
  
      assert issubclass(cls, nn.Module), "When using @basic_unit, the class must be a subclass of nn.Module."  # type: ignore
  
      cls = trace(cls)
      cls._nni_basic_unit = basic_unit_tag  # type: ignore
  
      _torchscript_patch(cls)
  
      return cls
  
file_path: nni/nni/nas/utils/_legacy.py
imports: import inspect
  import os
  import warnings
  from typing import Any, TypeVar, Type
  from nni.common.serializer import is_traceable, is_wrapped_with_trace, trace
label_desc: Use image node_id 4 for calling a global function with example usage: basic_unit(cls, basic_unit_tag) and returns: cls, cls, cls