code: def jagged_torch_function(func, *args, **kwargs):
      # SDPA has special kernels that handle nested tensors.
      # Dispatch to the correct implementation here
      if func is torch._C._nn.scaled_dot_product_attention:
          return jagged_scaled_dot_product_attention(*args, **kwargs)
  
      # Handle flatten() here because it's CompositeImplicit.
      if func.__name__ == "flatten":
  
          def _flatten_sig(input, start_dim=0, end_dim=-1):
              pass
  
          _, new_kwargs = normalize_function(
              _flatten_sig,
              args=args,
              kwargs=kwargs,
              normalize_to_only_use_kwargs=True,
          )
  
          inp = new_kwargs.pop("input")
  
          # NB: stay in outer dim space because we're going to redispatch on a NT input
          start_dim = _wrap_jagged_dim(
              inp.dim(),
              new_kwargs["start_dim"],
              "flatten",
              convert_to_inner_dim=False,
          )
          end_dim = _wrap_jagged_dim(
              inp.dim(),
              new_kwargs["end_dim"],
              "flatten",
              convert_to_inner_dim=False,
          )
  
          if start_dim == end_dim:
              return inp
  
          product = functools.reduce(
              operator.mul, inp.shape[start_dim : end_dim + 1]
          )
          new_shape = (
              *inp.shape[:start_dim],
              product,
              *inp.shape[end_dim + 1 :],
          )
  
          return inp.reshape(*new_shape)
  
      raise NotImplementedError(func)
  
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 13 for calling a global function with example usage: jagged_torch_function(func) and returns: jagged_scaled_dot_product_attention, inp, inp