stackformers 3.8.0__tar.gz → 3.8.2__tar.gz

{stackformers-3.8.0 → stackformers-3.8.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stackformers
-Version: 3.8.0
+Version: 3.8.2
 Summary: Typed, composable, SOLID transformer library for PyTorch
 Project-URL: Homepage, https://github.com/Red-Eyed/stackformers
 Project-URL: Repository, https://github.com/Red-Eyed/stackformers

{stackformers-3.8.0 → stackformers-3.8.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "stackformers"
-version = "3.8.0"
+version = "3.8.2"
 description = "Typed, composable, SOLID transformer library for PyTorch"
 readme = "README.md"
 license = { file = "LICENSE" }

{stackformers-3.8.0 → stackformers-3.8.2}/stackformers/attention/ops.py

@@ -159,7 +159,9 @@ def packed_attn_or_fallback(
                 stacklevel=2,
             )
 
-    b = int(q_seq.cu_seqlens.shape[0]) - 1
+    # Keep b as SymInt-friendly: int(SymInt) forces specialization under torch.export.
+    # Downstream consumers (_packed_heads_to_padded) accept SymInt for shape args.
+    b = q_seq.cu_seqlens.shape[0] - 1
     q_pad, q_mask = _packed_heads_to_padded(q, q_seq.cu_seqlens, b, q_seq.max_seqlen)
     k_pad, k_mask = _packed_heads_to_padded(k, k_seq.cu_seqlens, b, k_seq.max_seqlen)
     v_pad, _ = _packed_heads_to_padded(v, k_seq.cu_seqlens, b, k_seq.max_seqlen)

{stackformers-3.8.0 → stackformers-3.8.2}/stackformers/sequence.py

@@ -58,7 +58,7 @@ def to_seq_info(inp: SequenceInput) -> SequenceInfo:
 def make_padded_input(x: Tensor, mask: Bool[Tensor, "b n"]) -> PaddedInput:
     """Build PaddedInput with sequential 1-D positions (shape b n 1)."""
     n = x.shape[1]
-    pos = torch.arange(n, device=x.device, dtype=x.dtype)
+    pos = torch.arange(n, device=x.device, dtype=torch.float32)
     positions = pos.unsqueeze(0).unsqueeze(-1).expand(x.shape[0], -1, -1)  # b n 1
     return PaddedInput(x=x, mask=mask, abs_positions=positions)
 
@@ -69,7 +69,7 @@ def make_packed_input(x: Tensor, cu_seqlens: Int[Tensor, "bp1"], max_seqlen: int
     Delegates to :func:`position_ids_from_packed`, which is ``torch.export``-compatible.
     """
     seq = PackedSequence(cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
-    positions = position_ids_from_packed(seq).to(x.dtype).unsqueeze(-1)  # nt 1
+    positions = position_ids_from_packed(seq).to(torch.float32).unsqueeze(-1)  # nt 1
     return PackedInput(x=x, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, abs_positions=positions)
 
 

{stackformers-3.8.0 → stackformers-3.8.2}/uv.lock

@@ -811,7 +811,7 @@ wheels = [
 
 [[package]]
 name = "stackformers"
-version = "3.7.0"
+version = "3.8.1"
 source = { editable = "." }
 dependencies = [
     { name = "beartype" },