Equimo 2.6.0__tar.gz → 2.7.0__tar.gz

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

@@ -2,6 +2,44 @@
 
 All notable changes to Equimo are documented here. Equimo follows Semantic Versioning from version 2.0.0 onward.
 
+## [2.7.0] - 2026-10-26
+
+### Added
+
+- `equimo.core.ops` exports `softmax`, `log_softmax`, `softmax_cross_entropy_with_integer_labels`, and
+  `softmax_cross_entropy_with_logit_targets`, following "A Fix for Softmax and Cross-Entropy". They compute in at least
+  float32, return the input dtype, and keep the dominant logit's small gradient at high confidence. The logit-target
+  loss, for distillation, also preserves near-equality gradients, stops target gradients, and supports higher-order
+  derivatives.
+- Experimental fused attention: `equimo.core.layers.with_attention_implementation(model, "cudnn")` runs core,
+  prefix-tuned, and tabular attention through cuDNN's fused kernel for bfloat16 or float16 models on CUDA GPUs, and
+  attention modules accept the same choice as an `implementation` argument. On an RTX 4080, it shortened ViT and EoMT
+  training steps by 7-22% and reduced their temporary memory by 17-56%. It does not support attention dropout or
+  jax2onnx export; the default `"xla"` implementation is unchanged.
+- `semantic_crop_origins` places the fewest evenly spaced crops that cover a scaled image, for sliding-window evaluation
+  with `merge_semantic_crops`. With the shorter side scaled to the crop size, it matches the EoMT reference's windows,
+  except that its last crop always ends at the image edge: the reference's floating-point spacing stops a pixel short
+  for some sizes, such as a 1,347-pixel side with 448-pixel crops.
+
+### Changed
+
+- Equimo requires JAX and jaxlib 0.11.1 or newer. With it, a compiled step that captures both a frozen encoder and its
+  images as constants matches eager execution, which JAX 0.8.1 did not.
+- Attention pooling, attention-pooling heads, adapter fusion, the tabular decoder, and the `"softmax"` activation
+  compute softmax in at least float32 with the corrected derivative of `equimo.core.ops.softmax`, returning the model
+  dtype. Tabular attention computes its softmax in at least float32. Attention logits are cast before scaling, avoiding
+  a second bfloat16 rounding when the square root of the head dimension is not a power of two. Float32 forward results
+  and checkpoint structure are unchanged.
+
+### Fixed
+
+- Query classification uses stable cross-entropy, so confident correct predictions retain nonzero losses and target
+  gradients.
+- `selective_scan` returns the input dtype. Its float32 scan state leaked into bfloat16 outputs, unlike `ssd`.
+- `merge_semantic_crops` raises a `ValueError` when the crops leave part of the scaled image uncovered. It returned NaN
+  scores there instead, and any final resize spread them to every output pixel, so the argmax labelled the whole image
+  class 0.
+
 ## [2.6.0] - 2026-10-17
 
 ### Fixed

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Equimo
-Version: 2.6.0
+Version: 2.7.0
 Summary: Implementation of modern multimodal models in JAX/Equinox
 License-Expression: MIT AND Apache-2.0
 License-File: LICENSE.md
@@ -29,8 +29,8 @@ Classifier: Programming Language :: Python :: 3.14
 Requires-Dist: banax>=0.1.2
 Requires-Dist: einops>=0.8.0
 Requires-Dist: equinox>=0.13.3
-Requires-Dist: jax>=0.8.1
-Requires-Dist: jaxlib>=0.8.1
+Requires-Dist: jax>=0.11.1
+Requires-Dist: jaxlib>=0.11.1
 Requires-Dist: loguru>=0.7.3
 Requires-Dist: lz4>=4.4.3
 Requires-Dist: requests>=2.32.3
@@ -268,6 +268,7 @@ Minimal example:
 import jax
 import equimo.finetune as eqft
 import equimo.vision.models as em
+from equimo.core.ops import softmax_cross_entropy_with_integer_labels
 
 key = jax.random.PRNGKey(0)
 model = em.vit_tiny_patch16_224(num_classes=10, key=key)
