code: def jagged_binary_pointwise(func, *args, **kwargs):
      a, b = args[0], args[1]
      assert isinstance(a, NestedTensor) or isinstance(b, NestedTensor)
  
      mismatch_error_msg = "cannot call binary pointwise function {} with inputs of shapes {} and {}"
      # a is NT, b is NT
      if isinstance(a, NestedTensor) and isinstance(b, NestedTensor):
          # ex: (B, j0, D) + (B, j0, D)
          # ex: (B, j0, D) + (B, j0, 1)
          if raggedness_matches(a, b._size):
              return NestedTensor(
                  func(a._values, b._values, *args[2:], **kwargs),
                  **extract_kwargs(a),
              )
          raise RuntimeError(
              mismatch_error_msg.format(func.__name__, a._size, b._size)
          )
      # either a is NT or b is NT at this point
      a_is_nt = isinstance(a, NestedTensor)
      extracted_kwargs = (
          extract_kwargs(a) if a_is_nt else extract_kwargs(b)
      )
  
      # === Handle broadcasting across the batch / ragged dims ===
  
      # Easy case: take advantage of pre-existing broadcasting logic
      # ex: (B, j0, ?, ?) + (?) -> (B, j0, ?, ?)
      # ex: (B, j0, ?, ?) + (?, ?) -> (B, j0, ?, ?)
      # ex: (B, j0, ?, ?) + (1, 1, ?, ?) -> (B, j0, ?, ?)
      nt, t = (a, b) if a_is_nt else (b, a)
      # See Note: [ Squeezing leading ones ]
      if t.dim() > nt.dim():
          raise NotImplementedError(
              "NYI: broadcasting NT with T with larger dim"
          )
      t_squeezed = squeeze_leading_ones(t)
      if nt.dim() >= t_squeezed.dim() + 2:
          lhs, rhs = (
              (nt._values, t_squeezed)
              if a_is_nt
              else (t_squeezed, nt._values)
          )
          return NestedTensor(
              func(lhs, rhs, *args[2:], **kwargs), **extracted_kwargs
          )
  
      # Harder case: do manual broadcasting over unbound components
      # when NT dim == non-NT dim
      # ex: (B, j0, D_0, D_1) + (B, 1, D_0, D_1) -> (B, j0, D_0, D_1)
      if a.dim() == b.dim():
          # ex: (B, j0, D_0, D_1) + (1, 1, D_0, D_1) -> should
          # be (B, j0, D_0, D_1) but not yet supported
          if a.shape[0] != b.shape[0]:
              raise RuntimeError(
                  mismatch_error_msg.format(
                      func.__name__, a.shape, b.shape
                  )
              )
  
          # need to use offsets to broadcast across ragged dim properly
          # NB: inefficient fallback here; Triton codegen can help this
          # TODO: Make this work with autograd
          outputs = []
          for a_comp, b_comp in zip(a.unbind(), b.unbind()):
              outputs.append(func(a_comp, b_comp, *args[2:], **kwargs))
          new_values = torch.cat(outputs, dim=0)
          return NestedTensor(new_values, **extracted_kwargs)
  
      # ex: (B, j0, D_0, D_1) + (A, B, 1, D_0, D_1) -> error because this breaks the invariant
      # that ragged dim is wrt left-most batch dim
      raise RuntimeError(
          mismatch_error_msg.format(func.__name__, a.shape, b.shape)
      )
  
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 12 for calling a global function with example usage: jagged_binary_pointwise(func) and returns: NestedTensor, NestedTensor, NestedTensor