Equimo 2.5.1__tar.gz → 2.6.0__tar.gz

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

@@ -2,6 +2,26 @@
 
 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

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

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

@@ -1,6 +1,6 @@
 [project]
 name = "Equimo"
-version = "2.5.1"
+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.1 → equimo-2.6.0}/src/equimo/vision/models/eomt.py

@@ -83,8 +83,10 @@ class ScaleBlock(eqx.Module):
         _, 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=jnp.float32
+            "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)

{equimo-2.5.1 → 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.1/src/equimo/__init__.py

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