Equimo 2.4.0__tar.gz → 2.5.1__tar.gz

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

@@ -2,6 +2,38 @@
 
 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
+
+- Model forward and feature endpoints default `key` to `None` instead of a fixed `jax.random.PRNGKey(42)`. Deterministic
+  inference (`inference=True`) needs no key; a stochastic forward pass without a key raises a `ValueError` instead of
+  silently reusing the same randomness on every call.
+
+### Fixed
+
+- `extract_features`, `DenseProbe`, and the other fine-tuning wrappers no longer fail with their default `key=None` on
+  ViT-family models (including DINOv2, DINOv3, SigLIP2, EUPE, and LingBot-Vision), VisionParcae, MobileNetV3, AST, the
+  text encoder, TabPFN, and T0, whose feature endpoints required a PRNG key even in inference mode.
+- VisionParcae with `state_init="zero"` runs deterministic inference without a key. Key-dependent configurations,
+  including the default `state_init="like-init"`, raise an actionable error when the key is missing.
+
 ## [2.4.0] - 2026-09-24
 
 ### Added

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

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

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

@@ -190,7 +190,7 @@ class AudioSpectrogramTransformer(eqx.Module):
     def features(
         self,
         x: Float[Array, "time frequency"],
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         **kwargs,
     ) -> Float[Array, "seqlen dim"]:
@@ -207,7 +207,7 @@ class AudioSpectrogramTransformer(eqx.Module):
     def intermediate_features(
         self,
         x: Float[Array, "time frequency"],
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         indices: Sequence[int] | None = None,
         n_last_blocks: int | None = None,
@@ -251,7 +251,7 @@ class AudioSpectrogramTransformer(eqx.Module):
         self,
         x: Float[Array, "time frequency"],
         *,
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None,
         inference: Optional[bool],
     ) -> Float[Array, "seqlen dim"]:
         x = self.patch_embed(x)
@@ -266,7 +266,7 @@ class AudioSpectrogramTransformer(eqx.Module):
     def forward_features(
         self,
         x: Float[Array, "time frequency"],
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         **kwargs,
     ) -> dict:
@@ -283,7 +283,7 @@ class AudioSpectrogramTransformer(eqx.Module):
     def __call__(
         self,
         x: Float[Array, "time frequency"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         **kwargs,
     ) -> Float[Array, "num_classes"]:  # noqa: F821

{equimo-2.4.0 → equimo-2.5.1}/src/equimo/core/_prng.py

@@ -10,6 +10,11 @@ def split_for_mode(key: Any, num: int = 2, *, inference: bool | None) -> tuple:
 
     if inference is True:
         return (key,) * num
+    if key is None:
+        raise ValueError(
+            "A PRNG key is required for stochastic forward passes; pass key=... "
+            "or run with inference=True."
+        )
     return tuple(jr.split(key, num))
 
 

{equimo-2.4.0 → equimo-2.5.1}/src/equimo/language/models.py

@@ -68,7 +68,7 @@ class TransformerEncoderStack(eqx.Module):
     def __call__(
         self,
         x: Float[Array, "seqlen dim"],
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         mask: Optional[Float[Array, ""]] = None,
         *,
@@ -90,7 +90,7 @@ class TransformerEncoderStack(eqx.Module):
     def intermediate_features(
         self,
         x: Float[Array, "seqlen dim"],
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         mask: Optional[Float[Array, ""]] = None,
         indices: Sequence[int] | None = None,
@@ -209,7 +209,7 @@ class TextTransformerEncoder(eqx.Module):
         self,
         ids: Int[Array, "seqlen"],  # noqa: F821
         padding_mask: Float[Array, "seqlen"],  # noqa: F821
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
     ) -> Float[Array, "seqlen dim"]:
         x, valid_mask = self._prepare_tokens(ids, padding_mask)
@@ -226,7 +226,7 @@ class TextTransformerEncoder(eqx.Module):
         self,
         ids: Int[Array, "seqlen"],  # noqa: F821
         padding_mask: Float[Array, "seqlen"],  # noqa: F821
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         indices: Sequence[int] | None = None,
         n_last_blocks: int | None = None,
@@ -248,7 +248,7 @@ class TextTransformerEncoder(eqx.Module):
         self,
         ids: Int[Array, "seqlen"],  # noqa: F821
         padding_mask: Float[Array, "seqlen"],  # noqa: F821
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
     ) -> Float[Array, "dim"]:
         x = self.features(ids, padding_mask, key=key, inference=inference)

{equimo-2.4.0 → equimo-2.5.1}/src/equimo/tabular/models/tabpfn.py

@@ -217,7 +217,7 @@ class TabPFN(eqx.Module):
         x: Float[Array, "rows columns"],
         y: Array,
         n_train: int,
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         **kwargs,
     ) -> Float[Array, "rows dim"]:
