code: def _wrap_jagged_dims(ndim, dims, op_name):
      # ex: (2, 3, 4) -> (1, 2, 3)
      # ex: (0, 1, 4) -> (0, 3)
      from torch._prims_common import canonicalize_dims
  
      wrapped_dims = [canonicalize_dims(ndim, d) for d in dims]
      # This logic needs to be done after we canonicalize dims but before we
      # map to inner dims so we can print a nicer error message.
      zero_in_dims = 0 in wrapped_dims
      one_in_dims = 1 in wrapped_dims
      if zero_in_dims ^ one_in_dims:
          apply, not_apply = (
              ("batch", "ragged")
              if zero_in_dims
              else ("ragged", "batch")
          )
          raise RuntimeError(
              f"{op_name}(): applying over the {apply} dimension, but not the {not_apply}"
              " dimension is not supported for NestedTensor"
          )
      return (
          tuple(_outer_to_inner_dim(ndim, d) for d in dims if d != 0),
          zero_in_dims,
      )
  
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 3 for calling a global function with example usage: _wrap_jagged_dims(ndim, dims, op_name) and returns: tuple,  zero_in_dims