@@ -288,7 +289,7 @@ frozen = plan.frozen
 def loss_fn(trainable, batch):
     model = plan.combine(trainable)
     logits = jax.vmap(lambda x: model(x, key=key, inference=False))(batch["x"])
-    return cross_entropy(logits, batch["y"])  # supplied by user code
+    return softmax_cross_entropy_with_integer_labels(logits, batch["y"]).mean()
 ```
 
 Frozen leaves are absent from `plan.trainable`; they are not assigned a zero learning rate. `plan.labels` and
@@ -997,14 +998,52 @@ import jax
 import jax.numpy as jnp
 import equinox as eqx
 
-model_bf16 = jax.tree_util.tree_map(
-    lambda leaf: leaf.astype(jnp.bfloat16) if eqx.is_inexact_array(leaf) else leaf,
-    model_fp32,
-)
+
+def to_bfloat16(tree):
+    return jax.tree_util.tree_map(
+        lambda leaf: leaf.astype(jnp.bfloat16) if eqx.is_inexact_array(leaf) else leaf,
+        tree,
+    )
+
+
+model_bf16 = to_bfloat16(model_fp32)
+```
+
+For training, keep float32 parameters and cast inside the loss. Gradients then return in float32 and the optimizer
+updates float32 values; bfloat16 parameters drop updates smaller than about 0.4% of their magnitude.
+
+```python
+from equimo.core.ops import softmax_cross_entropy_with_integer_labels
+
+
+@eqx.filter_value_and_grad
+def loss_and_grads(model_fp32, images, labels, key):
+    model = to_bfloat16(model_fp32)
+    keys = jax.random.split(key, images.shape[0])
+    logits = jax.vmap(lambda x, k: model(x, key=k, inference=False))(images, keys)
+    return softmax_cross_entropy_with_integer_labels(logits, labels).mean()
+```
+
+Numerically sensitive operations cast locally to float32 and return the model dtype: normalization statistics, rotary
+factors, and every softmax. Attention between tokens applies `jax.nn.softmax` to float32 logits. Other softmaxes
+(attention pooling, attention-pooling heads, adapter fusion, the tabular decoder, and the `"softmax"` activation) use
+`equimo.core.ops.softmax`, whose derivative also stays accurate for highly confident distributions. For the same reason,
+use the [`equimo.core.ops` classification losses](docs/finetuning/classification_losses.md). The
+[performance notes](docs/performance.md#softmax-precision) give the measurements behind these choices.
+
+On NVIDIA GPUs with TF32 tensor cores (Ampere and newer), JAX's default precision may run float32 matrix multiplications
+as TF32. For float32 reference accuracy, run the computation under `jax.default_matmul_precision("highest")`.
+
+On CUDA GPUs, bfloat16 and float16 models can run attention through cuDNN's fused kernel, which saves memory and time on
+long sequences. This is experimental, requires inactive attention dropout, and cannot be exported with jax2onnx:
+
+```python
+from equimo.core.layers import with_attention_implementation
+
+model = with_attention_implementation(model_bf16, "cudnn")  # "xla" switches back
 ```
 
-Isolated `float32` upcasts are mandatory for numerically sensitive operations (softmax, layer norm variance). These are
-applied internally where needed.
+See the [fused-attention measurements](docs/performance.md#fused-attention).
 
 ## Contributing
 

{equimo-2.6.0 → equimo-2.7.0}/README.md

@@ -216,6 +216,7 @@ Minimal example:
 import jax
 import equimo.finetune as eqft
 import equimo.vision.models as em
+from equimo.core.ops import softmax_cross_entropy_with_integer_labels
 
 key = jax.random.PRNGKey(0)
 model = em.vit_tiny_patch16_224(num_classes=10, key=key)
@@ -236,7 +237,7 @@ frozen = plan.frozen
 def loss_fn(trainable, batch):
     model = plan.combine(trainable)
     logits = jax.vmap(lambda x: model(x, key=key, inference=False))(batch["x"])
-    return cross_entropy(logits, batch["y"])  # supplied by user code
+    return softmax_cross_entropy_with_integer_labels(logits, batch["y"]).mean()
 ```
 
 Frozen leaves are absent from `plan.trainable`; they are not assigned a zero learning rate. `plan.labels` and
@@ -945,14 +946,52 @@ import jax
 import jax.numpy as jnp
 import equinox as eqx
 
-model_bf16 = jax.tree_util.tree_map(
-    lambda leaf: leaf.astype(jnp.bfloat16) if eqx.is_inexact_array(leaf) else leaf,
-    model_fp32,
-)
+
+def to_bfloat16(tree):
+    return jax.tree_util.tree_map(
+        lambda leaf: leaf.astype(jnp.bfloat16) if eqx.is_inexact_array(leaf) else leaf,
+        tree,
+    )
+
+
+model_bf16 = to_bfloat16(model_fp32)
+```
+
+For training, keep float32 parameters and cast inside the loss. Gradients then return in float32 and the optimizer
+updates float32 values; bfloat16 parameters drop updates smaller than about 0.4% of their magnitude.
+
+```python
+from equimo.core.ops import softmax_cross_entropy_with_integer_labels
+
+
+@eqx.filter_value_and_grad
+def loss_and_grads(model_fp32, images, labels, key):
+    model = to_bfloat16(model_fp32)
+    keys = jax.random.split(key, images.shape[0])
+    logits = jax.vmap(lambda x, k: model(x, key=k, inference=False))(images, keys)
+    return softmax_cross_entropy_with_integer_labels(logits, labels).mean()
+```
+
+Numerically sensitive operations cast locally to float32 and return the model dtype: normalization statistics, rotary
+factors, and every softmax. Attention between tokens applies `jax.nn.softmax` to float32 logits. Other softmaxes
+(attention pooling, attention-pooling heads, adapter fusion, the tabular decoder, and the `"softmax"` activation) use
+`equimo.core.ops.softmax`, whose derivative also stays accurate for highly confident distributions. For the same reason,
+use the [`equimo.core.ops` classification losses](docs/finetuning/classification_losses.md). The
+[performance notes](docs/performance.md#softmax-precision) give the measurements behind these choices.
+
+On NVIDIA GPUs with TF32 tensor cores (Ampere and newer), JAX's default precision may run float32 matrix multiplications
+as TF32. For float32 reference accuracy, run the computation under `jax.default_matmul_precision("highest")`.
+
+On CUDA GPUs, bfloat16 and float16 models can run attention through cuDNN's fused kernel, which saves memory and time on
+long sequences. This is experimental, requires inactive attention dropout, and cannot be exported with jax2onnx:
+
+```python
+from equimo.core.layers import with_attention_implementation
+
+model = with_attention_implementation(model_bf16, "cudnn")  # "xla" switches back
 ```
 
-Isolated `float32` upcasts are mandatory for numerically sensitive operations (softmax, layer norm variance). These are
-applied internally where needed.
+See the [fused-attention measurements](docs/performance.md#fused-attention).
 
 ## Contributing
 

{equimo-2.6.0 → equimo-2.7.0}/docs/stability.md

@@ -30,6 +30,10 @@ tensor contract, and Equinox PyTree structure may change in a minor release whil
 adapter is developed. Such changes will be recorded in the changelog. Existing T0 archives remain covered by the
 checkpoint integrity and loading commitments below; this exception applies to the Python and raw-array API around them.
 
+Fused cuDNN attention, selected with `equimo.core.layers.with_attention_implementation` or `implementation="cudnn"`, is
+experimental. Its supported dtypes, masks, and module coverage follow cuDNN and `jax.nn.dot_product_attention` and may
+change in a minor release. The default `implementation="xla"` is stable.
+
 Optional integrations remain optional. Their public Equimo wrappers are stable, but the supported Python range may be
 narrower when an upstream dependency does not publish wheels. Equimo documents and tests those exceptions explicitly.
 

{equimo-2.6.0 → equimo-2.7.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "Equimo"
-version = "2.6.0"
+version = "2.7.0"
 description = "Implementation of modern multimodal models in JAX/Equinox"
 readme = "README.md"
 requires-python = ">=3.12,<3.15"
@@ -27,8 +27,8 @@ dependencies = [
   "banax>=0.1.2",
   "einops>=0.8.0",
   "equinox>=0.13.3",
-  "jax>=0.8.1",
-  "jaxlib>=0.8.1",
+  "jax>=0.11.1",
+  "jaxlib>=0.11.1",
   "loguru>=0.7.3",
   "lz4>=4.4.3",
   "requests>=2.32.3",
@@ -87,7 +87,7 @@ source-include = [
 
 [dependency-groups]
 dev = [
-    "jax2onnx>=0.15.1",
+    "jax2onnx>=0.17.0",
     "mdformat==1.0.0",
     "mdformat-footnote==0.1.3",
     "mdformat-gfm==1.0.0",

equimo-2.7.0/src/equimo/__init__.py

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

{equimo-2.6.0 → equimo-2.7.0}/src/equimo/core/layers/__init__.py

@@ -38,6 +38,7 @@ __all__ = [
     "register_ffn",
     "register_mixer",
     "register_norm",
+    "with_attention_implementation",
 ]
 
 from .activation import get_act, register_act
@@ -48,6 +49,7 @@ from .attention import (
     get_attn_block,
     register_attn,
     register_attn_block,
+    with_attention_implementation,
 )
 from .dropout import DropPath, DropPathAdd, get_dropout, register_dropout
 from .ffn import (

{equimo-2.6.0 → equimo-2.7.0}/src/equimo/core/layers/activation.py

@@ -3,6 +3,8 @@ from typing import Callable, Optional
 
 import jax
 
+from equimo.core.ops import softmax
+
 _ACT_REGISTRY: dict[str, Callable] = {}
 
 
@@ -65,5 +67,5 @@ _ACT_REGISTRY["elu"] = jax.nn.elu
 _ACT_REGISTRY["sigmoid"] = jax.nn.sigmoid
 _ACT_REGISTRY["hard_sigmoid"] = jax.nn.hard_sigmoid
 _ACT_REGISTRY["hard_swish"] = jax.nn.hard_swish
-_ACT_REGISTRY["softmax"] = jax.nn.softmax
+_ACT_REGISTRY["softmax"] = softmax
 _ACT_REGISTRY["relu6"] = jax.nn.relu6

{equimo-2.6.0 → equimo-2.7.0}/src/equimo/core/layers/attention.py

@@ -1,4 +1,5 @@
-from typing import Callable, List, Optional
+import dataclasses
+from typing import Callable, List, Literal, Optional, TypeVar
 
 import equinox as eqx
 import jax
@@ -59,6 +60,131 @@ def _broadcast_attention_mask(
         ) from error
 
 
+AttentionImplementation = Literal["xla", "cudnn"]
+_IMPLEMENTATIONS = ("xla", "cudnn")
+# cuDNN's fused attention backward requires sequence lengths divisible by 64.
+_CUDNN_SEQUENCE_MULTIPLE = 64
+
+_Model = TypeVar("_Model")
+
+
+def _check_implementation(implementation: str) -> None:
+    if implementation not in _IMPLEMENTATIONS:
+        raise ValueError(
+            f"implementation must be one of {_IMPLEMENTATIONS}, got {implementation!r}."
+        )
+
+
+def _check_cudnn_dropout(
+    dropout: eqx.nn.Dropout | None, inference: Optional[bool]
+) -> None:
+    """Reject attention-probability dropout, which the fused kernel cannot apply."""
+    if dropout is None:
+        return
+    if inference is None:
+        inference = dropout.inference
+    if not inference and dropout.p != 0:
+        raise ValueError(
+            "implementation='cudnn' does not support attention dropout. Set the "
+            "attention dropout rate to 0, call with inference=True, or use "
+            "implementation='xla'."
+        )
+
+
+def _fused_attention(
+    q: jax.Array,
+    k: jax.Array,
+    v: jax.Array,
+    allowed: jax.Array | None = None,
+    *,
+    implementation: AttentionImplementation = "cudnn",
+) -> jax.Array:
+    """Fused dot-product attention over ``(..., heads, seq, head_dim)`` inputs.
+
+    ``allowed`` broadcasts to ``(..., heads, query, key)``. Fully masked query rows
+    return zeros, as in the unfused path. Keys and values may have fewer heads
+    than queries if that number divides the query heads. ``implementation`` is
+    passed to :func:`jax.nn.dot_product_attention`; ``"xla"`` lets CPU tests
+    check the padding and masking.
+    """
+    if implementation == "cudnn" and q.dtype not in (jnp.bfloat16, jnp.float16):
+        raise ValueError(
+            "implementation='cudnn' requires bfloat16 or float16 attention inputs, "
+            f"got {q.dtype}. Cast the model or use implementation='xla'."
+        )
+    *leading, heads, query_len, head_dim = q.shape
+    key_len = k.shape[-2]
+    padded_query = query_len + -query_len % _CUDNN_SEQUENCE_MULTIPLE
+    padded_key = key_len + -key_len % _CUDNN_SEQUENCE_MULTIPLE
+
+    def batch_first(x: jax.Array, length: int) -> jax.Array:
+        x = jnp.swapaxes(x, -3, -2)
+        x = x.reshape(-1, *x.shape[-3:])
+        return jnp.pad(x, ((0, 0), (0, length - x.shape[1]), (0, 0), (0, 0)))
+
+    q = batch_first(q, padded_query)
+    k, v = batch_first(k, padded_key), batch_first(v, padded_key)
+    batch = q.shape[0]
+    if allowed is None:
+        out = jax.nn.dot_product_attention(
+            q,
+            k,
+            v,
+            query_seq_lengths=jnp.full((batch,), query_len, jnp.int32),
+            key_value_seq_lengths=jnp.full((batch,), key_len, jnp.int32),
+            implementation=implementation,
+        )
+    else:
+        # Materialize the mask for every head: a broadcast head axis fails under
+        # vmap (jax-ml/jax#34920). Padding adds masked queries and keys.
+        mask = jnp.broadcast_to(allowed, (*leading, heads, query_len, key_len))
+        mask = jnp.pad(
+            mask.reshape(batch, heads, query_len, key_len),
+            ((0, 0), (0, 0), (0, padded_query - query_len), (0, padded_key - key_len)),
+        )
+        out = jax.nn.dot_product_attention(
+            q, k, v, mask=mask, implementation=implementation
+        )
+        # cuDNN returns the mean value for fully masked rows instead of zeros.
+        has_key = jnp.swapaxes(jnp.any(mask, axis=-1), 1, 2)
+        out = jnp.where(has_key[..., None], out, 0)
+    out = out[:, :query_len].reshape(*leading, query_len, heads, head_dim)
+    return jnp.swapaxes(out, -3, -2)
+
+
+def with_attention_implementation(
+    model: _Model, implementation: AttentionImplementation
+) -> _Model:
+    """Return ``model`` with every attention module using ``implementation``.
+
+    ``"cudnn"`` runs core and tabular attention through NVIDIA cuDNN's fused
+    kernel (:func:`jax.nn.dot_product_attention`), which computes logits in
+    float32 without materializing them. It requires a CUDA GPU, bfloat16 or
+    float16 inputs, and inactive attention dropout, and it cannot be exported
+    with jax2onnx; switch back to ``"xla"`` before exporting. Parameters and
+    checkpoint layouts are unchanged. The setting is static, so compiled
+    functions retrace after switching.
+    """
+    _check_implementation(implementation)
+
+    def is_attention(node) -> bool:
+        return isinstance(node, eqx.Module) and any(
+            field.name == "implementation" for field in dataclasses.fields(node)
+        )
+
+    def switch(node):
+        if not is_attention(node):
+            return node
+        # tree_at cannot replace static fields, and copy.copy returns the same
+        # module, so rebuild the instance as Equinox's unflattening does.
+        switched = object.__new__(type(node))
+        switched.__dict__.update(node.__dict__)
+        object.__setattr__(switched, "implementation", implementation)
+        return switched
+
+    return jax.tree.map(switch, model, is_leaf=is_attention)
+
+
 @register_attn()
 class Attention(eqx.Module):
     """Multi-head self attention over tensors shaped ``(..., seq, dim)``.
@@ -67,11 +193,15 @@ class Attention(eqx.Module):
     batch-aligned with shape ``(..., query, key)`` or directly broadcastable to
     ``(..., heads, query, key)``. Fully masked query rows receive zero attention
     probability mass.
+
+    ``implementation="cudnn"`` selects the fused kernel described in
+    :func:`with_attention_implementation`.
     """
 
     dim: int = eqx.field(static=True)
     num_heads: int = eqx.field(static=True)
     head_dim: int = eqx.field(static=True)
+    implementation: AttentionImplementation = eqx.field(static=True)
 
     qkv: eqx.nn.Linear
     proj: eqx.nn.Linear
@@ -93,12 +223,15 @@ class Attention(eqx.Module):
         proj_drop: float = 0.0,
         norm_layer: str | type[eqx.Module] = "layernorm",
         eps: float = 1e-5,
+        implementation: AttentionImplementation = "xla",
         **kwargs,
     ):
         self.num_heads = num_heads
         self.head_dim = dim // num_heads
         self.dim = dim
         assert self.head_dim * num_heads == dim, "dim must be divisible by num_heads"
+        _check_implementation(implementation)
+        self.implementation = implementation
 
         key_qkv, key_proj = jr.split(key, 2)
         norm_layer = get_norm(norm_layer)
@@ -139,23 +272,24 @@ class Attention(eqx.Module):
                 )
             q, k = apply_rotary_qk(q, k, rotary)
 
-        attn = (
-            jnp.einsum("...hqd,...hkd->...hqk", q, k) / jnp.sqrt(self.head_dim)
-        ).astype(jnp.float32)
+        allowed = None
         if mask is not None:
             allowed = _broadcast_attention_mask(
                 mask,
                 x_ndim=x.ndim,
-                shape=attn.shape,
+                shape=(*q.shape[:-1], k.shape[-2]),
             )
-            attn = jnp.where(allowed, attn, -jnp.inf)
-        attn = jax.nn.softmax(attn, axis=-1)
-        if mask is not None:
-            attn = jnp.where(jnp.any(allowed, axis=-1, keepdims=True), attn, 0)
-        attn = attn.astype(x.dtype)
-        attn = self.attn_drop(attn, inference=inference, key=key1)
+        if self.implementation == "cudnn":
+            _check_cudnn_dropout(self.attn_drop, inference)
+            x = _fused_attention(q, k, v, allowed)
+        else:
+            attn = jnp.einsum("...hqd,...hkd->...hqk", q, k).astype(
+                jnp.float32
+            ) / jnp.sqrt(self.head_dim)
+            attn = jax.nn.softmax(attn, axis=-1, where=allowed).astype(x.dtype)
+            attn = self.attn_drop(attn, inference=inference, key=key1)
+            x = jnp.einsum("...hqk,...hkd->...hqd", attn, v)
 
-        x = jnp.einsum("...hqk,...hkd->...hqd", attn, v)
         x = jnp.swapaxes(x, -3, -2).reshape(*x.shape[:-3], x.shape[-2], self.dim)
         x = _apply_module_last_dim(self.proj, x)
         return self.proj_drop(x, inference=inference, key=key2)

equimo-2.7.0/src/equimo/core/ops/__init__.py

@@ -0,0 +1,18 @@
+from .scan import non_causal_linear_attn, segsum, selective_scan, ssd
+from .probability import (
+    log_softmax,
+    softmax,
+    softmax_cross_entropy_with_integer_labels,
+    softmax_cross_entropy_with_logit_targets,
+)
+
+__all__ = [
+    "non_causal_linear_attn",
+    "segsum",
+    "selective_scan",
+    "ssd",
+    "log_softmax",
+    "softmax",
+    "softmax_cross_entropy_with_integer_labels",
+    "softmax_cross_entropy_with_logit_targets",
+]