@@ -256,7 +256,7 @@ class TabPFN(eqx.Module):
         x: Float[Array, "rows columns"],
         y: Array,
         n_train: int,
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         indices: Sequence[int] | None = None,
         n_last_blocks: int | None = None,
@@ -322,7 +322,7 @@ class TabPFN(eqx.Module):
         x: Float[Array, "rows columns"],
         y: Array,
         n_train: int,
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         **kwargs,
     ) -> dict:
@@ -339,7 +339,7 @@ class TabPFN(eqx.Module):
         x: Float[Array, "rows columns"],
         y: Array,
         n_train: int,
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         **kwargs,
     ) -> Float[Array, "test classes"]:

{equimo-2.4.0 → equimo-2.5.1}/src/equimo/timeseries/models/t0.py

@@ -474,7 +474,7 @@ class T0(eqx.Module):
         group_ids: Int[Array, "variates time"],
         variate_type: Int[Array, "variates time"],
         *,
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
     ) -> Float[Array, "tokens embed_dim"]:
         """Return post-normalization latent tokens before quantile decoding.
@@ -503,7 +503,7 @@ class T0(eqx.Module):
         group_ids: Int[Array, "variates time"],
         variate_type: Int[Array, "variates time"],
         *,
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
     ) -> Float[Array, "variates patches patch_size quantiles"]:
         key_encoder, key_blocks, key_decoder = split_for_mode(

{equimo-2.4.0 → equimo-2.5.1}/src/equimo/vision/models/_features.py

@@ -8,7 +8,6 @@ its saved-checkpoint signature, unchanged.
 from typing import TYPE_CHECKING, Any, Collection, Optional, Sequence
 
 import jax
-import jax.random as jr
 from einops import reduce
 from jaxtyping import Array, Float, PRNGKeyArray
 
@@ -50,7 +49,7 @@ class DenseStageFeatures:
     def features(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         **kwargs,
     ) -> Float[Array, "dim height width"]:
@@ -66,7 +65,7 @@ class DenseStageFeatures:
     def intermediate_features(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         indices: Sequence[int] | None = None,
         n_last_blocks: int | None = None,
@@ -88,7 +87,7 @@ class DenseStageFeatures:
     def __call__(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         **kwargs,
     ) -> Float[Array, "num_classes"]:  # noqa: F821
@@ -116,7 +115,7 @@ class TokenStemFeatures:
     def features(
         self,
         x: Float[Array, "..."],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
     ) -> Float[Array, "..."]:
         key_pd, *keys = split_for_mode(key, 1 + len(self.blocks), inference=inference)
@@ -130,7 +129,7 @@ class TokenStemFeatures:
     def intermediate_features(
         self,
         x: Float[Array, "..."],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         indices: Sequence[int] | None = None,
         n_last_blocks: int | None = None,
@@ -153,7 +152,7 @@ class TokenStemFeatures:
     def __call__(
         self,
         x: Float[Array, "..."],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
     ) -> Float[Array, "..."]:
         """Features, token-mean pooling, and classification head."""

{equimo-2.4.0 → equimo-2.5.1}/src/equimo/vision/models/attnet.py

@@ -108,7 +108,7 @@ class AttNet(eqx.Module):
     def features(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         **kwargs,
     ) -> Float[Array, "num_classes"]:  # noqa: F821
@@ -119,7 +119,7 @@ class AttNet(eqx.Module):
     def intermediate_features(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         indices: Sequence[int] | None = None,
         n_last_blocks: int | None = None,
@@ -141,7 +141,7 @@ class AttNet(eqx.Module):
     def __call__(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         **kwargs,
     ) -> Float[Array, "num_classes"]:  # noqa: F821

{equimo-2.4.0 → equimo-2.5.1}/src/equimo/vision/models/deq.py

@@ -227,7 +227,7 @@ class BlockChunk(eqx.Module):
         x: Float[Array, "..."],
         *,
         z0: Float[Array, "..."] | None = None,
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None,
         inference: Optional[bool] = False,
         **kwargs,
     ) -> Tuple[Float[Array, "..."], list]:
@@ -362,7 +362,7 @@ class DEQ(eqx.Module):
     def features(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         **kwargs,
     ) -> Tuple[Float[Array, "num_classes"], list]:  # noqa: F821
@@ -380,7 +380,7 @@ class DEQ(eqx.Module):
     def intermediate_features(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         indices: Sequence[int] | None = None,
         n_last_blocks: int | None = None,
@@ -404,7 +404,7 @@ class DEQ(eqx.Module):
     def readout(
         self,
         z_star: Float[Array, "channels height width"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
     ):
         if len(self.fpi_index) == 0:
@@ -434,7 +434,7 @@ class DEQ(eqx.Module):
     def __call__(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         **kwargs,
     ) -> Tuple[Float[Array, "num_classes"], list]:  # noqa: F821

{equimo-2.4.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.4.0 → equimo-2.5.1}/src/equimo/vision/models/fastervit.py

@@ -393,7 +393,7 @@ class FasterViT(eqx.Module):
     def features(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
     ) -> Float[Array, "seqlen dim"]:
         """Extract features from input image.
