bidirectional-mamba 0.2.0__tar.gz

bidirectional_mamba-0.2.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 redEddie
+
+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.

bidirectional_mamba-0.2.0/MANIFEST.in

@@ -0,0 +1,17 @@
+include LICENSE
+include README.md
+include pyproject.toml
+
+recursive-include bidirectional_mamba *.py
+
+prune tests
+prune benchmarks
+prune .git
+prune .pytest_cache
+prune dist
+prune build
+prune *.egg-info
+
+global-exclude *.pyc
+global-exclude __pycache__
+global-exclude .DS_Store

bidirectional_mamba-0.2.0/PKG-INFO

@@ -0,0 +1,179 @@
+Metadata-Version: 2.4
+Name: bidirectional-mamba
+Version: 0.2.0
+Summary: Bidirectional Mamba-2 SSD block with SwiGLU FFN (Vim/VMamba-style merge).
+Author: redEddie
+License-Expression: MIT
+Project-URL: Repository, https://github.com/redEddie/bidirectional-mamba
+Keywords: mamba,ssm,bidirectional,vision-mamba,swiglu
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.1
+Requires-Dist: einops>=0.4
+Requires-Dist: mamba-ssm>=2.0
+Requires-Dist: causal-conv1d>=1.2
+Requires-Dist: matplotlib>=3.0
+Dynamic: license-file
+
+# bidirectional-mamba
+
+양방향 Mamba-2 SSD 블록에 SwiGLU FFN을 결합한 Vision-Mamba/VMamba 스타일 구현체입니다.  
+forward scan과 backward scan을 독립적으로 수행하고, 출력을 더한 뒤 gate와 out projection을 공유합니다.
+
+```
+zxbcdt    = in_proj(u)                  # [z | x | B | C | dt | dt_b], 하나의 projection
+y_fw      = SSD(conv  (x),     A,   D,   dt,       dt_bias  )
+y_bw      = SSD(conv_b(x_rev), A_b, D_b, dt_b_rev, dt_bias_b)_rev   # reversed
+y         = y_fw + y_bw                 # sum merge
+out       = out_proj(RMSNormGated(y, z))
+```
+
+블록 구조는 `pre-norm` 기반이며 `(mixer residual) + (SwiGLU FFN residual)`입니다.  
+스택 마지막에는 final RMSNorm이 붙습니다.
+
+## 설치
+
+CUDA + Triton 환경이 필요합니다 (`mamba-ssm`, `causal-conv1d` 필요). CPU-only는 지원하지 않습니다.
+
+```bash
+pip install torch einops matplotlib
+pip install mamba-ssm causal-conv1d --no-build-isolation
+pip install -e .
+```
+
+## 기본 사용법
+
+```python
+import torch
+from bidirectional_mamba import BiMambaStack
+
+stack = BiMambaStack(d_model=768, n_layers=4, d_state=64, headdim=64,
+                     chunk_size=49).cuda()
+
+x = torch.randn(2, 196, 768, device="cuda")   # [B, L, D]
+y = stack(x)                                  # [B, L, D]
+```
+
+## 주요 기능
+
+### 1. 방향별 분리 (`share_dt`, `share_bc`)
+
+기본적으로 forward/backward는 각자의 `A`, `D`, `dt_bias`, `conv1d`를 가집니다.  
+추가로 `dt` 투영과 `B/C` 투영도 방향별로 분리할 수 있습니다.
+
+```python
+# dt 분리: backward scan이 자신만의 dt를 가짐
+BiMambaStack(d_model=768, n_layers=4, share_dt=False)
+
+# B/C 분리: backward scan이 자신만의 B, C를 가짐
+BiMambaStack(d_model=768, n_layers=4, share_dt=False, share_bc=False)
+```
+
+### 2. Padding mask (`valid_mask`)
+
+가변 길이 배치를 다룰 때 반드시 `valid_mask`를 전달하세요.  
+padding이 없으면 backward scan이 padding 영역에서 시작해 실제 토큰까지 영향을 미칩니다.
+
+```python
+mask = torch.zeros(B, L, dtype=torch.bool, device="cuda")
+mask[i, :lengths[i]] = True
+
+y = stack(x, valid_mask=mask)
+```
+
+mask를 쓰면 padding 위치는 0을 출력하고 실제 토큰에는 영향을 주지 않습니다.
+
+### 3. 학습 건강 지표 (`return_stats`)
+
+`BiMambaStack`, `BiMambaBlock`, `BiMambaMixer` 모두 `return_stats=True`를 지원합니다.  
+loss 외에도 dt, A decay, B/C norm, forward/backward 기여도 등 Mamba 고유의 수치를 확인할 수 있습니다.
+
+```python
+out, stats = stack(x, return_stats=True)
+
+# stats 예시
+# {
+#   "layer0/mixer/dt/effective_mean": ...,
+#   "layer0/mixer/A/decay_mean": ...,
+#   "layer0/mixer/y_fw_bw_ratio": ...,
+#   ...
+# }
+```
+
+자세한 예제는 `benchmarks/log_mamba_stats.py`를 참고하세요.
+
+### 5. 구현 검증
+
+`benchmarks/verify_bidirectional.py`는 양방향 Mamba 구현이 제대로 되었는지 확인하는 독립 스크립트입니다.  
+flip equivariance, backward branch 사용 여부, explicit reference match, gradient flow, valid mask 등 8가지 기준을 검사합니다.
+
+```bash
+python benchmarks/verify_bidirectional.py
+```
+
+### 4. SSD token-mixing 행렬 (`mixing_matrix`)
+
+Mamba의 암묵적인 token-mixing 행렬 L을 attention map처럼 시각화할 수 있습니다.  
+forward scan은 하삼각, backward scan은 상삼각을 채웁니다.
+
+```python
+mixer = BiMambaMixer(d_model=192, d_state=32, headdim=32).cuda()
+L = mixer.mixing_matrix(x)  # [B, L, L]
+```
+
+시각화 예제는 `benchmarks/visualize_mixing_matrix.py`를 참고하세요.
+
+## CIFAR-10 벤치마크 결과
+
+**공통 설정**: patch_size=4, share_dt=False, share_bc=True, use_mlp=True, epochs=30, seed=0
+
+### Width scaling (n_layers=8 고정)
+
+| d_model | n_layers | headdim | params(추정) | test acc |
+|--------:|---------:|--------:|------------:|---------:|
+| 768     | 8        | 64      | ~67M        | 80.70%   |
+| 1024    | 8        | 128     | ~119M       | 81.76%   |
+| 4096    | 4        | 128     | ~944M       | 80.13%   |
+
+### Depth scaling (d_model=384 고정, LR per-depth scaling)
+
+| d_model | n_layers | headdim | params(추정) | test acc |
+|--------:|---------:|--------:|------------:|---------:|
+| 384     | 4        | 32      | ~8.5M       | 84.39%   |
+| 384     | 8        | 32      | ~17M        | 85.00%   |
+| 384     | 12       | 32      | ~25M        | 85.04%   |
+| 384     | 16       | 32      | ~34M        | 85.23%   |
+| 384     | 20       | 32      | ~42M        | 84.65%   |
+
+> d_model=384이 768/1024 width scaling보다 훨씬 적은 파라미터로 더 높은 성능을 냈습니다.  
+> 이는 현재 width scaling 설정(특히 4096)이 CIFAR-10에 비해 과도하게 크다는 것을 의미합니다.
+
+## 공개 API
+
+- `BiMambaMixer` — 양방향 SSD mixer (shared in/out projection, 방향별 conv + A + D + dt + dt_bias).
+- `BiMambaBlock` — pre-norm block: `mixer + SwiGLU FFN` residual.
+- `BiMambaStack` — `n`개 block + final RMSNorm.
+- `SwiGLUMLP` — 독립형 SwiGLU FFN (expand=8/3 기본).
+
+세 클래스 모두 `share_dt`, `share_bc`를 받고, `forward`는 `valid_mask`를 받습니다.  
+`BiMambaMixer/Block/Stack`은 추가로 `return_stats=True`를 받을 수 있습니다.
+
+## 참고
+
+- `A`/`A_b`, `D`/`D_b`, `dt_bias`/`dt_bias_b`, 두 `conv1d`, 그리고 (`share_dt=False`일 때) 두 `dt` 투영은 모두 묶여 있지 않습니다.
+- `in_proj`, `RMSNormGated` (z로 gate), `out_proj`는 공유됩니다.  
+  출력을 더한 뒤 gate를 적용하는 것은 `norm_before_gate=False`이므로 방향별로 따로 gate하는 것과 동일합니다.
+- `share_dt=True`는 이전 동작으로 돌아가 forward `dt`를 reversed해서 backward가 재사용합니다.  
+  softplus가 단조증가이고 per-head bias만 분리되므로 두 방향이 "어떤 토큰이 중요한지"에 대해 강제로 동의해야 합니다.
+- autoregressive 디코딩용 `step()`은 없습니다. 양방향 scan은 전체 시퀀스가 필요합니다.
+- `seq_idx` 기반 세그먼트 경계는 아직 구현되지 않았습니다.
+
+## 참고 문헌
+
+- **Vision Mamba (Vim)** — Zhu et al., *Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model*, 2024. [arXiv:2401.09417](https://arxiv.org/abs/2401.09417)
+- **VMamba** — Liu et al., *VMamba: Visual State Space Model*, 2024. [arXiv:2401.10166](https://arxiv.org/abs/2401.10166)
+- **SwiGLU FFN** — Shazeer, *GLU Variants Improve Transformer*, 2020. [arXiv:2002.05202](https://arxiv.org/abs/2002.05202)
+- **Mamba-2 / SSD** — Dao & Gu, *Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality*, 2024. [arXiv:2405.21060](https://arxiv.org/abs/2405.21060)

bidirectional_mamba-0.2.0/README.md

@@ -0,0 +1,159 @@
+# bidirectional-mamba
+
+양방향 Mamba-2 SSD 블록에 SwiGLU FFN을 결합한 Vision-Mamba/VMamba 스타일 구현체입니다.  
+forward scan과 backward scan을 독립적으로 수행하고, 출력을 더한 뒤 gate와 out projection을 공유합니다.
+
+```
+zxbcdt    = in_proj(u)                  # [z | x | B | C | dt | dt_b], 하나의 projection
+y_fw      = SSD(conv  (x),     A,   D,   dt,       dt_bias  )
+y_bw      = SSD(conv_b(x_rev), A_b, D_b, dt_b_rev, dt_bias_b)_rev   # reversed
+y         = y_fw + y_bw                 # sum merge
+out       = out_proj(RMSNormGated(y, z))
+```
+
+블록 구조는 `pre-norm` 기반이며 `(mixer residual) + (SwiGLU FFN residual)`입니다.  
+스택 마지막에는 final RMSNorm이 붙습니다.
+
+## 설치
+
+CUDA + Triton 환경이 필요합니다 (`mamba-ssm`, `causal-conv1d` 필요). CPU-only는 지원하지 않습니다.
+
+```bash
+pip install torch einops matplotlib
+pip install mamba-ssm causal-conv1d --no-build-isolation
+pip install -e .
+```
+
+## 기본 사용법
+
+```python
+import torch
+from bidirectional_mamba import BiMambaStack
+
+stack = BiMambaStack(d_model=768, n_layers=4, d_state=64, headdim=64,
+                     chunk_size=49).cuda()
+
+x = torch.randn(2, 196, 768, device="cuda")   # [B, L, D]
+y = stack(x)                                  # [B, L, D]
+```
+
+## 주요 기능
+
+### 1. 방향별 분리 (`share_dt`, `share_bc`)
+
+기본적으로 forward/backward는 각자의 `A`, `D`, `dt_bias`, `conv1d`를 가집니다.  
+추가로 `dt` 투영과 `B/C` 투영도 방향별로 분리할 수 있습니다.
+
+```python
+# dt 분리: backward scan이 자신만의 dt를 가짐
+BiMambaStack(d_model=768, n_layers=4, share_dt=False)
+
+# B/C 분리: backward scan이 자신만의 B, C를 가짐
+BiMambaStack(d_model=768, n_layers=4, share_dt=False, share_bc=False)
+```
+
+### 2. Padding mask (`valid_mask`)
+
+가변 길이 배치를 다룰 때 반드시 `valid_mask`를 전달하세요.  
+padding이 없으면 backward scan이 padding 영역에서 시작해 실제 토큰까지 영향을 미칩니다.
+
+```python
+mask = torch.zeros(B, L, dtype=torch.bool, device="cuda")
+mask[i, :lengths[i]] = True
+
+y = stack(x, valid_mask=mask)
+```
+
+mask를 쓰면 padding 위치는 0을 출력하고 실제 토큰에는 영향을 주지 않습니다.
+
+### 3. 학습 건강 지표 (`return_stats`)
+
+`BiMambaStack`, `BiMambaBlock`, `BiMambaMixer` 모두 `return_stats=True`를 지원합니다.  
+loss 외에도 dt, A decay, B/C norm, forward/backward 기여도 등 Mamba 고유의 수치를 확인할 수 있습니다.
+
+```python
+out, stats = stack(x, return_stats=True)
+
+# stats 예시
+# {
+#   "layer0/mixer/dt/effective_mean": ...,
+#   "layer0/mixer/A/decay_mean": ...,
+#   "layer0/mixer/y_fw_bw_ratio": ...,
+#   ...
+# }
+```
+
+자세한 예제는 `benchmarks/log_mamba_stats.py`를 참고하세요.
+
+### 5. 구현 검증
+
+`benchmarks/verify_bidirectional.py`는 양방향 Mamba 구현이 제대로 되었는지 확인하는 독립 스크립트입니다.  
+flip equivariance, backward branch 사용 여부, explicit reference match, gradient flow, valid mask 등 8가지 기준을 검사합니다.
+
+```bash
+python benchmarks/verify_bidirectional.py
+```
+
+### 4. SSD token-mixing 행렬 (`mixing_matrix`)
+
+Mamba의 암묵적인 token-mixing 행렬 L을 attention map처럼 시각화할 수 있습니다.  
+forward scan은 하삼각, backward scan은 상삼각을 채웁니다.
+
+```python
+mixer = BiMambaMixer(d_model=192, d_state=32, headdim=32).cuda()
+L = mixer.mixing_matrix(x)  # [B, L, L]
+```
+
+시각화 예제는 `benchmarks/visualize_mixing_matrix.py`를 참고하세요.
+
+## CIFAR-10 벤치마크 결과
+
+**공통 설정**: patch_size=4, share_dt=False, share_bc=True, use_mlp=True, epochs=30, seed=0
+
+### Width scaling (n_layers=8 고정)
+
+| d_model | n_layers | headdim | params(추정) | test acc |
+|--------:|---------:|--------:|------------:|---------:|
+| 768     | 8        | 64      | ~67M        | 80.70%   |
+| 1024    | 8        | 128     | ~119M       | 81.76%   |
+| 4096    | 4        | 128     | ~944M       | 80.13%   |
+
+### Depth scaling (d_model=384 고정, LR per-depth scaling)
+
+| d_model | n_layers | headdim | params(추정) | test acc |
+|--------:|---------:|--------:|------------:|---------:|
+| 384     | 4        | 32      | ~8.5M       | 84.39%   |
+| 384     | 8        | 32      | ~17M        | 85.00%   |
+| 384     | 12       | 32      | ~25M        | 85.04%   |
+| 384     | 16       | 32      | ~34M        | 85.23%   |
+| 384     | 20       | 32      | ~42M        | 84.65%   |
+
+> d_model=384이 768/1024 width scaling보다 훨씬 적은 파라미터로 더 높은 성능을 냈습니다.  
+> 이는 현재 width scaling 설정(특히 4096)이 CIFAR-10에 비해 과도하게 크다는 것을 의미합니다.
+
+## 공개 API
+
+- `BiMambaMixer` — 양방향 SSD mixer (shared in/out projection, 방향별 conv + A + D + dt + dt_bias).
+- `BiMambaBlock` — pre-norm block: `mixer + SwiGLU FFN` residual.
+- `BiMambaStack` — `n`개 block + final RMSNorm.
+- `SwiGLUMLP` — 독립형 SwiGLU FFN (expand=8/3 기본).
+
+세 클래스 모두 `share_dt`, `share_bc`를 받고, `forward`는 `valid_mask`를 받습니다.  
+`BiMambaMixer/Block/Stack`은 추가로 `return_stats=True`를 받을 수 있습니다.
+
+## 참고
+
+- `A`/`A_b`, `D`/`D_b`, `dt_bias`/`dt_bias_b`, 두 `conv1d`, 그리고 (`share_dt=False`일 때) 두 `dt` 투영은 모두 묶여 있지 않습니다.
+- `in_proj`, `RMSNormGated` (z로 gate), `out_proj`는 공유됩니다.  
+  출력을 더한 뒤 gate를 적용하는 것은 `norm_before_gate=False`이므로 방향별로 따로 gate하는 것과 동일합니다.
+- `share_dt=True`는 이전 동작으로 돌아가 forward `dt`를 reversed해서 backward가 재사용합니다.  
+  softplus가 단조증가이고 per-head bias만 분리되므로 두 방향이 "어떤 토큰이 중요한지"에 대해 강제로 동의해야 합니다.
+- autoregressive 디코딩용 `step()`은 없습니다. 양방향 scan은 전체 시퀀스가 필요합니다.
+- `seq_idx` 기반 세그먼트 경계는 아직 구현되지 않았습니다.
+
+## 참고 문헌
+
+- **Vision Mamba (Vim)** — Zhu et al., *Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model*, 2024. [arXiv:2401.09417](https://arxiv.org/abs/2401.09417)
+- **VMamba** — Liu et al., *VMamba: Visual State Space Model*, 2024. [arXiv:2401.10166](https://arxiv.org/abs/2401.10166)
+- **SwiGLU FFN** — Shazeer, *GLU Variants Improve Transformer*, 2020. [arXiv:2002.05202](https://arxiv.org/abs/2002.05202)
+- **Mamba-2 / SSD** — Dao & Gu, *Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality*, 2024. [arXiv:2405.21060](https://arxiv.org/abs/2405.21060)

bidirectional_mamba-0.2.0/bidirectional_mamba/__init__.py

@@ -0,0 +1,20 @@
+from ._mlp import SwiGLUMLP
+
+__version__ = "0.2.0"
+__all__ = ["SwiGLUMLP"]
+
+# Mamba-2 SSD stack (requires causal_conv1d + mamba_ssm with SSD kernels)
+try:
+    from .block import BiMambaBlock
+    from .mixer import BiMambaMixer
+    from .stack import BiMambaStack
+    __all__.extend(["BiMambaMixer", "BiMambaBlock", "BiMambaStack"])
+except Exception:
+    BiMambaMixer = BiMambaBlock = BiMambaStack = None  # type: ignore
+
+# Mamba-3 bidirectional mixer (requires latest state-spaces/mamba source)
+try:
+    from .mixer3 import BiMamba3Mixer
+    __all__.append("BiMamba3Mixer")
+except Exception:
+    BiMamba3Mixer = None  # type: ignore

bidirectional_mamba-0.2.0/bidirectional_mamba/_mlp.py

@@ -0,0 +1,16 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SwiGLUMLP(nn.Module):
+    def __init__(self, d_model: int, expand: float = 8 / 3, device=None, dtype=None):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        d_hidden = int(d_model * expand)
+        d_hidden = (d_hidden + 7) // 8 * 8
+        self.w1 = nn.Linear(d_model, d_hidden, bias=False, **factory_kwargs)
+        self.w2 = nn.Linear(d_model, d_hidden, bias=False, **factory_kwargs)
+        self.w3 = nn.Linear(d_hidden, d_model, bias=False, **factory_kwargs)
+
+    def forward(self, x):
+        return self.w3(F.silu(self.w1(x)) * self.w2(x))

bidirectional_mamba-0.2.0/bidirectional_mamba/block.py

@@ -0,0 +1,48 @@
+import torch.nn as nn
+
+try:
+    from mamba_ssm.ops.triton.layer_norm import RMSNorm
+except ImportError:
+    RMSNorm = nn.LayerNorm
+
+from ._mlp import SwiGLUMLP
+from .mixer import BiMambaMixer
+
+_NORM_EPSILON = 1e-5
+
+
+class BiMambaBlock(nn.Module):
+    """Pre-norm bidirectional Mamba-2 block: mixer residual + SwiGLU FFN residual.
+
+        x = x + BiMambaMixer(RMSNorm(x))
+        x = x + SwiGLU    (RMSNorm(x))
+    """
+
+    def __init__(self, d_model, d_state=64, headdim=64, expand=2,
+                 chunk_size=49, dt_min=0.001, dt_max=0.02,
+                 mlp_expand: float = 8 / 3, share_dt=False, share_bc=True,
+                 device=None, dtype=None):
+        super().__init__()
+        factory_kwargs = {"device": device, "dtype": dtype}
+        self.norm1 = RMSNorm(d_model, eps=_NORM_EPSILON, **factory_kwargs)
+        self.mixer = BiMambaMixer(
+            d_model, d_state=d_state, headdim=headdim, expand=expand,
+            chunk_size=chunk_size, dt_min=dt_min, dt_max=dt_max,
+            share_dt=share_dt, share_bc=share_bc, **factory_kwargs,
+        )
+        self.norm2 = RMSNorm(d_model, eps=_NORM_EPSILON, **factory_kwargs)
+        self.mlp = SwiGLUMLP(d_model, expand=mlp_expand, **factory_kwargs)
+
+    def forward(self, x, valid_mask=None, return_stats=False):
+        mixer_in = self.norm1(x)
+        if return_stats:
+            mixer_out, mixer_stats = self.mixer(
+                mixer_in, valid_mask=valid_mask, return_stats=True)
+        else:
+            mixer_out = self.mixer(mixer_in, valid_mask=valid_mask)
+            mixer_stats = None
+        x = x + mixer_out
+        x = x + self.mlp(self.norm2(x))
+        if return_stats:
+            return x, {"mixer": mixer_stats}
+        return x

bidirectional_mamba-0.2.0/bidirectional_mamba/mixer.py

@@ -0,0 +1,384 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from causal_conv1d import causal_conv1d_fn
+from mamba_ssm.ops.triton.layernorm_gated import RMSNorm as RMSNormGated
+from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined
+
+
+class BiMambaMixer(nn.Module):
+    """Bidirectional Mamba-2 SSD mixer with shared in/out projection.
+
+        zxbcdt    = in_proj(u)                      # [z | x | B | C | dt]
+        zxbcdt    = in_proj(u)                      # [z | x | B | C | dt | dt_b]
+        y_fw      = SSD(conv  (x),      A,   D,   dt,       dt_bias  )
+        y_bw      = SSD(conv_b(x_rev),  A_b, D_b, dt_b_rev, dt_bias_b)_rev
+        y         = y_fw + y_bw                     # sum merge (Vim/VMamba)
+        out       = out_proj(RMSNormGated(y, z))
+
+    Pass `valid_mask` whenever the batch is padded. Right-padding is harmless
+    for a causal model, but the backward scan starts inside the padding and
+    carries it into every real position, so an unmasked padded batch gives a
+    different answer for the same sequence depending on how much it was padded.
+    """
+
+    def __init__(self, d_model, d_state=64, headdim=64, expand=2, d_conv=4,
+                 ngroups=1, chunk_size=49, dt_min=0.001, dt_max=0.02,
+                 dt_init_floor=1e-4, A_init_range=(1, 16), share_dt=False,
+                 share_bc=True, device=None, dtype=None):
+        super().__init__()
+        factory_kwargs = {"device": device, "dtype": dtype}
+        self.d_state = d_state
+        self.headdim = headdim
+        self.ngroups = ngroups
+        self.chunk_size = chunk_size
+        self.share_dt = share_dt
+        self.share_bc = share_bc
+        d_inner = expand * d_model
+        self.d_ssm = d_inner
+        assert d_inner % headdim == 0
+        self.nheads = d_inner // headdim
+
+        # dt is the only token-dependent quantity the two directions would
+        # otherwise share outright: x, B and C at least pass through a
+        # per-direction conv, but dt does not, so with share_dt the backward
+        # scan gets the forward dt reversed and can only shift it by a
+        # per-head bias -- softplus is monotonic, so the two directions are
+        # forced to agree on which tokens matter. A second projection costs
+        # d_model * nheads (0.5% of the mixer at d_model=768).
+        self._dt_splits = 1 if share_dt else 2
+        # B/C projections can also be separated per direction. When share_bc
+        # is True we keep the v0.2.0 layout [z | xBC | dt ...] for checkpoint
+        # compatibility; when False we emit [z | x | BC | BC_b | dt ...].
+        self._bc_splits = 1 if share_bc else 2
+        d_in_proj = (2 * d_inner + self._bc_splits * 2 * ngroups * d_state
+                     + self._dt_splits * self.nheads)
+        self.in_proj = nn.Linear(d_model, d_in_proj, bias=False, **factory_kwargs)
+
+        conv_dim = d_inner + 2 * ngroups * d_state
+        self.conv1d = self._make_conv(conv_dim, d_conv, factory_kwargs)
+        self.conv1d_b = self._make_conv(conv_dim, d_conv, factory_kwargs)
+
+        self.dt_bias = self._init_dt_bias(dt_min, dt_max, dt_init_floor, factory_kwargs)
+        self.dt_bias_b = self._init_dt_bias(dt_min, dt_max, dt_init_floor, factory_kwargs)
+        self.A_log = self._init_A_log(A_init_range, device, dtype)
+        self.A_log_b = self._init_A_log(A_init_range, device, dtype)
+        self.D = self._init_D(device)
+        self.D_b = self._init_D(device)
+
+        self.norm = RMSNormGated(d_inner, eps=1e-5, norm_before_gate=False,
+                                 group_size=d_inner // ngroups, **factory_kwargs)
+        self.out_proj = nn.Linear(d_inner, d_model, bias=False, **factory_kwargs)
+
+    @staticmethod
+    def _make_conv(conv_dim, d_conv, factory_kwargs):
+        return nn.Conv1d(conv_dim, conv_dim, bias=True, kernel_size=d_conv,
+                         groups=conv_dim, padding=d_conv - 1, **factory_kwargs)
+
+    def _init_dt_bias(self, dt_min, dt_max, dt_init_floor, factory_kwargs):
+        dt = torch.exp(
+            torch.rand(self.nheads, **factory_kwargs) * (math.log(dt_max) - math.log(dt_min))
+            + math.log(dt_min)
+        ).clamp(min=dt_init_floor)
+        inv_dt = dt + torch.log(-torch.expm1(-dt))
+        p = nn.Parameter(inv_dt)
+        p._no_weight_decay = True
+        return p
+
+    def _init_A_log(self, A_init_range, device, dtype):
+        A = torch.empty(self.nheads, dtype=torch.float32, device=device).uniform_(*A_init_range)
+        p = nn.Parameter(torch.log(A).to(dtype=dtype))
+        p._no_weight_decay = True
+        return p
+
+    def _init_D(self, device):
+        p = nn.Parameter(torch.ones(self.nheads, device=device))
+        p._no_weight_decay = True
+        return p
+
+    def _scan_dir(self, xBC, dt, A, conv, D, dt_bias, flip, mask=None,
+                  return_intermediates=False):
+        if flip:
+            xBC, dt = xBC.flip(1), dt.flip(1)
+            if mask is not None:
+                mask = mask.flip(1)
+        dt_softplus = True
+        dt_input = dt
+        if mask is not None:
+            # Make masked steps inert in the recurrence. dt = 0 gives both
+            # decay = exp(0 * A) = 1, so the state passes through untouched,
+            # and dt * B * x = 0, so nothing is written. Zeroing xBC keeps the
+            # conv window from picking the masked region up, which leaves a
+            # position at the boundary seeing the same zeros causal_conv1d_fn
+            # already left-pads with at a true sequence start.
+            # softplus has to move out of the kernel to land on exactly zero.
+            xBC = xBC * mask
+            dt = F.softplus(dt.float() + dt_bias.float()).to(dt.dtype) * mask
+            dt_effective = dt
+            dt_bias, dt_softplus = None, False
+        else:
+            # The kernel applies softplus + bias internally. For diagnostics
+            # we materialise the effective step size without touching the
+            # tensor fed to the kernel.
+            dt_effective = F.softplus(dt.float() + dt_bias.float()).to(dt.dtype)
+        xBC = causal_conv1d_fn(
+            xBC.transpose(1, 2).contiguous(),
+            rearrange(conv.weight, "d 1 w -> d w"), bias=conv.bias, activation="silu",
+        ).transpose(1, 2)
+        x, B, C = torch.split(
+            xBC, [self.d_ssm, self.ngroups * self.d_state, self.ngroups * self.d_state], dim=-1)
+        y = mamba_chunk_scan_combined(
+            rearrange(x, "b l (h p) -> b l h p", p=self.headdim),
+            dt, A,
+            rearrange(B, "b l (g n) -> b l g n", g=self.ngroups),
+            rearrange(C, "b l (g n) -> b l g n", g=self.ngroups),
+            chunk_size=self.chunk_size, D=D, z=None, dt_bias=dt_bias,
+            dt_softplus=dt_softplus,
+        )
+        y = y.flip(1) if flip else y
+        if not return_intermediates:
+            return y
+        return y, {
+            "dt_raw": dt_input,
+            "dt_effective": dt_effective,
+            "x": x,
+            "B": B,
+            "C": C,
+            "xBC": xBC,
+        }
+
+    @staticmethod
+    def _broadcast_mask(mask, t):
+        """Broadcast `mask` to the same shape as `t` and return a bool tensor."""
+        if mask is None:
+            return None
+        while mask.dim() < t.dim():
+            mask = mask.unsqueeze(-1)
+        return mask.expand(t.shape).to(torch.bool)
+
+    @staticmethod
+    def _masked_stats(t, mask=None, dim=None):
+        """Return mean/std/min/max of `t` over unmasked positions.
+
+        `mask`: bool or float tensor broadcastable to `t`, True/non-zero on
+        real positions. If `mask` is None, all positions are used.
+        """
+        if mask is not None:
+            t = t[BiMambaMixer._broadcast_mask(mask, t)]
+        if t.numel() == 0:
+            nan = torch.tensor(float("nan"), device=t.device, dtype=t.dtype)
+            return nan, nan, nan, nan
+        mean = t.mean()
+        std = t.std(unbiased=False)
+        mn = t.min()
+        mx = t.max()
+        return mean.detach(), std.detach(), mn.detach(), mx.detach()
+
+    @staticmethod
+    def _masked_norm_mean(t, mask=None):
+        if mask is not None:
+            t = t[BiMambaMixer._broadcast_mask(mask, t)]
+        if t.numel() == 0:
+            return torch.tensor(float("nan"), device=t.device, dtype=t.dtype)
+        return t.norm(dim=-1).mean().detach()
+
+    def _compute_stats(self, z, A_decay, A_b_decay, fw, bw, mask=None):
+        """Build a detached statistics dict from scan intermediates.
+
+        `fw` and `bw` are the tuples returned by `_scan_dir` with
+        return_intermediates=True: (y, dict).
+        """
+        y_fw, fw_i = fw
+        y_bw, bw_i = bw
+
+        dt_raw_mean, dt_raw_std, dt_raw_min, dt_raw_max = self._masked_stats(
+            fw_i["dt_raw"], mask=mask)
+        dt_eff_mean, dt_eff_std, dt_eff_min, dt_eff_max = self._masked_stats(
+            fw_i["dt_effective"], mask=mask)
+        dt_b_raw_mean, dt_b_raw_std, _, _ = self._masked_stats(
+            bw_i["dt_raw"], mask=mask)
+        dt_b_eff_mean, dt_b_eff_std, dt_b_eff_min, dt_b_eff_max = self._masked_stats(
+            bw_i["dt_effective"], mask=mask)
+
+        A_mean, A_std, A_min, A_max = self._masked_stats(A_decay)
+        A_b_mean, A_b_std, _, _ = self._masked_stats(A_b_decay)
+
+        dt_bias_mean, dt_bias_std, _, _ = self._masked_stats(self.dt_bias)
+        dt_bias_b_mean, dt_bias_b_std, _, _ = self._masked_stats(self.dt_bias_b)
+
+        z_mean, z_std, _, _ = self._masked_stats(z, mask=mask)
+        y_fw_norm_mean = self._masked_norm_mean(y_fw, mask=mask)
+        y_bw_norm_mean = self._masked_norm_mean(y_bw, mask=mask)
+
+        stats = {
+            "dt/raw_mean": dt_raw_mean,
+            "dt/raw_std": dt_raw_std,
+            "dt/raw_min": dt_raw_min,
+            "dt/raw_max": dt_raw_max,
+            "dt/effective_mean": dt_eff_mean,
+            "dt/effective_std": dt_eff_std,
+            "dt/effective_min": dt_eff_min,
+            "dt/effective_max": dt_eff_max,
+            "dt_b/raw_mean": dt_b_raw_mean,
+            "dt_b/raw_std": dt_b_raw_std,
+            "dt_b/effective_mean": dt_b_eff_mean,
+            "dt_b/effective_std": dt_b_eff_std,
+            "dt_b/effective_min": dt_b_eff_min,
+            "dt_b/effective_max": dt_b_eff_max,
+            "A/decay_mean": A_mean,
+            "A/decay_std": A_std,
+            "A/decay_min": A_min,
+            "A/decay_max": A_max,
+            "A_b/decay_mean": A_b_mean,
+            "A_b/decay_std": A_b_std,
+            "dt_bias/mean": dt_bias_mean,
+            "dt_bias/std": dt_bias_std,
+            "dt_bias_b/mean": dt_bias_b_mean,
+            "dt_bias_b/std": dt_bias_b_std,
+            "x/norm_mean": self._masked_norm_mean(fw_i["x"], mask=mask),
+            "x_b/norm_mean": self._masked_norm_mean(bw_i["x"], mask=mask),
+            "B/norm_mean": self._masked_norm_mean(fw_i["B"], mask=mask),
+            "B_b/norm_mean": self._masked_norm_mean(bw_i["B"], mask=mask),
+            "C/norm_mean": self._masked_norm_mean(fw_i["C"], mask=mask),
+            "C_b/norm_mean": self._masked_norm_mean(bw_i["C"], mask=mask),
+            "conv_out/norm_mean": self._masked_norm_mean(fw_i["xBC"], mask=mask),
+            "conv_out_b/norm_mean": self._masked_norm_mean(bw_i["xBC"], mask=mask),
+            "z/mean": z_mean,
+            "z/std": z_std,
+            "y_fw/norm_mean": y_fw_norm_mean,
+            "y_bw/norm_mean": y_bw_norm_mean,
+        }
+        ratio = (y_fw_norm_mean / (y_bw_norm_mean + 1e-12)).item()
+        stats["y_fw_bw_ratio"] = torch.tensor(ratio, device=y_fw.device)
+        return stats
+
+    def _split_inputs(self, u):
+        """Split in_proj output into (z, xBC, xBC_b, dt, dt_b)."""
+        if self.share_bc:
+            sizes = ([self.d_ssm, self.d_ssm + 2 * self.ngroups * self.d_state]
+                     + [self.nheads] * self._dt_splits)
+            z, xBC, dt, *rest = torch.split(self.in_proj(u), sizes, dim=-1)
+            xBC_b = xBC
+            dt_b = dt if self.share_dt else rest[0]
+        else:
+            sizes = ([self.d_ssm, self.d_ssm]
+                     + [2 * self.ngroups * self.d_state] * self._bc_splits
+                     + [self.nheads] * self._dt_splits)
+            z, x, BC, BC_b, dt, *rest = torch.split(self.in_proj(u), sizes, dim=-1)
+            dt_b = dt if self.share_dt else rest[0]
+            xBC = torch.cat([x, BC], dim=-1)
+            xBC_b = torch.cat([x, BC_b], dim=-1)
+        return z, xBC, xBC_b, dt, dt_b
+
+    def _mixing_matrix_dir(self, intermediates, A, flip=False):
+        """Materialize the SSM token-mixing matrix for one scan direction.
+
+        `intermediates` is the dict returned by `_scan_dir(...,
+        return_intermediates=True)`.  The caller must pass the intermediates
+        in the scan's internal coordinates: for `flip=True`, `_scan_dir`
+        already flipped dt/B/C, so this function only flips the output matrix
+        back to the original sequence order.
+
+        Returns a [B, L, L, H] tensor where entry [b, t, s, h] is the
+        contribution of position s to position t for head h.
+        """
+        dt = intermediates["dt_effective"]                         # [B, L, H]
+        B = rearrange(intermediates["B"], "b l (g n) -> b l g n",
+                      g=self.ngroups)
+        C = rearrange(intermediates["C"], "b l (g n) -> b l g n",
+                      g=self.ngroups)
+        # Broadcast B/C from groups to heads
+        B_h = B.repeat_interleave(self.nheads // self.ngroups, dim=2)  # [B,L,H,N]
+        C_h = C.repeat_interleave(self.nheads // self.ngroups, dim=2)
+        # log decay per step, per head, per state dim
+        log_decay = (dt.unsqueeze(-1) *                               # [B,L,H,1]
+                     A.view(1, 1, self.nheads, 1).float())            # [1,1,H,1]
+        log_decay_cumsum = torch.cumsum(log_decay, dim=1)             # [B,L,H,N]
+        B, L, H, N = B_h.shape
+        cum_t = log_decay_cumsum.unsqueeze(2)                         # [B,L,1,H,N]
+        cum_s = log_decay_cumsum.unsqueeze(1)                         # [B,1,L,H,N]
+        decay_diff = torch.exp(cum_t - cum_s)                         # [B,L,L,H,N]
+        # C_t @ (decay) @ B_s, summed over state dim
+        contrib = (C_h.unsqueeze(2) * B_h.unsqueeze(1) * decay_diff).sum(-1)
+        # multiply by dt_s
+        contrib = contrib * dt.unsqueeze(1)                           # [B,1,L,H]
+        # causal/anti-causal mask applied by caller via flip
+        tri = torch.tril(torch.ones(L, L, device=contrib.device)).view(1, L, L, 1)
+        L_mat = contrib * tri
+        if flip:
+            L_mat = L_mat.flip(1).flip(2)
+        return L_mat
+
+    def mixing_matrix(self, u, valid_mask=None, average_heads=True,
+                      normalize_rows=False):
+        """Materialize the bidirectional SSM token-mixing matrix.
+
+        Returns a tensor of shape `[B, L, L]` (or `[B, L, L, H]` if
+        `average_heads=False`) where entry `[b, t, s]` is the SSM contribution
+        of input token s to output token t.  The forward scan fills the lower
+        triangle and the backward scan fills the upper triangle.
+
+        This is the analogue of an attention map for SSD: it shows how far
+        each position attends to past/future tokens, and whether recency bias
+        is present (diagonal-near weights dominating).
+        """
+        A = -torch.exp(self.A_log.float())
+        A_b = -torch.exp(self.A_log_b.float())
+        z, xBC, xBC_b, dt, dt_b = self._split_inputs(u)
+        mask = None if valid_mask is None else valid_mask.unsqueeze(-1).to(xBC.dtype)
+        fw = self._scan_dir(xBC, dt, A, self.conv1d, self.D, self.dt_bias,
+                            flip=False, mask=mask, return_intermediates=True)
+        bw = self._scan_dir(xBC_b, dt_b, A_b, self.conv1d_b, self.D_b,
+                            self.dt_bias_b, flip=True, mask=mask,
+                            return_intermediates=True)
+        L = self._mixing_matrix_dir(fw[1], A, flip=False) + \
+            self._mixing_matrix_dir(bw[1], A_b, flip=True)
+        if average_heads:
+            L = L.mean(-1)                                            # [B,L,L]
+        if normalize_rows:
+            denom = L.abs().sum(dim=-1, keepdim=True).clamp_min(1e-12)
+            L = L / denom
+        return L.detach()
+
+    def forward(self, u, valid_mask=None, return_stats=False):
+        """`u`: [B, L, d_model]. `valid_mask`: [B, L] bool, True on real tokens.
+
+        Masked positions neither influence real ones nor carry an output; the
+        result on the real positions is identical to running them alone.
+
+        If `return_stats=True`, returns `(output, stats_dict)` where the stats
+        are detached diagnostics about dt, A decay, B/C norms, forward/backward
+        scan contributions, etc.  This keeps intermediate tensors alive and is
+        intended for logging, not for every training step.
+        """
+        A = -torch.exp(self.A_log.float())
+        A_b = -torch.exp(self.A_log_b.float())
+        z, xBC, xBC_b, dt, dt_b = self._split_inputs(u)
+        mask = None if valid_mask is None else valid_mask.unsqueeze(-1).to(xBC.dtype)
+
+        if return_stats:
+            fw = self._scan_dir(xBC, dt, A, self.conv1d, self.D, self.dt_bias,
+                                flip=False, mask=mask, return_intermediates=True)
+            bw = self._scan_dir(xBC_b, dt_b, A_b, self.conv1d_b, self.D_b,
+                                self.dt_bias_b, flip=True, mask=mask,
+                                return_intermediates=True)
+            y_fw, y_bw = fw[0], bw[0]
+            stats = self._compute_stats(z, A, A_b, fw, bw, mask=mask)
+        else:
+            y_fw = self._scan_dir(xBC, dt, A, self.conv1d, self.D, self.dt_bias,
+                                  flip=False, mask=mask)
+            y_bw = self._scan_dir(xBC_b, dt_b, A_b, self.conv1d_b, self.D_b,
+                                  self.dt_bias_b, flip=True, mask=mask)
+            stats = None
+
+        y = y_fw + y_bw
+        y = self.norm(rearrange(y, "b l h p -> b l (h p)"), z)
+        out = self.out_proj(y)
+        out = out if mask is None else out * mask
+        if return_stats:
+            return out, stats
+        return out

bidirectional_mamba-0.2.0/bidirectional_mamba/mixer3.py

@@ -0,0 +1,58 @@
+"""Bidirectional Mamba-3 mixer (prototype).
+
+Requires the latest ``state-spaces/mamba`` source install (mamba-ssm >= 2.3.x)
+which provides ``mamba_ssm.modules.mamba3.Mamba3``.  Only the SISO path is used
+here because MIMO mode needs TileLang, which is not required for encoders.
+
+This is intentionally a separate module so that the rest of the package still
+loads on environments that only have Mamba-2 installed.
+"""
+
+import torch
+import torch.nn as nn
+
+from mamba_ssm.modules.mamba3 import Mamba3
+
+
+class BiMamba3Mixer(nn.Module):
+    """Bidirectional Mamba-3 SISO mixer: forward scan + reversed backward scan.
+
+    Uses two independent ``Mamba3`` instances, so the parameter count is roughly
+    twice that of a single causal Mamba-3 block.  This is a feasibility prototype;
+    a production version would share ``in_proj`` and separate only the
+    direction-dependent quantities (B, C, dt, A).
+    """
+
+    def __init__(self, d_model, d_state=64, headdim=64, expand=2,
+                 chunk_size=64, dt_min=0.001, dt_max=0.1,
+                 dt_init_floor=1e-4, A_floor=1e-4, rope_fraction=0.5,
+                 device=None, dtype=None):
+        super().__init__()
+        factory_kwargs = {"device": device, "dtype": dtype}
+        self.d_model = d_model
+        self.d_inner = int(expand * d_model)
+        assert self.d_inner % headdim == 0
+        self.nheads = self.d_inner // headdim
+
+        common = {
+            "d_state": d_state,
+            "headdim": headdim,
+            "expand": expand,
+            "chunk_size": chunk_size,
+            "dt_min": dt_min,
+            "dt_max": dt_max,
+            "dt_init_floor": dt_init_floor,
+            "A_floor": A_floor,
+            "rope_fraction": rope_fraction,
+            "is_mimo": False,
+            "is_outproj_norm": False,
+            **factory_kwargs,
+        }
+        self.mamba_fw = Mamba3(d_model, **common)
+        self.mamba_bw = Mamba3(d_model, **common)
+
+    def forward(self, u):
+        """`u`: [B, L, d_model]. Returns: [B, L, d_model]."""
+        y_fw = self.mamba_fw(u)
+        y_bw = self.mamba_bw(u.flip(1)).flip(1)
+        return y_fw + y_bw

bidirectional_mamba-0.2.0/bidirectional_mamba/stack.py

@@ -0,0 +1,53 @@
+import torch.nn as nn
+
+try:
+    from mamba_ssm.ops.triton.layer_norm import RMSNorm
+except ImportError:
+    RMSNorm = nn.LayerNorm
+
+from .block import BiMambaBlock, _NORM_EPSILON
+
+
+class BiMambaStack(nn.Module):
+    """Stack of `n_layers` BiMambaBlocks + final RMSNorm.
+
+    Input/output: `[B, L, d_model]`.
+    """
+
+    def __init__(self, d_model, n_layers=4, d_state=64, headdim=64,
+                 expand=2, chunk_size=49, dt_min=0.001, dt_max=0.02,
+                 mlp_expand: float = 8 / 3, share_dt=False, share_bc=True,
+                 device=None, dtype=None):
+        super().__init__()
+        factory_kwargs = {"device": device, "dtype": dtype}
+        self.layers = nn.ModuleList([
+            BiMambaBlock(
+                d_model, d_state=d_state, headdim=headdim, expand=expand,
+                chunk_size=chunk_size, dt_min=dt_min, dt_max=dt_max,
+                mlp_expand=mlp_expand, share_dt=share_dt, share_bc=share_bc,
+                **factory_kwargs,
+            )
+            for _ in range(n_layers)
+        ])
+        self.norm_f = RMSNorm(d_model, eps=_NORM_EPSILON, **factory_kwargs)
+
+    def forward(self, x, valid_mask=None, return_stats=False):
+        """`valid_mask`: [B, L] bool, True on real tokens. Required for padded
+        batches -- see `BiMambaMixer`.
+
+        If `return_stats=True`, returns `(output, stats_dict)` with per-layer
+        mixer diagnostics prefixed by `layer{i}/` plus `final_norm/...`.
+        """
+        stats = {} if return_stats else None
+        for i, layer in enumerate(self.layers):
+            if return_stats:
+                x, layer_stats = layer(x, valid_mask=valid_mask, return_stats=True)
+                stats[f"layer{i}/mixer"] = layer_stats["mixer"]
+            else:
+                x = layer(x, valid_mask=valid_mask)
+        x = self.norm_f(x)
+        if return_stats:
+            stats["final_norm/mean"] = x.mean().detach()
+            stats["final_norm/std"] = x.std(unbiased=False).detach()
+            return x, stats
+        return x

bidirectional_mamba-0.2.0/bidirectional_mamba.egg-info/SOURCES.txt

@@ -0,0 +1,10 @@
+LICENSE
+MANIFEST.in
+README.md
+pyproject.toml
+bidirectional_mamba/__init__.py
+bidirectional_mamba/_mlp.py
+bidirectional_mamba/block.py
+bidirectional_mamba/mixer.py
+bidirectional_mamba/mixer3.py
+bidirectional_mamba/stack.py

bidirectional_mamba-0.2.0/pyproject.toml

@@ -0,0 +1,31 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "bidirectional-mamba"
+version = "0.2.0"
+description = "Bidirectional Mamba-2 SSD block with SwiGLU FFN (Vim/VMamba-style merge)."
+readme = "README.md"
+requires-python = ">=3.10"
+license = "MIT"
+authors = [{ name = "redEddie" }]
+keywords = ["mamba", "ssm", "bidirectional", "vision-mamba", "swiglu"]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "torch>=2.1",
+    "einops>=0.4",
+    "mamba-ssm>=2.0",
+    "causal-conv1d>=1.2",
+    "matplotlib>=3.0",
+]
+
+[project.urls]
+Repository = "https://github.com/redEddie/bidirectional-mamba"
+
+[tool.setuptools.packages.find]
+include = ["bidirectional_mamba*"]
+exclude = ["tests*"]

bidirectional_mamba-0.2.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+