Introduction

scipy.optimize.linear_sum_assignment is useful in TensorFlow or Keras when you need an optimal one-to-one matching, but it is important to treat it as a discrete preprocessing step rather than as a differentiable model operation. The usual pattern is to compute a cost matrix in TensorFlow, run the Hungarian assignment in SciPy, and then bring the matched indices back into TensorFlow for the actual loss calculation. That design keeps the assignment logic correct without pretending the assignment itself is part of the gradient graph.

What linear_sum_assignment Does

The function solves the linear sum assignment problem, also called the Hungarian matching problem. Given a cost matrix, it chooses one row-column pairing per item so that total cost is minimized.

A simple SciPy example:

1import numpy as np
2from scipy.optimize import linear_sum_assignment
3
4cost = np.array(
5    [
6        [4.0, 1.0, 3.0],
7        [2.0, 0.0, 5.0],
8        [3.0, 2.0, 2.0],
9    ]
10)
11
12row_idx, col_idx = linear_sum_assignment(cost)
13
14print(row_idx)  # [0 1 2]
15print(col_idx)  # [1 0 2]
16print(cost[row_idx, col_idx].sum())  # 5.0

In machine learning, this appears in object detection, clustering alignment, sequence alignment, and evaluation code where predictions and labels need a one-to-one match.

The Key Constraint in TensorFlow

SciPy runs on NumPy arrays and Python control flow. TensorFlow training graphs run on tensors and differentiable ops. That mismatch leads to the central design rule:

• convert the cost matrix to NumPy
• run linear_sum_assignment
• use the matched indices to gather tensors back inside TensorFlow

That is normal. The Hungarian algorithm makes a discrete combinatorial choice, so gradients flow through the matched tensors, not through the assignment decision itself.

Example: Matching Predictions to Targets

Here is a minimal TensorFlow-friendly pattern:

1import numpy as np
2import tensorflow as tf
3from scipy.optimize import linear_sum_assignment
4
5def hungarian_match(cost_matrix: tf.Tensor):
6    row_idx, col_idx = linear_sum_assignment(cost_matrix.numpy())
7    return tf.convert_to_tensor(row_idx), tf.convert_to_tensor(col_idx)
8
9pred = tf.constant([[0.9, 0.1], [0.2, 0.8], [0.6, 0.4]], dtype=tf.float32)
10target = tf.constant([[1.0, 0.0], [0.0, 1.0], [1.0, 0.0]], dtype=tf.float32)
11
12# Pairwise L1 cost between each prediction and each target
13cost = tf.reduce_sum(
14    tf.abs(pred[:, None, :] - target[None, :, :]),
15    axis=-1,
16)
17
18row_idx, col_idx = hungarian_match(cost)
19
20matched_pred = tf.gather(pred, row_idx)
21matched_target = tf.gather(target, col_idx)
22
23loss = tf.reduce_mean(tf.keras.losses.mse(matched_target, matched_pred))
24print(loss.numpy())

The assignment is computed in SciPy, but the final loss remains a TensorFlow tensor.

Using It Inside Keras Training

If you need this inside a model training loop, custom training code is usually clearer than forcing everything into a stock loss function.

1class MatchingModel(tf.keras.Model):
2    def train_step(self, data):
3        x, y = data
4
5        with tf.GradientTape() as tape:
6            pred = self(x, training=True)
7
8            cost = tf.reduce_sum(
9                tf.abs(pred[:, None, :] - y[None, :, :]),
10                axis=-1,
11            )
12
13            row_idx, col_idx = linear_sum_assignment(cost.numpy())
14
15            matched_pred = tf.gather(pred, row_idx)
16            matched_y = tf.gather(y, col_idx)
17
18            loss = tf.reduce_mean(tf.keras.losses.mse(matched_y, matched_pred))
19
20        grads = tape.gradient(loss, self.trainable_variables)
21        self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
22        return {"loss": loss}

This works best when the matching dimension is small enough that .numpy() conversion is not a bottleneck.

What About tf.numpy_function?

You can wrap SciPy with tf.numpy_function or tf.py_function if you need the call inside a TensorFlow pipeline boundary:

1def assignment_op(cost):
2    row_idx, col_idx = linear_sum_assignment(cost)
3    return row_idx.astype(np.int32), col_idx.astype(np.int32)
4
5row_idx, col_idx = tf.numpy_function(
6    assignment_op,
7    [cost],
8    [tf.int32, tf.int32],
9)

This can make integration easier, but it does not make the operation differentiable. It also makes shape handling and debugging less transparent, so it should be a deliberate integration choice rather than a default.

Common Pitfalls

• Expecting gradients through the assignment itself. Fix: treat matching as a discrete step and backpropagate only through the matched loss.
• Calling .numpy() inside graph-only code. Fix: use eager execution for the matching section or wrap the SciPy call with tf.numpy_function.
• Forgetting the cost matrix must be two-dimensional. Fix: build one assignment problem per example if the data is batched.
• Ignoring Python round-trip overhead for large cost matrices. Fix: measure the NumPy conversion cost before using the pattern at large scale.
• Trying to force the assignment into a stock loss API when the training logic is custom. Fix: use a custom train_step when the matching logic is central to training.

Summary

• 'linear_sum_assignment is useful in TensorFlow when you need one-to-one matching.'
• Run SciPy on a NumPy cost matrix, then feed the matched indices back into TensorFlow.
• Keep the final loss differentiable even though the matching step is not.
• Prefer custom training logic when matching is part of training.
• Use tf.numpy_function only when you need graph integration and understand the shape and gradient trade-offs.