Equimo 2.6.0__tar.gz → 2.8.0__tar.gz

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

@@ -2,6 +2,140 @@
 
 All notable changes to Equimo are documented here. Equimo follows Semantic Versioning from version 2.0.0 onward.
 
+## [2.8.0] - 2026-11-10
+
+### Added
+
+- SAM 3.1 in `equimo.vision.models.sam3` (`Sam3`, `sam31_multiplex`, `decode_concepts`, the prompt and output types and
+  an experimental video tracker), `sam31_text` and `ClipBpeTokenizer` in `equimo.language` with a `clip` extra, and
+  converted weights under the SAM License. See `docs/sam3.md`.
+- SAM 3.1 prompting helpers: `encode_concept`, `source_size` for prompts given in original pixels, default positive
+  interactive labels with validation, and `decode_concepts` keeping every query by default.
+- `Sam3VideoTracker.remove_object`.
+- `equimo.core.parameters` describes how optimizers should treat each leaf, for every training workflow:
+  `trainable_mask` (an Equinox filter excluding buffers), `weight_decay_mask` (also usable as an Optax `mask`),
+  `DEFAULT_NO_DECAY_TAGS`, `buffer_mask`, `module_tags`, and the field markers `buffer()` and `tagged(...)`. The path
+  tagger moved to `equimo.core.tags`; `equimo.finetune` and `equimo.finetune.tags` re-export the same functions.
+
+### Changed
+
+- RoPE and Swin position embeddings compute the values that follow from their configuration (`VisionRoPE.freqs`,
+  `DinoRoPE.periods`, `RoPE.rotations`, and the `PosEmbMLPSwinv1D`/`PosEmbMLPSwinv2D` coordinate and index tables) when
+  used instead of storing them as PyTree leaves, so optimizers no longer see them and weight decay can no longer change
+  them. They remain readable as attributes but cannot be replaced with `eqx.tree_at`; set them through constructor
+  arguments. The structural signatures of models using them change. Checkpoints, fine-tuning bundles, and PMT artifacts
+  saved by earlier versions still load; `load_weights` raises `ValueError` when a stored value differs from the computed
+  one by more than rounding, for example after weight decay, and `load_delta` raises `FineTuneBundleError` for a bundle
+  whose trained delta changed such a value. Earlier versions cannot read checkpoints, bundles, or PMT artifacts that 2.8
+  saves for these models. The value hashes that identify a model (fine-tuning bundle base-checkpoint hashes, merge
+  checks, PMT encoder digests) cover the static fields of these modules instead, such as RoPE `base` and
+  `periods_dtype`, so a base model that differs only in those settings is rejected.
+- Frozen PEFT arrays are declared buffers and excluded from every trainability mode: VeRA `vera_A`/`vera_B`, LoRA-FA
+  `frozen_A`/`correction_matrix`, RandLoRA bases, PiSSA `base_weight_delta`, OFT `unmerged_weight`, and a fixed
+  AdaptFormer `scale`. `mode="full"` previously trained the VeRA bases, the LoRA-FA correction matrix, and the OFT
+  backup weight. `make_param_labels` without `trainable_paths` gives buffers the label `None` instead of a decay label.
+- The default no-weight-decay policy (`LLRDConfig.no_weight_decay_tags`, `DEFAULT_NO_DECAY_TAGS`) also excludes layer
+  scales (`block.layer_scale`, including PMD lateral scales) and state-space dynamics (`ssm`: Mamba `A_log`, `D`,
+  `dt_bias`; Parcae `A_log`, `dt_bias`, `B`, `C`). Leaves of registered normalization layers are tagged `norm` whatever
+  their field name, including norms inside `eqx.nn.Sequential`, and leaves named `*_bias` are tagged `bias`. The path
+  tagger also tags the Swin-v2 continuous position-bias MLP (`pos_emb_funct`, used by FasterViT HAT blocks)
+  `embedding.position`. These tags change tag-based selection (`TargetSpec` tags, `train_norm`, `train_bias`, and the
+  `norm`, `head_plus_norm`, `bias`, and surgical modes), parameter roles, plan labels, and plan fingerprints.
+  `CANONICAL_TAGS` gains `ssm`, `buffer`, and `vector`.
+- The path tagger tags arrays with at most one axis `vector`, and the default no-weight-decay policy excludes them, as
+  timm, MAE, and OpenCLIP do. With the tags above, built-in backbones need no `vector` tag; PEFT and probe-head vectors
+  no longer decay: DoRA magnitudes, AdaLoRA singular values, VeRA and RandLoRA scales, IA3 vectors, FourierFT
+  coefficients, scale/shift parameters, adapter residual scales, and the `AttentionPool` and
+  `AttentionPoolingClassifierHead` query vectors. Their plan labels change from `*_decay` to `*_no_decay`.
+- Position-embedding tables are evaluated while tracing, so `jax.jit` programs and ONNX exports hold them as constants.
+  MLLA attention rotates queries and keys with one RoPE call, so its programs hold one table per block. Exporting `RoPE`
+  (MLLA) or the Swin tables to ONNX needs a jax2onnx release that includes enpasos/jax2onnx#291; jax2onnx 0.17.1 and
+  0.18.0 fail on them.
+- `VisionParcae.weight_decay_mask()` follows the default policy, so it also leaves biases, norms, layer scales, vectors,
+  and position embeddings undecayed.
+- `update_ema_model` averages only trainable leaves; integer arrays and buffers take the current model's values.
+- `l2_sp_loss`, `adapter_norm_loss`, and `task_vector_norm_loss` skip buffers.
+- `VisionTransformer` raises `TypeError` for unknown constructor arguments instead of ignoring them.
+- `DropPath` keeps or drops the whole unbatched input with one draw, matching `DropPathAdd`; it previously dropped rows
+  of the leading axis (tokens or channels) independently.
+- `PromptedModel.__call__` raises `ValueError` when the base model has no `head`; use `features()` for headless models.
+- `apply_adalora` raises `ValueError` for a nonzero `dropout`, which AdaLoRA wrappers cannot apply.
+- `load_image` raises `ValueError` when only one of `mean` and `std` is given.
+- Tokenizer files are cached under a directory derived from the full URL, so pinned revisions no longer share a cache
+  entry; existing tokenizers are downloaded once more.
+
+### Fixed
+
+- `FasterViT` runs models with three or more stages. Their window stages previously raised `ValueError` on every forward
+  endpoint.
+- `SHViT.__call__` pools before its final norm. It previously raised `ValueError` with the default `norm_layer`.
+- `LowFormer` drops the stem entry of `att_strides`, so each stage uses its own stride and the presets run at 224
+  pixels. The last attention stage of every preset changes shape: its checkpoints saved by earlier versions no longer
+  load.
+- `FeatureSpec(pooling="attention")` returns the feature dtype instead of float32 for low-precision features.
+- `apply_eva_lora` skips targets it already wrapped instead of raising for missing activation artifacts.
+- `MllaBlock` (used by `Mlla` and `Vssd`) applies its output projection to the gated attention output and keeps the
+  first convolutional position encoding in the residual stream, as in MLLA and VSSD. Outputs of these models change.
+- EUPE ViTs use the DINOv3 patch RoPE configuration, including training-time coordinate rescaling, which their config
+  previously set under an ignored key. Inference outputs are unchanged.
+- EUPE ViTs compute their RoPE periods in bfloat16, as the EUPE checkpoints store them. Pretrained outputs are
+  unchanged; models built with `pretrained=False` previously used float32 periods.
+- LoRA-FA with `custom_vjp=False`, RandLoRA, and AdaptFormer with a fixed scale stop gradients to their frozen arrays.
+- `GroupSpec.tags` holds the tags shared by every leaf of the group instead of the first leaf's tags, so tag-based
+  optimizer rules, such as Rollfast's `GroupRule(tag=...)`, no longer depend on which leaf a group happened to start
+  with.
+- `apply_lora_rank_pattern` keeps the `metadata` of LoRA wrappers, such as PiSSA's initialization record.
+- Plans with layer-wise learning-rate decay, or other non-uniform multipliers, on PEFT models no longer raise
+  "inconsistent learning rate": PEFT labels then also name the block, head, or embedding (`lora_A_block_03_decay`
+  instead of `lora_A_decay`). Labels under uniform rates are unchanged.
+- `apply_adalora` honors `scaling="alpha_over_sqrt_r"`.
+- The `dynamictanh` norm (`DyT`) can be used as a block `norm_layer`.
+- Adapters, AdaptFormer, scale/shift, and `AttentionPool` return the activation dtype instead of promoting bfloat16 to
+  float32.
+- `save_delta` without `base_model` accepts a `TrainableSpec(mode="peft")` that trains no head, norm, or bias.
+- `SentencePieceTokenizer` checks for the `language` extra before downloading, writes downloads atomically, and times
+  out after 60 seconds.
+- PyTorch conversion renames only whole `scale` and `kernel` path components, keeps dictionary keys in parameter names,
+  and rejects an unknown `source` even when a `torch_model` is given.
+
+## [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.8.0}/LICENSES/pretrained/README.md

