aiaccel.torch.functional package#

Submodules#

aiaccel.torch.functional.linear_sum_assignment module#

aiaccel.torch.functional.linear_sum_assignment.linear_sum_assignment(cost_matrix: Tensor, maximize: bool = False) → tuple[Tensor, Tensor][source]#

Solve the linear sum assignment problem for a batch of cost matrices.

Parameters:

  • cost_matrix (torch.Tensor) – A tensor of shape (…, m, n) representing the cost matrix for each assignment problem.

  • maximize (bool) – If True, the problem is treated as a maximization problem. If False, it is treated as a minimization problem. Defaults to False.

Returns:

A tuple containing two tensors:

  • row_indices: Indices of the rows assigned to each column.

  • col_indices: Indices of the columns assigned to each row.

Return type:

tuple

Module contents#

aiaccel.torch.functional.linear_sum_assignment(cost_matrix: Tensor, maximize: bool = False) → tuple[Tensor, Tensor][source]#

Solve the linear sum assignment problem for a batch of cost matrices.

Parameters:

  • cost_matrix (torch.Tensor) – A tensor of shape (…, m, n) representing the cost matrix for each assignment problem.

  • maximize (bool) – If True, the problem is treated as a maximization problem. If False, it is treated as a minimization problem. Defaults to False.

Returns:

A tuple containing two tensors:

  • row_indices: Indices of the rows assigned to each column.

  • col_indices: Indices of the columns assigned to each row.

Return type:

tuple