Assignment Operations

def hungarian_algorithm(cost_matrix):
    """
    The Hungarian algorithm for the linear assignment problem.
    
    Solves the problem of finding a minimum-cost assignment given a cost matrix.
    Given an n×m cost matrix, finds an assignment that minimizes the total cost
    while ensuring each row and column is assigned at most once.
    
    Args:
        cost_matrix: A matrix of costs (jax.Array with shape (n, m))
    
    Returns:
        A tuple (i, j) where i is an array of row indices and j is an array 
        of column indices representing the optimal assignment.
        The cost of the assignment is cost_matrix[i, j].sum().
    """

def base_hungarian_algorithm(cost_matrix):
    """
    Base implementation of the Hungarian algorithm.
    
    Lower-level implementation that provides the core Hungarian algorithm
    functionality without additional conveniences.
    
    Args:
        cost_matrix: A matrix of costs (jax.Array with shape (n, m))
    
    Returns:
        A tuple (i, j) where i is an array of row indices and j is an array 
        of column indices representing the optimal assignment.
    """

import optax
import jax.numpy as jnp

# Example 1: Basic assignment problem
cost_matrix = jnp.array([
    [8, 4, 7],
    [5, 2, 3],
    [9, 6, 7],
    [9, 4, 8],
])

# Find optimal assignment
row_indices, col_indices = optax.assignment.hungarian_algorithm(cost_matrix)

# Calculate total cost
total_cost = cost_matrix[row_indices, col_indices].sum()
print(f"Optimal assignment cost: {total_cost}")
print(f"Row assignments: {row_indices}")
print(f"Column assignments: {col_indices}")

# Example 2: Larger assignment problem
cost_matrix = jnp.array([
    [90, 80, 75, 70],
    [35, 85, 55, 65],
    [125, 95, 90, 95],
    [45, 110, 95, 115],
    [50, 100, 90, 100],
])

row_indices, col_indices = optax.assignment.hungarian_algorithm(cost_matrix)
total_cost = cost_matrix[row_indices, col_indices].sum()
print(f"Optimal assignment cost: {total_cost}")

import optax.assignment
# or
from optax.assignment import hungarian_algorithm, base_hungarian_algorithm