code: def check_schema(schema_str: str, func, *args, **kwargs) -> None:
      named_arg_types = schema_str.split(", ")
      num_optional_args = sum(
          [x.endswith("?") for x in named_arg_types]
      )
      min_args = len(named_arg_types) - num_optional_args
  
      # special case: ellipses allows for any number of unchecked args at the end
      if named_arg_types[-1] == "...":
          named_arg_types = named_arg_types[:-1]
      else:
          if not (
              len(args) >= min_args
              and len(args) <= len(named_arg_types)
          ):
              raise ValueError(
                  f"NestedTensor {func.__name__}({schema_str}): expected at least {min_args} "
                  f"arguments and at most {len(named_arg_types)} arguments, but got: "
                  f"{len(args)} arguments"
              )
  
      arg_type_check_fns = {
          "t": lambda x: isinstance(x, torch.Tensor)
          and not isinstance(x, NestedTensor),
          "jt": lambda x: isinstance(x, NestedTensor)
          and x._lengths is None
          and x._ragged_idx
          == 1,  # ops with "jt" require contiguous JT only
          "jt_all": lambda x: isinstance(
              x, NestedTensor
          ),  # ops with "jt_all" can accept all kinds of JT
          "any": lambda x: True,
      }
      for i, named_arg_type in enumerate(named_arg_types):
          name, arg_type = named_arg_type.split(": ")
          is_optional = arg_type.endswith("?")
          normalized_arg_type = (
              arg_type[:-1] if is_optional else arg_type
          )
          if normalized_arg_type not in arg_type_check_fns.keys():
              raise AssertionError(
                  f"Unknown arg type: {normalized_arg_type}"
              )
  
          if i >= len(args):
              if not is_optional:
                  raise ValueError(
                      f"NestedTensor {func.__name__}({schema_str}) "
                      f"missing required argument: {name}"
                  )
              continue
  
          if not arg_type_check_fns[normalized_arg_type](args[i]):
              type_to_desc = {
                  "t": "tensor",
                  "jt": "contiguous jagged layout NestedTensor",
                  "jt_all": "jagged layout NestedTensor",
                  "any": "<any type>",
              }
  
              raise ValueError(
                  f"NestedTensor {func.__name__}({schema_str}): expected {name} to be a "
                  f"{type_to_desc[arg_type]}"
              )
  
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 4 for calling a global function with example usage: check_schema(schema_str, func) without return types