Equimo 2.5.0__tar.gz → 2.6.0__tar.gz

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

@@ -2,6 +2,42 @@
 
 All notable changes to Equimo are documented here. Equimo follows Semantic Versioning from version 2.0.0 onward.
 
+## [2.6.0] - 2026-10-17
+
+### Fixed
+
+- `anneal_mask_state` and `anneal_pmt_mask_state` return exactly zero mask probability from the declared end step. On
+  GPUs, the float32 progress quotient could round just below one (600 / 600 is 0.99999994 on an H100), which left a
+  probability near 3e-7 at that step and made `mask_free_eomt` and `mask_free_pmt` refuse a model whose annealing had
+  ended.
+- Masked attention in EoMT and PMT keeps each query mask with exactly the annealed probability. A query kept its mask
+  when its uniform draw was at most the probability, so at probability zero a draw of exactly 0.0 (about one in 8.4
+  million per query) still applied it. Only draws exactly equal to the probability change outcome.
+
+## [2.5.2] - 2026-10-16
+
+### Fixed
+
+- The EoMT and PMT mask-head `ScaleBlock` accumulates its upscaling matmul in the input precision when that exceeds
+  float32. In 2.5.1, float64 models computed the upscale in float32 on CPU and failed on GPU, where cuBLAS has no
+  float64-input, float32-output GEMM.
+
+## [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.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Equimo
-Version: 2.5.0
+Version: 2.6.0
 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.6.0}/pyproject.toml

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

equimo-2.6.0/src/equimo/__init__.py

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

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

@@ -76,7 +76,21 @@ 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]
+        # Accumulate in at least float32 without demoting float64 inputs.
+        accumulation_dtype = jnp.promote_types(jnp.result_type(weight, x), jnp.float32)
+        upscaled = jnp.einsum(
+            "oiab,ihw->ohawb", weight, x, preferred_element_type=accumulation_dtype
+        )
+        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.6.0}/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.6.0}/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 → equimo-2.6.0}/src/equimo/vision/segmentation.py

@@ -63,7 +63,8 @@ def _query_attention_mask(
     """
     num_queries = mask_logits.shape[0]
     allowed = resize_mask_logits(mask_logits, grid) > 0
-    keep = jr.uniform(key, (num_queries,)) <= probability
+    # Strict: uniform draws include exactly zero, which must not keep a mask at probability zero.
+    keep = jr.uniform(key, (num_queries,)) < probability
     allowed = jnp.where(keep[:, None, None], allowed, True)
     patch_start = num_queries + num_prefix_tokens
     length = patch_start + grid[0] * grid[1]
@@ -90,7 +91,10 @@ def _anneal_mask_probabilities(
     start = jnp.asarray(start_steps, dtype=jnp.float32)
     end = jnp.asarray(end_steps, dtype=jnp.float32)
     progress = jnp.clip((step.astype(jnp.float32) - start) / (end - start), 0, 1)
-    return (1 - progress) ** power
+    # GPU float32 division is not correctly rounded (600 / 600 is 0.99999994 on an H100),
+    # so the end step itself, not the quotient reaching one, removes the masks.
+    ended = step >= jnp.asarray(end_steps, dtype=jnp.int32)
+    return jnp.where(ended, 0.0, (1 - progress) ** power)
 
 
 def semantic_scores(

equimo-2.5.0/src/equimo/__init__.py

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