code: def mean_dim(func, *args, **kwargs):
      _, new_kwargs = normalize_function(
          func,
          args=args,
          kwargs=kwargs,
          normalize_to_only_use_kwargs=True,
      )
  
      inp = new_kwargs.pop("input")
      # NB: mean expects dim as a single item list of ints for some reason
      new_kwargs["dim"] = [
          _wrap_jagged_dim(inp.dim(), new_kwargs["dim"][0], "mean")
      ]
  
      return NestedTensor(
          func(inp._values, **new_kwargs), **extract_kwargs(inp)
      )
  
imports: import functools
  import math
  import operator
  import torch
  from torch.nested._internal.sdpa import jagged_scaled_dot_product_attention
  from .nested_tensor import NestedTensor
  from typing import *
  from torch.fx.operator_schemas import normalize_function
label_desc: Use image node_id 46 for calling a global function with example usage: mean_dim(func) and returns: NestedTensor