Equimo 0.1.0__py3-none-any.whl

Equimo-0.1.0.dist-info/LICENSE.md

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Clément POIRET
+
+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.

Equimo-0.1.0.dist-info/METADATA

@@ -0,0 +1,104 @@
+Metadata-Version: 2.1
+Name: Equimo
+Version: 0.1.0
+Summary: Implementation of popular vision models in Jax
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE.md
+Requires-Dist: einops>=0.8.0
+Requires-Dist: equinox>=0.11.5
+Requires-Dist: jax>=0.4.25
+Requires-Dist: jaxlib>=0.4.25
+
+# Equimo: Modern Vision Models in JAX/Equinox
+
+**WARNING**: This is a research library implementing recent computer vision models. The implementations are based on paper descriptions and may not be exact replicas of the original implementations. Use with caution in production environments.
+
+Equimo (Equinox Image Models) provides JAX/Equinox implementations of recent computer vision models, currently focusing (but not limited to) on transformer and state-space architectures.
+
+## Features
+
+- Pure JAX/Equinox implementations
+- Focus on recent architectures (2023-2024 papers)
+- Modular design for easy experimentation
+- Extensive documentation and type hints
+
+## Installation
+
+### From PyPI
+
+```bash
+pip install equimo
+```
+
+### From Source
+
+```bash
+git clone https://github.com/clementpoiret/equimo.git
+cd equimo
+pip install -e .
+```
+
+## Implemented Models
+
+| Model | Paper | Year | Status |
+|-------|-------|------|--------|
+| FasterViT | [FasterViT: Fast Vision Transformers with Hierarchical Attention](https://arxiv.org/abs/2306.06189) | 2023 | ✅ |
+| Castling-ViT | [Castling-ViT: Compressing Self-Attention via Switching Towards Linear-Angular Attention During Vision Transformer Inference](https://arxiv.org/abs/2211.10526) | 2023 | Partial* |
+| MLLA | [Mamba-like Linear Attention](https://arxiv.org/abs/2405.16605) | 2024 | ✅ |
+| PartialFormer | [Efficient Vision Transformers with Partial Attention](https://eccv.ecva.net/virtual/2024/poster/1877) | 2024 | ✅ |
+| SHViT | [SHViT: Single-Head Vision Transformer with Memory Efficient Macro Design](https://arxiv.org/abs/2401.16456) | 2024 | ✅ |
+| VSSD | [VSSD: Vision Mamba with Non-Causal State Space Duality](https://arxiv.org/abs/2407.18559) | 2024 | ✅ |
+
+*: Only contains the [Linear Angular Attention](https://github.com/clementpoiret/Equimo/blob/f8fcc79e45ca65e9deb1d970c4286c0b8562f9c2/equimo/layers/attention.py#L1407) module. It is straight forward to build a ViT around it, but may require an additional `__call__` kwarg to control the `sparse_reg` bool.
+
+## Basic Usage
+
+```python
+import jax
+
+import equimo.models as em
+
+# Create a model (e.g. `faster_vit_0_224`)
+key = jax.random.PRNGKey(0)
+model = em.FasterViT(
+    img_size=224,
+    in_channels=3,
+    dim=64,
+    in_dim=64,
+    depths=[2, 3, 6, 5],
+    num_heads=[2, 4, 8, 16],
+    hat=[False, False, True, False],
+    window_size=[7, 7, 7, 7],
+    ct_size=2,
+    key=key,
+)
+
+# Generate random input
+x = jax.random.normal(key, (3, 224, 224))
+
+# Run inference
+output = model(x, enable_dropout=False, key=key)
+```
+
+## Contributing
+
+Contributions are welcome! Please feel free to submit a Pull Request. For major changes, please open an issue first to discuss what you would like to change.
+
+## License
+
+This project is licensed under the MIT License - see the LICENSE file for details.
+
+## Citation
+
+If you use Equimo in your research, please cite:
+
+```bibtex
+@software{equimo2024,
+  author = {Clément POIRET},
+  title = {Equimo: Modern Vision Models in JAX/Equinox},
+  year = {2024},
+  publisher = {GitHub},
+  url = {https://github.com/clementpoiret/equimo}
+}
+```

Equimo-0.1.0.dist-info/RECORD

@@ -0,0 +1,29 @@
+equimo/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+equimo/utils.py,sha256=A6s6mRAvdnKWgTniZRYmSNNc0b7slAPFy5XwtSkl9dQ,3507
+equimo/layers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+equimo/layers/attention.py,sha256=BMMBIgMAgC24BZSeICqRoYJOd07QzhK6sLgXyJSZpqA,46132
+equimo/layers/convolution.py,sha256=HlnnoRkDXKOuUVWlNVvm2eCT_CPU3a68-FZFU_4dwYQ,10529
+equimo/layers/downsample.py,sha256=z3l6iiww8-k976q-WlbYk1iaSeyniMg7RD_Q-BwZV-Y,1691
+equimo/layers/dropout.py,sha256=wXVccY87q5080nJKmQycw75QYwna8V-zUb0OMpoNYAw,7457
+equimo/layers/ffn.py,sha256=fBzeQMXjAaUmSZ7F3ann5QT13apF5bNhPVyXvbiaYeY,5469
+equimo/layers/generic.py,sha256=o2RzVxKj74Lkr6ADtIq8XD0ujLBbChDV5ALMHK3GaBY,2159
+equimo/layers/mamba.py,sha256=WupMpWZ4OmtVMyhP_iu5_s3CAUBvnDDqzPmkg-GKTm8,5567
+equimo/layers/norm.py,sha256=TwqT2DmV7ljEiIHTtBz3s8adx99VHItivy6UmJzjPtM,2425
+equimo/layers/patch.py,sha256=xHTzIalPXuq_EisMZP3RpR2Ks5R7qIz6tF0FXlGCbyg,12634
+equimo/layers/posemb.py,sha256=eAa7yc4MxeGJGadApJrFnEYh4h275eF53YxV4_QM0Go,13392
+equimo/layers/sharing.py,sha256=brqIRjXv1rytqIejt1AB0sDXdL1378c0Ubax5ZnNZfs,4003
+equimo/layers/squeeze_excite.py,sha256=BVhuX2s_sQStdKd-Xg3sn76d7r1AW8JmRlVHtN7LTtM,4142
+equimo/models/__init__.py,sha256=6-CNU6Ww5_abbfOJJrcdVgsCf584i3QbgL0w6ZsjhVA,211
+equimo/models/emamodel.py,sha256=v4eYSUJVl5nV6P2_HNc1L2715DUJfuWQfJwogCxvS5M,2916
+equimo/models/fastervit.py,sha256=Vj-IcVmtGUcjE-i8YnUJAqrQUYmy6MhIWSgHMyuy7Go,15061
+equimo/models/mlla.py,sha256=9EY51VUJT1SzwqpdRz-UrCLJaI5kmqmDUE27DJa9jio,5519
+equimo/models/partialformer.py,sha256=-QyvgFGvJTU67DbFQGTLV8K6si51E3y4FGECyVUrngk,12813
+equimo/models/shvit.py,sha256=3ljwoAMCPxA5y09pSVqrO1kKIYHegRUIYtmVyeFlpO4,11028
+equimo/models/vit.py,sha256=ZWsLlvwyDJMLd7twMhJgOcHQyaXBjMgcD49-cQyM-d8,15885
+equimo/models/vssd.py,sha256=_UByQTnfmRi_IyNQgAh9Ddipr7QG-d5rFK_srkT4tls,6284
+equimo/ops/scan.py,sha256=1v80gkIkmQ0QoUj5YJpVQQ7ZjF3AfPfQchGLqSOqf4g,8729
+Equimo-0.1.0.dist-info/LICENSE.md,sha256=NDoiESQBdng_tbHOUj6BvMHpxu5P8JtsRbqVz-4rNIM,1072
+Equimo-0.1.0.dist-info/METADATA,sha256=6v5O8AiLaWdj9yX0F_uBTFAIK1OFEBlirVWnjEhLLfY,3412
+Equimo-0.1.0.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+Equimo-0.1.0.dist-info/top_level.txt,sha256=wfqgodqy_2pxpHWYoY89Fzr2I0-zPOER8zQ5bHWUxH0,7
+Equimo-0.1.0.dist-info/RECORD,,

Equimo-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.6.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

Equimo-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+equimo