code: def lazy_matcher(cost: jnp.ndarray) -> jnp.ndarray:
      """Computes lazy Matching on cost matrix for a batch of datapoints.
  
      This matcher ignores input and matches i-th to i-th input.
  
      Args:
        cost: Cost matrix for the matching of shape [B, N, M].
  
      Returns:
        An assignment of size [B, 2, min(N,M)].
      """
      batch_size, n, m = cost.shape
      return jax.lax.broadcast(
          jnp.arange(0, min(n, m)), (batch_size, 2)
      )
  
label_desc: Use image node_id 1 for calling a global function with example usage: lazy_matcher(cost) and returns: jax