google-deepmind · copybara-service · Oct 8, 2024 · Oct 4, 2024
diff --git a/docs/api/assignment.rst b/docs/api/assignment.rst
@@ -0,0 +1,12 @@
+Assignment problem
+==================
+
+.. currentmodule:: optax.assignment
+
+.. autosummary::
+    hungarian_algorithm
+
+
+Hungarian algorithm
+~~~~~~~~~~~~~~~~~~~
+.. autofunction:: hungarian_algorithm
diff --git a/docs/gallery.rst b/docs/gallery.rst
@@ -189,6 +189,27 @@
     </div>
 
 
+.. raw:: html
+
+    <div class="sphx-glr-thumbcontainer" tooltip="Solving the linear assignment problem.">
+
+.. only:: html
+
+  .. image:: /images/examples/linear_assignment_problem.png
+    :alt:
+
+  :doc:`_collections/examples/linear_assignment_problem`
+
+.. raw:: html
+
+      <div class="sphx-glr-thumbnail-title">Solving the linear assignment problem.</div>
+    </div>
+
+.. raw:: html
+
+    </div>
+
+
 Contrib Examples
 ----------------
 

diff --git a/docs/images/examples/linear_assignment_problem.png b/docs/images/examples/linear_assignment_problem.png
diff --git a/docs/index.rst b/docs/index.rst
@@ -54,6 +54,7 @@ for instructions on installing JAX.
    :caption: 📖 Reference
    :maxdepth: 2
 
+   api/assignment
    api/optimizers
    api/transformations
    api/combining_optimizers

diff --git a/examples/linear_assignment_problem.ipynb b/examples/linear_assignment_problem.ipynb
diff --git a/optax/__init__.py b/optax/__init__.py
@@ -17,6 +17,7 @@
 # pylint: disable=wrong-import-position
 # pylint: disable=g-importing-member
 
+from optax import assignment
 from optax import contrib
 from optax import losses
 from optax import monte_carlo
@@ -306,6 +307,7 @@
     "apply_updates",
     "ApplyEvery",
     "ApplyIfFiniteState",
+    "assignment",
     "centralize",
     "chain",
     "clip_by_block_rms",

diff --git a/optax/assignment/__init__.py b/optax/assignment/__init__.py
@@ -0,0 +1,19 @@
+# Copyright 2024 DeepMind Technologies Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""The assignment sub-package."""
+
+# pylint:disable=g-importing-member
+
+from optax.assignment._hungarian_algorithm import hungarian_algorithm