code: def hungarian_cover_matcher(
      weights: jnp.ndarray, eps: float = 1e-8
  ) -> jnp.ndarray:
      """Computes the minimum linear sum assignment using the Hungarian algorithm.
  
      Args:
        weights: A float32 [b, n, m] array, where each inner matrix represents
        weights to be use for matching.
        eps: Small number to test for equality to 0.
  
      Returns:
        Jobs and workers matching indices as [b, 2, m] .
      """
      b, n, m = weights.shape
      should_transpose = n > m
      if should_transpose:
          weights = jnp.transpose(weights, axes=(0, 2, 1))
          n, m = m, n
  
      # TODO(agritsenko): Figure out a more efficient way of correctly handling
      # rectangular cost matrices.
      if n != m:  # So n < m based on the code block above.
          pad_n = 1  # `m - n` is guaranteed to be correct, but 1 also works.
          pad_values = (
              jnp.max(weights, axis=(1, 2), keepdims=True) * 1.1
          )
          pad_values = jnp.broadcast_to(pad_values, (b, pad_n, m))
          weights = jnp.concatenate((weights, pad_values), axis=1)
          n += pad_n
      else:
          pad_n = 0
  
      weights = _prepare(weights)
      adj_matrix = jnp.abs(weights) < eps
      state, assignment = _maximum_bipartite_matching(adj_matrix)
      workers_cover, jobs_cover = _compute_cover(state, assignment)
  
      def _cover_incomplete(arg):
          workers_cover, jobs_cover, _, _ = arg
          cover_sum = jnp.sum(workers_cover, dtype=jnp.int32) + jnp.sum(
              jobs_cover, dtype=jnp.int32
          )
          return cover_sum < b * n
  
      def _update_weights_and_match(arg):
          workers_cover, jobs_cover, weights, assignment = arg
          weights = _update_weights_using_cover(
              workers_cover, jobs_cover, weights
          )
          adj_matrix = jnp.abs(weights) < eps
          state, assignment = _maximum_bipartite_matching(
              adj_matrix, assignment
          )
          workers_cover, jobs_cover = _compute_cover(state, assignment)
          return workers_cover, jobs_cover, weights, assignment
  
      (
          workers_cover,
          jobs_cover,
          weights,
          assignment,
      ) = jax.lax.while_loop(
          _cover_incomplete,
          _update_weights_and_match,
          (workers_cover, jobs_cover, weights, assignment),
      )
  
      workers_ind = jnp.broadcast_to(jnp.arange(n), (b, n))
      jobs_ind = jnp.argmax(assignment, axis=2)
  
      # Remove padded indices.
      workers_ind = workers_ind[:, : n - pad_n]
      jobs_ind = jobs_ind[:, : n - pad_n]
  
      if not should_transpose:
          ind = jnp.stack([workers_ind, jobs_ind], axis=1)
      else:
          ind = jnp.stack([jobs_ind, workers_ind], axis=1)
      return ind
  
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 8 for calling a global function with example usage: hungarian_cover_matcher(weights, eps) and returns: ind, unknown, workers_cover,  jobs_cover,  weights,  assignment