code: def _find_augmenting_path(
      assignment: jnp.ndarray, adj_matrix: jnp.ndarray
  ) -> Dict[str, jnp.ndarray]:
      """Finds an augmenting path given an assignment and an adjacency matrix.
  
      The augmenting path search starts from the unassigned workers, then goes on
      to find jobs (via an unassigned pairing), then back again to workers (via an
      existing pairing), and so on. The path alternates between unassigned and
      existing pairings. Returns the state after the search.
  
      Note: In the state the worker and job, indices are 1-indexed so that we can
      use 0 to represent unreachable nodes. State contains the following keys:
  
      - jobs: A [b, 1, m] array containing the highest index unassigned worker that
        can reach this job through a path.
      - jobs_from_worker: A [b, n] array containing the worker reached immediately
        before this job.
      - workers: A [b, n, 1] array containing the highest index unassigned worker
        that can reach this worker through a path.
      - workers_from_job: A [b, m] array containing the job reached immediately
        before this worker.
      - new_jobs: A bool [b, m] array containing True if the unassigned job can be
        reached via a path.
  
      State can be used to recover the path via backtracking.
  
      Args:
        assignment: A 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.
        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).
  
      Returns:
        A state dictionary, which represents the outcome of running an augmenting
        path search on the graph given the assignment.
      """
      b, n, m = assignment.shape
      unassigned_workers = jnp.logical_not(
          jnp.any(assignment, axis=2, keepdims=True)
      )
      unassigned_jobs = jnp.logical_not(
          jnp.any(assignment, axis=1, keepdims=True)
      )
  
      unassigned_pairings = jnp.logical_and(
          adj_matrix, jnp.logical_not(assignment)
      ).astype(jnp.int32)
      existing_pairings = assignment.astype(jnp.int32)
  
      # Initialize unassigned workers to have non-zero ids, assigned workers will
      # have ids = 0.
      worker_indices = jnp.arange(1, n + 1, dtype=jnp.int32)
      init_workers = jnp.tile(
          worker_indices[jnp.newaxis, :, jnp.newaxis], (b, 1, 1)
      )
      init_workers = init_workers * unassigned_workers.astype(jnp.int32)
  
      state = {
          "jobs": jnp.zeros((b, 1, m), dtype=jnp.int32),
          "jobs_from_worker": jnp.zeros((b, m), dtype=jnp.int32),
          "workers": init_workers,
          "workers_from_job": jnp.zeros((b, n), dtype=jnp.int32),
      }
  
      def _has_active_workers(arg):
          """Check if there are still active workers."""
          _, cur_workers = arg
          return jnp.sum(cur_workers) > 0
  
      def _augment_step(arg):
          """Performs one search step."""
          state, cur_workers = arg
  
          # Find potential jobs using current workers.
          potential_jobs = cur_workers * unassigned_pairings
          curr_jobs = jnp.max(potential_jobs, axis=1, keepdims=True)
          curr_jobs_from_worker = 1 + jnp.argmax(potential_jobs, axis=1)
  
          # Remove already accessible jobs from curr_jobs.
          default_jobs = jnp.zeros_like(state["jobs"])
          curr_jobs = jnp.where(
              state["jobs"] > 0, default_jobs, curr_jobs
          )
          curr_jobs_from_worker = (
              curr_jobs_from_worker
              * (curr_jobs > 0).astype(jnp.int32)[:, 0, :]
          )
  
          # Find potential workers from current jobs.
          potential_workers = curr_jobs * existing_pairings
          cur_workers = jnp.max(
              potential_workers, axis=2, keepdims=True
          )
          cur_workers_from_job = 1 + jnp.argmax(
              potential_workers, axis=2
          )
  
          # Remove already accessible workers from cur_workers.
          default_workers = jnp.zeros_like(state["workers"])
          cur_workers = jnp.where(
              state["workers"] > 0, default_workers, cur_workers
          )
          cur_workers_from_job = (
              cur_workers_from_job
              * (cur_workers > 0).astype(jnp.int32)[:, :, 0]
          )
  
          # Update state so that we can backtrack later.
          state["jobs"] = jnp.maximum(state["jobs"], curr_jobs)
          state["jobs_from_worker"] = jnp.maximum(
              state["jobs_from_worker"], curr_jobs_from_worker
          )
          state["workers"] = jnp.maximum(state["workers"], cur_workers)
          state["workers_from_job"] = jnp.maximum(
              state["workers_from_job"], cur_workers_from_job
          )
  
          return state, cur_workers
  
      state, _ = jax.lax.while_loop(
          _has_active_workers, _augment_step, (state, init_workers)
      )
  
      # Compute new jobs, this is useful for determnining termnination of the
      # maximum bi-partite matching and initialization for backtracking.
      new_jobs = jnp.logical_and(state["jobs"] > 0, unassigned_jobs)
      state["new_jobs"] = new_jobs[:, 0, :]
      return state
  
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 3 for calling a global function with example usage: _find_augmenting_path(assignment, adj_matrix) and returns: state, unknown, state,  cur_workers