Equimo 2.3.0__tar.gz → 2.4.0__tar.gz

{equimo-2.3.0 → equimo-2.4.0}/CHANGELOG.md

@@ -2,6 +2,15 @@
 
 All notable changes to Equimo are documented here. Equimo follows Semantic Versioning from version 2.0.0 onward.
 
+## [2.4.0] - 2026-09-24
+
+### Added
+
+- JAX-native query-segmentation assignment, validity-aware point sampling, class/mask/Dice losses, auxiliary-output
+  composition, and inspectable batch/device reductions, with pinned upstream numerical references.
+- Padding-aware PMT BatchNorm and configurable collective axes for synchronized statistics, preserving default model
+  checkpoint layouts.
+
 ## [2.3.0] - 2026-09-24
 
 ### Added

{equimo-2.3.0 → equimo-2.4.0}/NOTICE

@@ -23,6 +23,18 @@ Source: https://github.com/tensorflow/models/blob/b41a080d61cd4298fea9894804e965
 Equimo contributors modified the upstream code's formatting and type
 annotations.
 
+Query assignment
+----------------
+
+src/equimo/vision/query_training/_assignment.py adapts the Hungarian algorithm
+from Optax, licensed under Apache-2.0 (LICENSES/Apache-2.0.txt).
+Copyright 2024 DeepMind Technologies Limited. All Rights Reserved.
+Source: https://github.com/google-deepmind/optax/blob/225a7079f4630bf75bee94ec78de9aa69c60fab3/optax/assignment/_hungarian_algorithm.py
+
+Equimo contributors added target-slot validity, bounded loops, numerical-failure reporting,
+fixed-capacity target-ordered outputs, explicit float32 arithmetic and stopped
+gradients, and documented deterministic tie handling.
+
 LingBot-Vision
 -------------
 

{equimo-2.3.0 → equimo-2.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Equimo
-Version: 2.3.0
+Version: 2.4.0
 Summary: Implementation of modern multimodal models in JAX/Equinox
 License-Expression: MIT AND Apache-2.0
 License-File: LICENSE.md

{equimo-2.3.0 → equimo-2.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "Equimo"
-version = "2.3.0"
+version = "2.4.0"
 description = "Implementation of modern multimodal models in JAX/Equinox"
 readme = "README.md"
 requires-python = ">=3.12,<3.15"

equimo-2.4.0/src/equimo/__init__.py

@@ -0,0 +1 @@
+__version__ = "2.4.0"

{equimo-2.3.0 → equimo-2.4.0}/src/equimo/vision/models/pmd.py

@@ -29,6 +29,7 @@ from equimo.vision.segmentation import (
 
 
 _MISSING = object()
+NormOption = float | bool | str | tuple[str, ...]
 
 
 @dataclass(frozen=True)
@@ -57,7 +58,7 @@ class PMTConfig:
     num_blocks: int = 6
     norm_layer: str = "groupnorm"
     norm_max_group: int = 32
-    norm_kwargs: tuple[tuple[str, float | bool], ...] | Mapping[str, float | bool] = ()
+    norm_kwargs: tuple[tuple[str, NormOption], ...] | Mapping[str, NormOption] = ()
     masked_attention: bool = True
 
     def __post_init__(self) -> None:
@@ -96,7 +97,7 @@ class PMTConfig:
             raise ValueError("Unsupported PMT lateral norm_layer.")
         allowed = {
             "groupnorm": {"eps", "channelwise_affine"},
-            "batchnorm": {"eps", "momentum"},
+            "batchnorm": {"eps", "momentum", "axis_name"},
             "layernorm": {"eps", "use_weight", "use_bias"},
             "rmsnorm": {"eps", "use_weight", "use_bias"},
             "none": set(),
@@ -179,11 +180,14 @@ def anneal_pmt_mask_state(
 
 
 class ReferenceBatchNorm(eqx.nn.StatefulLayer):
-    """Single-host counterpart of the author's lateral SyncBatchNorm.
+    """Padding-aware counterpart of the author's lateral SyncBatchNorm.
 
     Call over one image inside ``jax.vmap(..., axis_name="pmt_batch")`` in
     training. Statistics cover batch and tokens, including prefix tokens.
     Inference uses the stored running mean and unbiased running variance.
+    Supply multiple axis names to synchronize over a local batch and devices.
+    False ``example_valid`` excludes an image from moments and counts. A global
+    batch with no valid observations returns zeros and leaves state unchanged.
     """
 
     weight: jax.Array
@@ -193,8 +197,16 @@ class ReferenceBatchNorm(eqx.nn.StatefulLayer):
     count_index: eqx.nn.StateIndex
     eps: float = eqx.field(static=True)
     momentum: float = eqx.field(static=True)
+    axis_name: str | tuple[str, ...] = eqx.field(static=True)
 
-    def __init__(self, dim: int, *, eps: float = 1e-5, momentum: float = 0.1):
+    def __init__(
+        self,
+        dim: int,
+        *,
+        eps: float = 1e-5,
+        momentum: float = 0.1,
+        axis_name: str | tuple[str, ...] = "pmt_batch",
+    ):
         if dim <= 0 or eps <= 0 or not 0 < momentum <= 1:
             raise ValueError("Invalid reference BatchNorm dimensions or options.")
         self.weight = jnp.ones((dim,), dtype=jnp.float32)
@@ -204,39 +216,70 @@ class ReferenceBatchNorm(eqx.nn.StatefulLayer):
         self.count_index = eqx.nn.StateIndex(jnp.asarray(0, dtype=jnp.int32))
         self.eps = eps
         self.momentum = momentum
+        names = (axis_name,) if isinstance(axis_name, str) else axis_name
+        if (
+            not isinstance(names, tuple)
+            or not names
+            or any(not isinstance(name, str) or not name for name in names)
+            or len(set(names)) != len(names)
+        ):
+            raise ValueError(
+                "BatchNorm axis_name requires nonempty distinct axis names."
+            )
+        self.axis_name = axis_name
 
     def __call__(
-        self, x: jax.Array, state: eqx.nn.State, *, inference: bool
+        self,
+        x: jax.Array,
+        state: eqx.nn.State,
+        *,
+        inference: bool,
+        example_valid: bool | jax.Array = True,
     ) -> tuple[jax.Array, eqx.nn.State]:
         if x.ndim != 2 or x.shape[0] != self.weight.shape[0]:
             raise ValueError("Reference BatchNorm expects (channels, tokens).")
-        value = x.astype(jnp.float32)
+        valid = jnp.asarray(example_valid)
+        if valid.shape != () or valid.dtype != jnp.bool_:
+            raise ValueError("BatchNorm example_valid must be a boolean scalar.")
+        value = jnp.where(valid, x, 0).astype(jnp.float32)
         if inference:
             mean = state.get(self.mean_index)
             variance = state.get(self.variance_index)
         else:
-            if x.shape[1] * lax.axis_size("pmt_batch") < 2:
+            if example_valid is True and x.shape[1] * lax.axis_size(self.axis_name) < 2:
                 raise ValueError("BatchNorm requires at least two observations.")
-            mean = lax.pmean(jnp.mean(value, axis=1), "pmt_batch")
-            variance = lax.pmean(
-                jnp.mean(jnp.square(value - mean[:, None]), axis=1), "pmt_batch"
+            count = lax.psum(valid.astype(jnp.float32) * x.shape[1], self.axis_name)
+            count = eqx.error_if(
+                count, count == 1, "BatchNorm requires at least two valid observations."
+            )
+            denominator = jnp.maximum(count, 1)
+            mean = lax.psum(jnp.sum(value, axis=1), self.axis_name) / denominator
+            centered = jnp.where(valid, value - mean[:, None], 0)
+            variance = (
+                lax.psum(jnp.sum(jnp.square(centered), axis=1), self.axis_name)
+                / denominator
             )
-            count = lax.psum(jnp.asarray(x.shape[1], dtype=jnp.float32), "pmt_batch")
-            unbiased = variance * count / (count - 1)
+            unbiased = variance * count / jnp.maximum(count - 1, 1)
+            old_mean = state.get(self.mean_index)
+            old_variance = state.get(self.variance_index)
+            mean_update = (1 - self.momentum) * old_mean + self.momentum * mean
+            variance_update = (
+                1 - self.momentum
+            ) * old_variance + self.momentum * unbiased
             state = state.set(
-                self.mean_index,
-                (1 - self.momentum) * state.get(self.mean_index) + self.momentum * mean,
+                self.mean_index, jnp.where(count > 0, mean_update, old_mean)
             )
             state = state.set(
-                self.variance_index,
-                (1 - self.momentum) * state.get(self.variance_index)
-                + self.momentum * unbiased,
+                self.variance_index, jnp.where(count > 0, variance_update, old_variance)
+            )
+            state = state.set(
+                self.count_index,
+                state.get(self.count_index) + (count > 0).astype(jnp.int32),
             )
-            state = state.set(self.count_index, state.get(self.count_index) + 1)
         normalized = (value - mean[:, None]) * lax.rsqrt(variance[:, None] + self.eps)
-        return (normalized * self.weight[:, None] + self.bias[:, None]).astype(
-            x.dtype
-        ), state
+        return jnp.where(
+            valid, normalized * self.weight[:, None] + self.bias[:, None], 0
+        ).astype(x.dtype), state
 
 
 class PMDAttention(Attention):
@@ -295,10 +338,13 @@ class Lateral(eqx.Module):
         state: eqx.nn.State | None,
         *,
         inference: bool,
+        example_valid: bool | jax.Array = True,
     ) -> tuple[jax.Array, eqx.nn.State | None]:
         if self.norm_layer == "batchnorm":
             assert state is not None
-            normalized, state = self.norm(tokens.T, state, inference=inference)
+            normalized, state = self.norm(
+                tokens.T, state, inference=inference, example_valid=example_valid
+            )
             normalized = normalized.T
         elif self.norm_layer in ("groupnorm", "none"):
             normalized = self.norm(tokens.T).T
@@ -397,6 +443,7 @@ class PlainMaskDecoder(eqx.Module):
         inference: bool = True,
         key: jax.Array | None = None,
         mask_state: PMTMaskState | None = None,
+        example_valid: bool | jax.Array = True,
     ) -> PMTOutput | tuple[PMTOutput, eqx.nn.State | None]:
         """Run PMD; explicit state gives the same tuple interface for every norm."""
         supplied = state is not _MISSING
@@ -431,7 +478,9 @@ class PlainMaskDecoder(eqx.Module):
         )
         lateral = []
         for layer, tokens in zip(self.lateral, features.levels, strict=True):
-            result, running_state = layer(tokens, running_state, inference=inference)
+            result, running_state = layer(
+                tokens, running_state, inference=inference, example_valid=example_valid
+            )
             lateral.append(result)
         tokens = jnp.sum(jnp.stack(lateral), axis=0)
         tokens = jnp.concatenate((self.queries.astype(tokens.dtype), tokens), axis=0)

{equimo-2.3.0 → equimo-2.4.0}/src/equimo/vision/models/pmt.py

@@ -31,9 +31,10 @@ class PMT(eqx.Module):
 
     ``norm_layer="batchnorm"`` reproduces the author's lateral architecture.
     Construct that model with ``eqx.nn.make_with_state(PMT)`` and pass its state
-    on each call. Training must vmap across real images with
-    ``axis_name="pmt_batch"``. Normalization runs after ``encode`` so cached
-    frozen features can be reused while the decoder learns.
+    on each call. Training uses a named batch axis (default ``pmt_batch``).
+    Pass ``example_valid=False`` for padding; ``norm_kwargs['axis_name']`` may
+    include device axes for synchronized statistics. Normalization runs after
+    ``encode`` so cached frozen features can be reused while the decoder learns.
     """
 
     backbone: VisionTransformer
@@ -158,6 +159,7 @@ class PMT(eqx.Module):
         inference: bool = True,
         key: jax.Array | None = None,
         mask_state: PMTMaskState | None = None,
+        example_valid: bool | jax.Array = True,
     ) -> PMTOutput | tuple[PMTOutput, eqx.nn.State | None]:
         """Decode retained features with the selected normalization state."""
         if (
@@ -169,7 +171,12 @@ class PMT(eqx.Module):
         ):
             raise ValueError("PMT features disagree with encoder or input view.")
         return self.head(
-            features, state, inference=inference, key=key, mask_state=mask_state
+            features,
+            state,
+            inference=inference,
+            key=key,
+            mask_state=mask_state,
+            example_valid=example_valid,
         )
 
     def __call__(
@@ -180,6 +187,7 @@ class PMT(eqx.Module):
         inference: bool = True,
         key: jax.Array | None = None,
         mask_state: PMTMaskState | None = None,
+        example_valid: bool | jax.Array = True,
     ) -> PMTOutput | tuple[PMTOutput, eqx.nn.State | None]:
         """Return mask/class predictions; state is optional with GroupNorm."""
         return self.decode(
@@ -188,6 +196,7 @@ class PMT(eqx.Module):
             inference=inference,
             key=key,
             mask_state=mask_state,
+            example_valid=example_valid,
         )
 
 

equimo-2.4.0/src/equimo/vision/query_training/__init__.py

@@ -0,0 +1,40 @@
+"""Optimizer-independent, JAX-native query-segmentation training utilities."""
+
+from ._assignment import QueryAssignment, match_queries
+from ._sampling import PointSamples, sample_mask_points, sample_query_points
+from ._loss import (
+    LossTerm,
+    MatchingCosts,
+    PredictionLoss,
+    PredictionLossContext,
+    QueryLossConfig,
+    QueryLossContext,
+    QueryLossResult,
+    QueryTargets,
+    prepare_query_loss,
+    query_matching_costs,
+    query_prediction_loss,
+    query_segmentation_loss,
+    reduce_query_losses,
+)
+
+__all__ = [
+    "LossTerm",
+    "MatchingCosts",
+    "PointSamples",
+    "PredictionLoss",
+    "PredictionLossContext",
+    "QueryAssignment",
+    "QueryLossConfig",
+    "QueryLossContext",
+    "QueryLossResult",
+    "QueryTargets",
+    "match_queries",
+    "prepare_query_loss",
+    "query_matching_costs",
+    "query_prediction_loss",
+    "query_segmentation_loss",
+    "reduce_query_losses",
+    "sample_mask_points",
+    "sample_query_points",
+]

equimo-2.4.0/src/equimo/vision/query_training/_assignment.py

@@ -0,0 +1,130 @@
+# Copyright 2024 DeepMind Technologies Limited. All Rights Reserved.
+# Copyright 2026 Equimo contributors.
+# SPDX-License-Identifier: Apache-2.0
+"""Padded rectangular assignment, adapted from Optax's Hungarian algorithm.
+
+Reference: google-deepmind/optax, revision
+225a7079f4630bf75bee94ec78de9aa69c60fab3,
+optax/assignment/_hungarian_algorithm.py. See NOTICE for modifications.
+"""
+
+from __future__ import annotations
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from jax import lax
+
+
+class QueryAssignment(eqx.Module):
+    """Fixed-capacity matches in target order.
+
+    ``query_indices[t]`` is the query matched to target slot ``t``. Invalid
+    slots have query index ``-1``. ``is_valid`` is false for nonfinite
+    active costs or internal arithmetic overflow; no matches are returned in
+    that case. Indices and decisions
+    are nondifferentiable. Ties use the lowest available search index, keeping
+    an existing predecessor on equal costs; this is not a lexicographic solver.
+    """
+
+    query_indices: jax.Array
+    valid: jax.Array
+    is_valid: jax.Array
+
+
+def match_queries(costs: jax.Array, target_valid: jax.Array) -> QueryAssignment:
+    """Minimize ``(Q,T)`` costs with ``T <= Q``, entirely on the JAX device.
+
+    Every valid target is assigned a distinct query. Invalid target columns
+    are ignored, including their contents. Complexity is O(T**2 * Q), with
+    O(T * Q) cost storage. No host solver, callback, or differentiation through
+    assignment is used. Nonfinite active costs or internal arithmetic overflow
+    produce ``is_valid=False``. Every augmenting-path loop is explicitly bounded.
+    """
+    if costs.ndim != 2 or target_valid.shape != (costs.shape[1],):
+        raise ValueError("Expected (queries, targets) costs and (targets,) validity.")
+    queries, targets = costs.shape
+    if queries == 0 or targets > queries:
+        raise ValueError("Assignment requires 0 <= target capacity <= query count.")
+    if target_valid.dtype != jnp.bool_:
+        raise TypeError("Target validity must be boolean.")
+    matrix = lax.stop_gradient(costs.astype(jnp.float32).T)
+    finite = jnp.all(jnp.where(target_valid[:, None], jnp.isfinite(matrix), True))
+    active = target_valid & finite
+    matrix = jnp.where(active[:, None], matrix, 0)
+    if targets == 0:
+        return QueryAssignment(jnp.zeros((0,), dtype=jnp.int32), active, finite)
+
+    def insert(row, state):
+        def augment(state):
+            row_potential, column_potential, owner, feasible = state
+            owner = owner.at[0].set(row + 1)
+
+            def search(carry):
+                u, v, visited, distance, previous, column, steps = carry
+                visited = visited.at[column].set(True)
+                available = ~visited[1:]
+                reduced = matrix[owner[column] - 1] - u[owner[column]] - v[1:]
+                improves = available & (reduced < distance)
+                previous = jnp.where(improves, column, previous)
+                distance = jnp.where(improves, reduced, distance)
+                candidate = jnp.where(available, distance, jnp.inf)
+                next_column = jnp.argmin(candidate).astype(jnp.int32) + 1
+                delta = candidate[next_column - 1]
+                # Out-of-range scatter updates are dropped; unvisited columns
+                # must not repeatedly add to the dummy row's potential.
+                indices = jnp.where(visited, owner, targets + 1)
+                u = u.at[indices].add(delta, mode="drop")
+                v = jnp.where(visited, v - delta, v)
+                distance = jnp.where(available, distance - delta, distance)
+                return u, v, visited, distance, previous, next_column, steps + 1
+
+            def searching(carry):
+                *_, column, steps = carry
+                return (owner[column] != 0) & (steps <= row)
+
+            initial = (
+                row_potential,
+                column_potential,
+                jnp.zeros((queries + 1,), dtype=jnp.bool_),
+                jnp.full((queries,), jnp.inf, dtype=jnp.float32),
+                jnp.zeros((queries,), dtype=jnp.int32),
+                jnp.asarray(0, dtype=jnp.int32),
+                jnp.asarray(0, dtype=jnp.int32),
+            )
+            u, v, _, _, previous, column, _ = lax.while_loop(searching, search, initial)
+            feasible &= (
+                (owner[column] == 0)
+                & jnp.all(jnp.isfinite(u))
+                & jnp.all(jnp.isfinite(v))
+            )
+
+            def backtrack(carry):
+                owner, column, steps = carry
+                predecessor = previous[column - 1]
+                return owner.at[column].set(owner[predecessor]), predecessor, steps + 1
+
+            owner, column, _ = lax.while_loop(
+                lambda carry: feasible & (carry[1] != 0) & (carry[2] <= row),
+                backtrack,
+                (owner, column, jnp.asarray(0, dtype=jnp.int32)),
+            )
+            return u, v, owner, feasible & (column == 0)
+
+        return lax.cond(active[row] & state[-1], augment, lambda value: value, state)
+
+    initial = (
+        jnp.zeros((targets + 1,), dtype=jnp.float32),
+        jnp.zeros((queries + 1,), dtype=jnp.float32),
+        jnp.zeros((queries + 1,), dtype=jnp.int32),
+        finite,
+    )
+    _, _, owner, feasible = lax.fori_loop(0, targets, insert, initial)
+    indices = jnp.full((targets,), -1, dtype=jnp.int32)
+    destinations = jnp.where(owner[1:] > 0, owner[1:] - 1, targets)
+    indices = indices.at[destinations].set(
+        jnp.arange(queries, dtype=jnp.int32), mode="drop"
+    )
+    return QueryAssignment(
+        jnp.where(feasible, indices, -1), active & feasible, feasible
+    )