code: def _maximum_bipartite_matching(
      adj_matrix: jnp.ndarray, assignment: Optional[jnp.ndarray] = None
  ) -> Tuple[Dict[str, jnp.ndarray], jnp.ndarray]:
      """Performs maximum bipartite matching using augmented paths.
  
      Args:
        adj_matrix: A bool [b, n, m] array, where each element of the inner matrix
          represents whether the worker (row) can be matched to the job (column).
        assignment: An optional bool [b, n, m] array, where each element of the
          inner matrix represents whether the worker has been matched to the job.
          This may be a partial assignment. If specified, this assignment will be
          used to seed the iterative algorithm.
  
      Returns:
        A state dict representing the final augmenting path state search, and
        a maximum bipartite matching assignment array. Note that the state outcome
        can be used to compute a minimum vertex cover for the bipartite graph.
      """
      if assignment is None:
          assignment = _greedy_assignment(adj_matrix)
      state = _find_augmenting_path(assignment, adj_matrix)
  
      def _has_new_jobs(arg):
          state, _ = arg
          return jnp.any(state["new_jobs"])
  
      def _improve_assignment_and_find_new_path(arg):
          state, assignment = arg
          assignment = _improve_assignment(assignment, state)
          state = _find_augmenting_path(assignment, adj_matrix)
          return state, assignment
  
      state, assignment = jax.lax.while_loop(
          _has_new_jobs,
          _improve_assignment_and_find_new_path,
          (state, assignment),
      )
  
      return state, assignment
  
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 5 for calling a global function with example usage: _maximum_bipartite_matching(adj_matrix, assignment) and returns: state,  assignment, jnp, state,  assignment