code: def cpu_matcher(matching_fn):
      """Wraps matching function to be usable within jitted functions.
  
      Args:
        matching_fn: function; A matching function that aligns the predictions of
          the model with targets.
  
      Returns:
        Matching function with host callback that can be jitted.
      """
  
      # The callback function can only take a single argument.
      def slice_and_match(args):
          cost, ncol = args
          return slicer(cost, ncol, matching_fn)
  
      @jax.custom_vjp
      def matching_fn_hcb(cost, n_cols=None):
          *b, n, m = cost.shape
          return jax.pure_callback(
              slice_and_match,
              jax.ShapeDtypeStruct(b + [2, min(n, m)], jnp.int32),
              (cost, n_cols),
              vectorized=True,
          )
  
      # Define forward and backward passes.
      def matching_fn_hcb_vjp_fwd(cost, n_cols):
          return matching_fn_hcb(cost, n_cols), None
  
      def matching_fn_hcb_vjp_bwd(*_):
          return (None,)  # Return no gradient.
  
      matching_fn_hcb.defvjp(
          matching_fn_hcb_vjp_fwd, matching_fn_hcb_vjp_bwd
      )
  
      # Note: When called from TPU, errors in the callback will NOT cause the code
      # to fail, but will produce nonsensical matching outputs. Test the callback
      # carefully with jax.jit(callback, backend='cpu') first.
      return matching_fn_hcb
  
imports: import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: cpu_matcher(matching_fn) and returns: matching_fn_hcb, slicer, jax, matching_fn_hcb,  None, None