@@ -23,6 +23,7 @@ still current.
 | dinov3     | `dinov3_`         |     8 | DINOv3 License                                               | [snapshot](dinov3-License.md)                           | [DINOv3](https://github.com/facebookresearch/dinov3)                                     | [upstream license](https://github.com/facebookresearch/dinov3/blob/main/LICENSE.md)                              | 2026-08-21 |
 | eupe       | `eupe_`           |     6 | FAIR Noncommercial Research License v1                       | [snapshot](eupe-FAIR-Noncommercial-Research-License.md) | [EUPE](https://github.com/facebookresearch/EUPE)                                         | [upstream license](https://github.com/facebookresearch/EUPE/blob/main/LICENSE.md)                                | 2026-08-21 |
 | lingbot    | `lingbot_`        |     4 | Apache 2.0                                                   | [snapshot](lingbot-vision-Apache-2.0.txt)               | [LingBot-Vision checkpoints](https://huggingface.co/collections/robbyant/lingbot-vision) | [upstream license](https://github.com/robbyant/lingbot-vision/blob/main/LICENSE)                                 | 2026-09-23 |
+| sam31      | `sam31_`          |     2 | SAM License                                                  | [snapshot](sam-License.md)                              | [SAM 3.1](https://github.com/facebookresearch/sam3)                                      | [upstream license](https://github.com/facebookresearch/sam3/blob/main/LICENSE)                                   | 2026-11-09 |
 | siglip2    | `siglip2_`        |    15 | Apache 2.0                                                   | [snapshot](siglip2-Apache-2.0.txt)                      | [SigLIP 2 model](https://huggingface.co/google/siglip2-base-patch16-256)                 | [current license declaration](https://huggingface.co/google/siglip2-base-patch16-256)                            | 2026-08-21 |
 | t0         | `t0_`             |     1 | Apache 2.0                                                   | [snapshot](t0-alpha-Apache-2.0.txt)                     | [T0-alpha model](https://huggingface.co/theforecastingcompany/t0-alpha)                  | [upstream license](https://github.com/theforecastingcompany/tfc-t0/blob/main/LICENSE)                            | 2026-08-21 |
 | tabpfn     | `tabpfn_`         |     8 | TABPFN-3 License v1.0                                        | [snapshot](tabpfn-3-License-v1.0.txt)                   | [TabPFN-3 model](https://huggingface.co/Prior-Labs/tabpfn_3)                             | [upstream license](https://huggingface.co/Prior-Labs/tabpfn_3/blob/main/LICENSE)                                 | 2026-08-21 |
@@ -80,6 +81,15 @@ the upstream vision and text checkpoint parameters to its JAX/Equinox parameter
 source is the [`google-deepmind/tips` project](https://github.com/google-deepmind/tips), which states that it is not an
 official Google product. No endorsement by the upstream authors is implied.
 
+### SAM 3.1
+
+The `sam31_multiplex` and `sam31_text` archives are converted from Meta's SAM 3.1 checkpoint (`facebook/sam3.1`,
+`sam3.1_multiplex.pt`) and remain SAM Materials under the [SAM License](sam-License.md). Equimo converted the parameters
+to its JAX/Equinox layout and archive format; they are not an official Meta product and no endorsement by Meta is
+implied. The license requires that a copy of it accompany any distribution, that publications using the materials
+acknowledge them (section 1.b.ii), and that uses comply with trade controls (sections 1.b.iii and 1.b.v). Equimo's SAM
+3.1 code is not SAM Materials in the author's reading; see `docs/licensing/sam3-provenance.md`.
+
 ### T0-alpha
 
 Equimo converted the upstream T0-alpha checkpoint to its JAX/Equinox parameter layout and archive format. The Apache-2.0

equimo-2.8.0/LICENSES/pretrained/sam-License.md

@@ -0,0 +1,96 @@
+SAM License\
+Last Updated: November 19, 2025
+
+“Agreement” means the terms and conditions for use, reproduction, distribution and modification of the SAM Materials set
+forth herein.
+
+“SAM Materials” means, collectively, Documentation and the models, software and algorithms, including machine-learning
+model code, trained model weights, inference-enabling code, training-enabling code, fine-tuning enabling code, and other
+elements of the foregoing distributed by Meta and made available under this Agreement.
+
+“Documentation” means the specifications, manuals and documentation accompanying SAM Materials distributed by Meta.
+
+“Licensee” or “you” means you, or your employer or any other person or entity (if you are entering into this Agreement
+on such person or entity’s behalf), of the age required under applicable laws, rules or regulations to provide legal
+consent and that has legal authority to bind your employer or such other person or entity if you are entering in this
+Agreement on their behalf.
+
+“Meta” or “we” means Meta Platforms Ireland Limited (if you are located in or, if you are an entity, your principal
+place of business is in the EEA or Switzerland) or Meta Platforms, Inc. (if you are located outside of the EEA or
+Switzerland).
+
+“Sanctions” means any economic or trade sanctions or restrictions administered or enforced by the United States
+(including the Office of Foreign Assets Control of the U.S. Department of the Treasury (“OFAC”), the U.S. Department of
+State and the U.S. Department of Commerce), the United Nations, the European Union, or the United Kingdom.
+
+“Trade Controls” means any of the following: Sanctions and applicable export and import controls.
+
+By using or distributing any portion or element of the SAM Materials, you agree to be bound by this Agreement.
+
+1\. License Rights and Redistribution.
+
+a. Grant of Rights. You are granted a non-exclusive, worldwide, non-transferable and royalty-free limited license under
+Meta’s intellectual property or other rights owned by Meta embodied in the SAM Materials to use, reproduce, distribute,
+copy, create derivative works of, and make modifications to the SAM Materials.
+
+b. Redistribution and Use. i. Distribution of SAM Materials, and any derivative works thereof, are subject to the terms
+of this Agreement. If you distribute or make the SAM Materials, or any derivative works thereof, available to a third
+party, you may only do so under the terms of this Agreement and you shall provide a copy of this Agreement with any such
+SAM Materials.
+
+ii. If you submit for publication the results of research you perform on, using, or otherwise in connection with SAM
+Materials, you must acknowledge the use of SAM Materials in your publication.
+
+iii. Your use of the SAM Materials must comply with applicable laws and regulations, including Trade Control Laws and
+applicable privacy and data protection laws. iv. Your use of the SAM Materials will not involve or encourage others to
+reverse engineer, decompile or discover the underlying components of the SAM Materials. v. You are not the target of
+Trade Controls and your use of SAM Materials must comply with Trade Controls. You agree not to use, or permit others to
+use, SAM Materials for any activities subject to the International Traffic in Arms Regulations (ITAR) or end uses
+prohibited by Trade Controls, including those related to military or warfare purposes, nuclear industries or
+applications, espionage, or the development or use of guns or illegal weapons. 2. User Support. Your use of the SAM
+Materials is done at your own discretion; Meta does not process any information nor provide any service in relation to
+such use. Meta is under no obligation to provide any support services for the SAM Materials. Any support provided is “as
+is”, “with all faults”, and without warranty of any kind.
+
+3\. Disclaimer of Warranty. UNLESS REQUIRED BY APPLICABLE LAW, THE SAM MATERIALS AND ANY OUTPUT AND RESULTS THEREFROM ARE
+PROVIDED ON AN “AS IS” BASIS, WITHOUT WARRANTIES OF ANY KIND, AND META DISCLAIMS ALL WARRANTIES OF ANY KIND, BOTH
+EXPRESS AND IMPLIED, INCLUDING, WITHOUT LIMITATION, ANY WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR
+FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING OR
+REDISTRIBUTING THE SAM MATERIALS AND ASSUME ANY RISKS ASSOCIATED WITH YOUR USE OF THE SAM MATERIALS AND ANY OUTPUT AND
+RESULTS.
+
+4\. Limitation of Liability. IN NO EVENT WILL META OR ITS AFFILIATES BE LIABLE UNDER ANY THEORY OF LIABILITY, WHETHER IN
+CONTRACT, TORT, NEGLIGENCE, PRODUCTS LIABILITY, OR OTHERWISE, ARISING OUT OF THIS AGREEMENT, FOR ANY LOST PROFITS OR ANY
+DIRECT OR INDIRECT, SPECIAL, CONSEQUENTIAL, INCIDENTAL, EXEMPLARY OR PUNITIVE DAMAGES, EVEN IF META OR ITS AFFILIATES
+HAVE BEEN ADVISED OF THE POSSIBILITY OF ANY OF THE FOREGOING.
+
+5\. Intellectual Property.
+
+a. Subject to Meta’s ownership of SAM Materials and derivatives made by or for Meta, with respect to any derivative
+works and modifications of the SAM Materials that are made by you, as between you and Meta, you are and will be the
+owner of such derivative works and modifications.
+
+b. If you institute litigation or other proceedings against Meta or any entity (including a cross-claim or counterclaim
+in a lawsuit) alleging that the SAM Materials, outputs or results, or any portion of any of the foregoing, constitutes
+infringement of intellectual property or other rights owned or licensable by you, then any licenses granted to you under
+this Agreement shall terminate as of the date such litigation or claim is filed or instituted. You will indemnify and
+hold harmless Meta from and against any claim by any third party arising out of or related to your use or distribution
+of the SAM Materials.
+
+6\. Term and Termination. The term of this Agreement will commence upon your acceptance of this Agreement or access to
+the SAM Materials and will continue in full force and effect until terminated in accordance with the terms and
+conditions herein. Meta may terminate this Agreement if you are in breach of any term or condition of this Agreement.
+Upon termination of this Agreement, you shall delete and cease use of the SAM Materials. Sections 3, 4 and 7 shall
+survive the termination of this Agreement.
+
+7\. Governing Law and Jurisdiction. This Agreement will be governed and construed under the laws of the State of
+California without regard to choice of law principles, and the UN Convention on Contracts for the International Sale of
+Goods does not apply to this Agreement. The courts of California shall have exclusive jurisdiction of any dispute
+arising out of this Agreement.
+
+8\. Modifications and Amendments. Meta may modify this Agreement from time to time; provided that they are similar in
+spirit to the current version of the Agreement, but may differ in detail to address new problems or concerns. All such
+changes will be effective immediately. Your continued use of the SAM Materials after any modification to this Agreement
+constitutes your agreement to such modification. Except as provided in this Agreement, no modification or addition to
+any provision of this Agreement will be binding unless it is in writing and signed by an authorized representative of
+both you and Meta.

{equimo-2.6.0 → equimo-2.8.0}/NOTICE

@@ -66,3 +66,45 @@ This product includes software derived from Chronos-2 by Amazon.com, Inc. or
 its affiliates, Copyright Amazon.com, Inc. All Rights Reserved, licensed under
 the Apache License 2.0.
 Source: https://github.com/amazon-science/chronos-forecasting
+
+CLIP tokenizer and vocabulary
+-----------------------------
+
+ClipBpeTokenizer in src/equimo/language/tokenizers.py is adapted from
+clip/simple_tokenizer.py in OpenAI's CLIP release, and
+src/equimo/language/data/bpe_simple_vocab_16e6.txt.gz is the vocabulary from
+that release.
+
+Source: https://github.com/openai/CLIP
+
+MIT License
+
+Copyright (c) 2021 OpenAI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+SAM 3.1 weights
+---------------
+
+Equimo's sam31_multiplex and sam31_text archives are converted from Meta's
+SAM 3.1 checkpoint and remain under the SAM License; see
+LICENSES/pretrained/sam-License.md. No SAM Materials are included in this
+distribution: the archives are downloaded separately. Equimo's SAM 3.1 code is
+under Equimo's license; docs/licensing/sam3-provenance.md records how it was
+written.

equimo-2.6.0/README.md → equimo-2.8.0/PKG-INFO

@@ -1,3 +1,59 @@
+Metadata-Version: 2.4
+Name: Equimo
+Version: 2.8.0
+Summary: Implementation of modern multimodal models in JAX/Equinox
+License-Expression: MIT AND Apache-2.0
+License-File: LICENSE.md
+License-File: LICENSES/Apache-2.0.txt
+License-File: LICENSES/pretrained/README.md
+License-File: LICENSES/pretrained/ast-BSD-3-Clause.txt
+License-File: LICENSES/pretrained/convnext-Apache-2.0.txt
+License-File: LICENSES/pretrained/convnextv2-CC-BY-NC-4.0.txt
+License-File: LICENSES/pretrained/dinov2-Apache-2.0.txt
+License-File: LICENSES/pretrained/dinov3-License.md
+License-File: LICENSES/pretrained/eupe-FAIR-Noncommercial-Research-License.md
+License-File: LICENSES/pretrained/lingbot-vision-Apache-2.0.txt
+License-File: LICENSES/pretrained/sam-License.md
+License-File: LICENSES/pretrained/siglip2-Apache-2.0.txt
+License-File: LICENSES/pretrained/t0-alpha-Apache-2.0.txt
+License-File: LICENSES/pretrained/tabpfn-3-License-v1.0.txt
+License-File: LICENSES/pretrained/tips-CC-BY-4.0.txt
+License-File: LICENSES/tfc-t0-APACHE-2.0.txt
+License-File: NOTICE
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+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.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
+Requires-Dist: torch>=2.11,<2.12 ; extra == 'audio'
+Requires-Dist: torchaudio>=2.11,<2.12 ; extra == 'audio'
+Requires-Dist: ftfy>=6.1,<7 ; extra == 'clip'
+Requires-Dist: regex>=2024.11.6 ; extra == 'clip'
+Requires-Dist: matplotlib>=3.10.1 ; extra == 'extras'
+Requires-Dist: pillow>=11.1.0 ; extra == 'extras'
+Requires-Dist: scikit-learn>=1.6.1 ; extra == 'extras'
+Requires-Dist: tensorflow>=2.18.1 ; extra == 'language'
+Requires-Dist: tensorflow-text>=2.18.1 ; extra == 'language'
+Requires-Dist: torch>=2.6.0 ; extra == 'torch'
+Requires-Dist: timm>=1.0.15 ; extra == 'torch'
+Requires-Python: >=3.12, <3.15
+Provides-Extra: audio
+Provides-Extra: clip
+Provides-Extra: extras
+Provides-Extra: language
+Provides-Extra: torch
+Description-Content-Type: text/markdown
+
 # Equimo: Modern Multimodal Models in JAX/Equinox
 
 **WARNING**: This is a research library implementing recent model architectures. The implementations are based on paper
@@ -7,6 +63,12 @@ Equimo provides JAX/Equinox implementations of recent architectures across modal
 modality today; language has first-class text encoders/tokenizers, audio includes AST spectrogram models, tabular
 includes TabPFN-3 core models, and the experimental time-series namespace includes the raw T0 backbone.
 
+![SAM 3.1 tracking the concept "person" through a film scene](https://rhizome-labs.com/equimo/sam31-video-tracking.gif)
+
+*SAM 3.1 video tracking from the text prompt `"person"`, with
+[the code in the guide](docs/sam3.md#tracking-a-film-clip). Excerpt from* Pulp Fiction *(1994, Miramax), which belongs
+to its rights holders and is not covered by Equimo's licenses.*
+
 ## Features
 
 - Pure JAX/Equinox implementations
@@ -56,7 +118,7 @@ Equimo is organized by modality, with reusable building blocks separated from mo
 
 | Namespace              | Purpose                                                                                 |
 | ---------------------- | --------------------------------------------------------------------------------------- |
-| `equimo.core`          | Shared layers, scan ops, implicit/DEQ utilities, EMA helpers                            |
+| `equimo.core`          | Shared layers, scan ops, implicit/DEQ utilities, EMA helpers, parameter tags and masks  |
 | `equimo.vision`        | Vision models, vision layers, and image IO                                              |
 | `equimo.language`      | Text encoders and tokenizers                                                            |
 | `equimo.audio`         | Audio models, layers, and checkpoint-linked AST waveform preprocessing                  |
@@ -100,6 +162,7 @@ Beyond a standard ViT (e.g., DINOv2 or SigLIP), Equimo provides other SotA archi
 | FreeNet        | [FreeNet: Liberating Depth-Wise Separable Operations for Building Faster Mobile Vision Architectures](https://ojs.aaai.org/index.php/AAAI/article/view/33041)   | 2025 | ✅‡       |
 | EUPE           | [Efficient Universal Perception Encoder](https://arxiv.org/abs/2603.22387)                                                                                      | 2026 | ✅        |
 | LingBot-Vision | [Vision Pretraining for Dense Spatial Perception](https://arxiv.org/abs/2607.05247)                                                                             | 2026 | ✅§       |
+| SAM 3.1        | [SAM 3: Segment Anything with Concepts](https://arxiv.org/abs/2511.16719)                                                                                       | 2025 | ✅¶       |
 | ViT-5          | [ViT-5: Vision Transformers for The Mid-2020s](https://arxiv.org/abs/2602.08071)                                                                                | 2026 | ✅        |
 
 \*: Only contains the Linear Angular Attention module. It is straightforward to build a ViT around it, but may require
@@ -116,6 +179,9 @@ network from these blocks.
 tokens, and self-attention without a key bias; the Giant variant uses a SwiGLU FFN. Pretrained Small, Base, Large, and
 Giant backbones are available — see [pretrained models](#list-of-pretrained-models).
 
+¶: SAM 3.1 finds every instance of a concept in an image, segments objects from clicks, and tracks through video
+(experimental). See the [SAM 3.1 guide](docs/sam3.md). Its weights are under the SAM License.
+
 ## Implemented Audio Models
 
 | Model | Paper                                                                  | Year | Status |
@@ -216,6 +282,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 +303,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
@@ -746,6 +813,7 @@ The following models have pretrained weights available in Equimo:
 - [TabPFN-3](https://arxiv.org/abs/2605.13986)
 - [T0-alpha](https://huggingface.co/theforecastingcompany/t0-alpha)
 - [LingBot-Vision](https://arxiv.org/abs/2607.05247) (Small, Base, Large, and Giant)
+- [SAM 3.1](https://arxiv.org/abs/2511.16719) (the `sam31_multiplex` vision and `sam31_text` text models)
 
 Before using any pretrained weights, review the [pretrained-model license index](LICENSES/pretrained/README.md). It
 links the current upstream terms and records when each bundled license snapshot was last checked. Upstream licenses can
@@ -753,7 +821,7 @@ change, so the copies in this repository may be outdated; always check the lates
 terms are welcome as issues or pull requests.
 
 Model identifiers map to filenames in Equimo's
-[Hugging Face repository](https://huggingface.co/poiretclement/equimo/tree/2ce35f6512516a67ff99c473423e6ca48aa0ae01/models/default).
+[Hugging Face repository](https://huggingface.co/poiretclement/equimo/tree/f4b4cf5bc2d47a98fd0df11eddb08fd3bb5dded6/models/default).
 
 Published ConvNeXt size factories load the default classifier with `pretrained=True`:
 
@@ -945,14 +1013,69 @@ 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()
+```
+
+Give the optimizer only the trainable leaves and Equimo's weight-decay mask. `trainable_mask` leaves out buffers, such
+as frozen PEFT bases, which weight decay would otherwise change; `weight_decay_mask` leaves biases, norms, layer scales,
+state-space dynamics, every other array with at most one axis, position embeddings, and learned class, register,
+distillation, and mask tokens undecayed:
+
+```python
+import optax
+from equimo.core.parameters import trainable_mask, weight_decay_mask
+
+params, static = eqx.partition(model_fp32, trainable_mask(model_fp32))
+optimizer = optax.adamw(3e-4, weight_decay=0.05, mask=weight_decay_mask)
+opt_state = optimizer.init(params)
+```
+
+Differentiate with respect to `params` and rebuild the model with `eqx.combine(params, static)`. Fine-tuning plans apply
+the same policy.
+
+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