@@ -422,7 +422,7 @@ class FasterViT(eqx.Module):
     def intermediate_features(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         indices: Sequence[int] | None = None,
         n_last_blocks: int | None = None,
@@ -451,7 +451,7 @@ class FasterViT(eqx.Module):
     def __call__(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
     ) -> Float[Array, "num_classes"]:  # noqa: F821
         """Process input image through the full network.

{equimo-2.4.0 → equimo-2.5.1}/src/equimo/vision/models/lowformer.py

@@ -223,7 +223,7 @@ class LowFormer(eqx.Module):
     def features(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         **kwargs,
     ) -> Float[Array, "seqlen dim"]:
@@ -251,7 +251,7 @@ class LowFormer(eqx.Module):
     def intermediate_features(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         indices: Sequence[int] | None = None,
         n_last_blocks: int | None = None,
@@ -280,7 +280,7 @@ class LowFormer(eqx.Module):
     def __call__(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         **kwargs,
     ) -> Float[Array, "num_classes"]:  # noqa: F821

{equimo-2.4.0 → equimo-2.5.1}/src/equimo/vision/models/mobilenet.py

@@ -75,7 +75,7 @@ class MobileNetv3(eqx.Module):
     def features(
         self,
         x: Float[Array, "..."],
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
     ) -> Float[Array, "..."]:
         key_conv1, key_layers = split_for_mode(key, 2, inference=inference)
@@ -93,7 +93,7 @@ class MobileNetv3(eqx.Module):
     def intermediate_features(
         self,
         x: Float[Array, "..."],
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         indices: Sequence[int] | None = None,
         n_last_blocks: int | None = None,
@@ -124,7 +124,7 @@ class MobileNetv3(eqx.Module):
     def __call__(
         self,
         x: Float[Array, "..."],
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
     ) -> Float[Array, "..."]:
         key_f, key_d = split_for_mode(key, 2, inference=inference)

{equimo-2.4.0 → equimo-2.5.1}/src/equimo/vision/models/parcae.py

@@ -1067,7 +1067,7 @@ class VisionParcae(eqx.Module):
         H: int,
         W: int,
         inference: Optional[bool],
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None,
         rotary=None,
         checkpoint: bool = False,
         **kwargs,
@@ -1288,7 +1288,7 @@ class VisionParcae(eqx.Module):
         H: int,
         W: int,
         inference: Optional[bool],
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None,
         rotary=None,
         checkpoint: bool = False,
         **kwargs,
@@ -1315,7 +1315,7 @@ class VisionParcae(eqx.Module):
         H: int,
         W: int,
         inference: Optional[bool],
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None,
         num_steps: int | Array | None = None,
         num_steps_pair: tuple[int | Array, int | Array]
         | list[int | Array]
@@ -1344,11 +1344,17 @@ class VisionParcae(eqx.Module):
             key=key_rope,
         )
 
-        keys = jnp.stack(
-            split_for_mode(
-                key_loop,
-                self.max_recurrence,
-                inference=deterministic_inference,
+        # A missing key is only accepted for deterministic inference, where no
+        # recurrent step consumes randomness.
+        keys = (
+            None
+            if key_loop is None
+            else jnp.stack(
+                split_for_mode(
+                    key_loop,
+                    self.max_recurrence,
+                    inference=deterministic_inference,
+                )
             )
         )
 
@@ -1372,9 +1378,13 @@ class VisionParcae(eqx.Module):
             grad_scan_bound = min(self.max_recurrence, self.mean_backprop_depth)
             no_grad_scan_bound = self.max_recurrence - grad_scan_bound
 
-        def no_grad_body(carry, xs):
+        def step_key(index):
+            if keys is None:
+                return None
+            return jax.lax.dynamic_index_in_dim(keys, index, axis=0, keepdims=False)
+
+        def no_grad_body(carry, t):
             h_cur, last_cur = carry
-            t, key_step = xs
             active = t < no_grad_steps
 
             def inactive(c):
@@ -1388,7 +1398,7 @@ class VisionParcae(eqx.Module):
                     H=H,
                     W=W,
                     inference=inference,
-                    key=key_step,
+                    key=step_key(t),
                     rotary=rotary,
                     checkpoint=False,
                     **kwargs,
@@ -1401,11 +1411,7 @@ class VisionParcae(eqx.Module):
 
         if no_grad_scan_bound:
             no_grad_ts = jnp.arange(no_grad_scan_bound, dtype=jnp.int32)
-            (h, last_h), _ = jax.lax.scan(
-                no_grad_body,
-                (h, h),
-                (no_grad_ts, keys[:no_grad_scan_bound]),
-            )
+            (h, last_h), _ = jax.lax.scan(no_grad_body, (h, h), no_grad_ts)
             h = jax.lax.stop_gradient(h)
             last_h = jax.lax.stop_gradient(last_h)
         else:
@@ -1424,16 +1430,13 @@ class VisionParcae(eqx.Module):
                     no_grad_steps + t,
                     jnp.asarray(self.max_recurrence - 1, jnp.int32),
                 )
-                key_step = jax.lax.dynamic_index_in_dim(
-                    keys, key_index, axis=0, keepdims=False
-                )
                 h_next = self._recurrent_step(
                     h_active,
                     e,
                     H=H,
                     W=W,
                     inference=inference,
-                    key=key_step,
+                    key=step_key(key_index),
                     rotary=rotary,
                     checkpoint=self.recurrent_checkpoint,
                     **kwargs,
@@ -1463,7 +1466,7 @@ class VisionParcae(eqx.Module):
         self,
         x: Float[Array, "channels height width"],
         *,
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None,
         mask: Optional[Int[Array, "embed_h embed_w"]] = None,
         inference: Optional[bool] = None,
         num_steps: int | Array | None = None,
@@ -1474,6 +1477,12 @@ class VisionParcae(eqx.Module):
         **kwargs,
     ) -> tuple[Float[Array, "seq dim"], dict]:
         deterministic_inference = inference is True and self.state_init == "zero"
+        if key is None and not deterministic_inference:
+            raise ValueError(
+                "VisionParcae requires a PRNG key unless inference=True and "
+                f"state_init='zero' (got inference={inference!r}, "
+                f"state_init={self.state_init!r})."
+            )
         key_embed, key_prelude, key_loop, key_coda = split_for_mode(
             key, 4, inference=deterministic_inference
         )
@@ -1525,7 +1534,7 @@ class VisionParcae(eqx.Module):
     def features(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None = None,
         mask: Optional[Int[Array, "embed_h embed_w"]] = None,
         inference: Optional[bool] = None,
         num_steps: int | Array | None = None,
@@ -1556,7 +1565,7 @@ class VisionParcae(eqx.Module):
     def forward_features(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         **kwargs,
     ) -> dict:
@@ -1590,7 +1599,7 @@ class VisionParcae(eqx.Module):
     def intermediate_features(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray,
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         indices: Sequence[int] | None = None,
         n_last_blocks: int | None = None,
@@ -1686,7 +1695,7 @@ class VisionParcae(eqx.Module):
         *,
         H: int | None = None,
         W: int | None = None,
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = True,
         **kwargs,
     ) -> Float[Array, "num_classes"]:  # noqa: F821
@@ -1723,7 +1732,7 @@ class VisionParcae(eqx.Module):
     def __call__(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         num_steps: int | Array | None = None,
         num_steps_pair: tuple[int | Array, int | Array]

{equimo-2.4.0 → equimo-2.5.1}/src/equimo/vision/models/partialformer.py

@@ -305,7 +305,7 @@ class PartialFormer(eqx.Module):
     def features(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = ...,
+        key: PRNGKeyArray | None = ...,
         inference: Optional[bool] = ...,
         return_qa: Literal[False] = ...,
     ) -> Float[Array, "seqlen dim"]: ...
@@ -314,7 +314,7 @@ class PartialFormer(eqx.Module):
     def features(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = ...,
+        key: PRNGKeyArray | None = ...,
         inference: Optional[bool] = ...,
         *,
         return_qa: Literal[True],
@@ -323,7 +323,7 @@ class PartialFormer(eqx.Module):
     def features(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         return_qa: bool = False,
     ) -> (
@@ -364,7 +364,7 @@ class PartialFormer(eqx.Module):
     def intermediate_features(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
         indices: Sequence[int] | None = None,
         n_last_blocks: int | None = None,
@@ -401,7 +401,7 @@ class PartialFormer(eqx.Module):
     def __call__(
         self,
         x: Float[Array, "channels height width"],
-        key: PRNGKeyArray = jr.PRNGKey(42),
+        key: PRNGKeyArray | None = None,
         inference: Optional[bool] = None,
     ) -> Float[Array, "num_classes"]:  # noqa: F821
         """Process input image through the full network.