code: def map_only(__type_or_types: TypeAny) -> MapOnlyFn[FnAny[Any]]:
      """
      Suppose you are writing a tree_map over tensors, leaving everything
      else unchanged.  Ordinarily you would have to write:
  
          def go(t):
              if isinstance(t, Tensor):
                  return ...
              else:
                  return t
  
      With this function, you only need to write:
  
          @map_only(Tensor)
          def go(t):
              return ...
  
      You can also directly use 'tree_map_only'
      """
  
      def wrapper(func: Callable[[T], Any]) -> Callable[[Any], Any]:
          @functools.wraps(func)
          def wrapped(x: T) -> Any:
              if isinstance(x, __type_or_types):
                  return func(x)
              return x
  
          return wrapped
  
      return wrapper
  
imports: import functools
  import warnings
  from typing import Any, Callable, Iterable, List, Optional, overload, Tuple, Type, TypeVar, Union
  import torch
  import optree
  from optree import PyTreeSpec
label_desc: Use image node_id 15 for calling a global function with example usage: map_only(__type_or_types) and returns: wrapper, wrapped, x, func