Skip to content

Fix: Replace tl.range flatten with Python range for Triton 3.x compatibility#19

Open
Jah-yee wants to merge 4 commits into
thinking-machines-lab:mainfrom
Jah-yee:pr-flatten
Open

Fix: Replace tl.range flatten with Python range for Triton 3.x compatibility#19
Jah-yee wants to merge 4 commits into
thinking-machines-lab:mainfrom
Jah-yee:pr-flatten

Conversation

@Jah-yee
Copy link
Copy Markdown

@Jah-yee Jah-yee commented Apr 15, 2026

Good day,

This PR fixes Issue #14: Incompatible with all public PyTorch/Triton releases.

Problem

The library cannot be used with any publicly available PyTorch/Triton combination due to the parameter in tl.range() being removed before the Triton 3.x public release.

Root Cause

The code used tl.range(..., flatten=True) which was available only in internal/nightly Triton builds but never shipped in any public release.

Solution

Replace tl.range(..., flatten=True) with Python's built-in range():

  • Python's range() works inside @triton.jit decorated functions
  • It's compatible with all Triton versions (2.x and 3.x)
  • The loop logic remains the same

Testing

  • Code compiles successfully with Triton 3.2.0

Related

Fixes: TypeError: range.init() got an unexpected keyword argument 'flatten'

Thank you for your work on this excellent library!

Warmly, RoomWithOutRoof

RoomWithOutRoof and others added 4 commits April 15, 2026 22:49
- Add comprehensive tests in tests/test_batch_invariant_ops.py
- Fix mean_dim to handle empty dimension (return zeros instead of NaN)

Good day,

Thank you for your work on this excellent library!

Warmly,
RoomWithOutRoof
Added softmax along with log_softmax, expanding batch-invariant operations.

Good day!

Warmly,
RoomWithoutRoof
Fixes Issue thinking-machines-lab#15: Fix NaN in mean_kernel for empty dimension

When the reduction dimension is empty (N==0), return 0 to avoid NaN.

Good day!

Warmly, RoomWithoutRoof
…ibility

The flatten parameter was removed in Triton 3.x public release.
This changes tl.range(..., flatten=True) to Python's built-in range(),
which works inside @triton.jit functions and is compatible with all versions.

Fixes: TypeError: range.__init__() got an unexpected keyword argument 'flatten'
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant