Equimo 2.5.0__tar.gz → 2.5.1__tar.gz

{equimo-2.5.0 → equimo-2.5.1}/CHANGELOG.md

@@ -2,6 +2,22 @@
 
 All notable changes to Equimo are documented here. Equimo follows Semantic Versioning from version 2.0.0 onward.
 
+## [2.5.1] - 2026-10-16
+
+### Changed
+
+- The EoMT and PMT mask-head `ScaleBlock` computes its 2 × 2, stride-2 transposed convolution as a matmul followed by
+  depth-to-space. This avoids an untuned cuDNN fallback for the bfloat16 weight gradient that made GPU training steps
+  several times slower at batch sizes of 12 and above. Parameters and checkpoints are unchanged; outputs match the
+  previous implementation within floating-point rounding.
+
+### Fixed
+
+- `match_queries` and the query-segmentation losses no longer rely on dropping out-of-range scatter updates. Under
+  `vmap` with `--xla_gpu_deterministic_ops` or `--xla_gpu_exclude_nondeterministic_ops`, jaxlib 0.11.2 and earlier write
+  such updates into the next batch row (openxla/xla#49380), which silently produced non-optimal query assignments and
+  wrong classification targets.
+
 ## [2.5.0] - 2026-10-15
 
 ### Changed

{equimo-2.5.0 → equimo-2.5.1}/PKG-INFO

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

{equimo-2.5.0 → equimo-2.5.1}/pyproject.toml

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

equimo-2.5.1/src/equimo/__init__.py

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

{equimo-2.5.0 → equimo-2.5.1}/src/equimo/vision/models/eomt.py

@@ -76,7 +76,19 @@ class ScaleBlock(eqx.Module):
         self.norm = LayerNorm2d(channels, eps=1e-6)
 
     def __call__(self, x: jax.Array) -> jax.Array:
-        return self.norm(self.conv2(jax.nn.gelu(self.conv1(x), approximate=False)))
+        # The 2x2 stride-2 transposed convolution has no overlapping windows, so
+        # it is a per-pixel matmul followed by depth-to-space. Lowering it as a
+        # convolution sends its bfloat16 weight gradient to an untuned cuDNN
+        # kernel that dominates training steps at larger batch sizes.
+        _, height, width = x.shape
+        # Equinox correlates the stored kernel with the dilated input.
+        weight = self.conv1.weight[:, :, ::-1, ::-1]
+        upscaled = jnp.einsum(
+            "oiab,ihw->ohawb", weight, x, preferred_element_type=jnp.float32
+        )
+        upscaled = upscaled.reshape(-1, 2 * height, 2 * width) + self.conv1.bias
+        upscaled = upscaled.astype(x.dtype)
+        return self.norm(self.conv2(jax.nn.gelu(upscaled, approximate=False)))
 
 
 class MaskEmbedding(eqx.Module):

{equimo-2.5.0 → equimo-2.5.1}/src/equimo/vision/query_training/_assignment.py

@@ -71,10 +71,11 @@ def match_queries(costs: jax.Array, target_valid: jax.Array) -> QueryAssignment:
                 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")
+                # Unvisited columns add zero instead of being dropped as
+                # out-of-range updates: deterministic GPU scatters in jaxlib
+                # <= 0.11.2 write those into the next batch row
+                # (openxla/xla#49380).
+                u = u.at[owner].add(jnp.where(visited, delta, 0))
                 v = jnp.where(visited, v - delta, v)
                 distance = jnp.where(available, distance - delta, distance)
                 return u, v, visited, distance, previous, next_column, steps + 1
@@ -120,11 +121,12 @@ def match_queries(costs: jax.Array, target_valid: jax.Array) -> QueryAssignment:
         finite,
     )
     _, _, owner, feasible = lax.fori_loop(0, targets, insert, initial)
-    indices = jnp.full((targets,), -1, dtype=jnp.int32)
+    # Unmatched queries write to a spare slot rather than out of range, which
+    # deterministic GPU scatters mishandle (openxla/xla#49380).
+    indices = jnp.full((targets + 1,), -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"
-    )
+    indices = indices.at[destinations].set(jnp.arange(queries, dtype=jnp.int32))
+    indices = indices[:targets]
     return QueryAssignment(
         jnp.where(feasible, indices, -1), active & feasible, feasible
     )

{equimo-2.5.0 → equimo-2.5.1}/src/equimo/vision/query_training/_loss.py

@@ -410,9 +410,11 @@ def query_prediction_loss(
     )
     is_valid &= jnp.all(counts <= 1)
     no_object = classes.shape[1] - 1
-    assigned_labels = jnp.full((classes.shape[0],), no_object, dtype=jnp.int32)
+    # Unmatched targets write to a spare slot rather than out of range, which
+    # deterministic GPU scatters mishandle (openxla/xla#49380).
+    assigned_labels = jnp.full((classes.shape[0] + 1,), no_object, dtype=jnp.int32)
     destinations = jnp.where(matched, indices, classes.shape[0])
-    assigned_labels = assigned_labels.at[destinations].set(labels, mode="drop")
+    assigned_labels = assigned_labels.at[destinations].set(labels)[:-1]
     weights = (
         jnp.where(assigned_labels == no_object, config.no_object_weight, 1.0)
         * targets.supervised

equimo-2.5.0/src/equimo/__init__.py

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