code: def _prepare(weights: jnp.ndarray) -> jnp.ndarray:
      """Prepare the cost matrix.
  
      To speed up computational efficiency of the algorithm, all weights are shifted
      to be non-negative. Each element is reduced by the row / column minimum. Note
      that neither operation will effect the resulting solution but will provide a
      better starting point for the greedy assignment. Note this corresponds to the
      pre-processing and step 1 of the Hungarian algorithm from Wikipedia.
  
      Args:
        weights: A float32 [b, n, m] array, where each inner matrix represents
        weights to be use for matching.
  
      Returns:
        A prepared weights array of the same shape and dtype.
      """
      # Since every worker needs a job and every job needs a worker, we can subtract
      # the minimum from each.
      assert weights.ndim == 3
  
      weights = weights - jnp.min(weights, axis=2, keepdims=True)
      weights = weights - jnp.min(weights, axis=1, keepdims=True)
      return weights
  
file_path: scenic/scenic/model_lib/matchers/hungarian_cover.py
imports: from typing import Tuple, Optional, Dict
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: _prepare(weights) and returns: weights