code: def _wrap_jagged_dim(ndim, dim, op_name, convert_to_inner_dim=True):
      from torch._prims_common import canonicalize_dims
  
      wrapped = canonicalize_dims(ndim, dim)
      if wrapped < 2:
          raise RuntimeError(
              f"{op_name}(): not supported for NestedTensor on dim=0 or dim=1"
          )
      return (
          _outer_to_inner_dim(ndim, wrapped)
          if convert_to_inner_dim
          else wrapped
      )
  
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 2 for calling a global function with example usage: _wrap_jagged_dim(ndim, dim, op_name, convert_to_inner_dim) and returns: unknown