code: def slicer(cost, n_present_col, matching_fn):
      """Maps matching_fn over examples after removing padding to speed up matching.
  
      Args:
        cost: Cost matrix or batch of cost matrices with any number of batch
          dimensions. Requires n_row >= n_col.
        n_present_col: Number of non-padding columns of the cost matrices, or None
          if padding should not be removed.
        matching_fn: A matching function that operates on a single cost matrix.
  
      Returns:
        Matchings of shape [batch, 2, n_col].
  
      Raises:
        ValueError if n_row < n_col and n_present_col is not None.
      """
      batch_shape = cost.shape[:-2]
      cost = cost.reshape(-1, *cost.shape[-2:])
  
      if n_present_col is None:
          matches = np.stack([matching_fn(c) for c in cost])
          return matches.reshape(*batch_shape, *matches.shape[1:])
  
      n_present_col = n_present_col.reshape(-1)
      assert cost.shape[:1] == n_present_col.shape, (
          cost.shape,
          n_present_col.shape,
      )
  
      batch, n_row, n_col = cost.shape
      if n_row < n_col:
          raise ValueError(
              f"Slicer requires that n_row ({n_row}) >= n_col ({n_col})."
          )
  
      eye = np.eye(n_row, dtype=bool)
      matches = []
      for i in range(batch):
          present_col = max(
              n_present_col[i], 1
          )  # One col even if all are padded.
          cost_m = cost[
              i, :, :present_col
          ]  # Slicing should avoid a copy.
  
          row, col = matching_fn(cost_m)
  
          # Add padded matches (if padding was done correctly these can be random).
          unmatched_row = np.where(~eye[row].max(axis=0))[
              0
          ]  # Faster than setdiff1d.
          unmatched_row = unmatched_row.astype(np.int32)
          unmatched_col = np.arange(present_col, n_col, dtype=np.int32)
  
          # Assume n_row >= n_col >= n_present_col.
          n_common = n_col - present_col
          unmatched_row = unmatched_row[:n_common]
  
          # Reconstruct the matching.
          row = np.concatenate([row, unmatched_row], axis=0)
          col = np.concatenate([col, unmatched_col], axis=0)
  
          matches.append(np.stack([row, col], axis=0))
  
      matches = np.stack(matches)
  
      return matches.reshape(*batch_shape, *matches.shape[1:])
  
imports: import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: slicer(cost, n_present_col, matching_fn) and returns: matches, matches