[pallas] Fix Pallas Mosaic 64-bit mode loop indexing#38154
Merged
Conversation
copybara-service
Bot
force-pushed
the
test_925836768
branch
6 times, most recently
from
d7aea04 to
8c9752d
Compare
This is another attempt for #37970: Prior to this change, when running in 64-bit mode, we get the error `'arith.cmpi' op requires all operands to have the same type` because in the pipeline loop some of the operations get promoted to 64-bit mode. That change was rolled back due to failure attempting to apply `asarray` to tracers. I do not have a reproducing failure from the user, but an LLM suggested that this can happen in presence of disable_jit, and in fact created a failing test case for a pipeline with disable_jit. We had used `jnp.int32(0)` and `jnp.int32(num_steps)` which create tracers instead of constants, which in turn confuses the disable_jit machinery. A better approach is to use `np.int32(0)`, which creates constants of the right type. For `num_steps` we should use `np.int32(num_steps) if isinstance(num_steps, int) else num_steps`. While fixing this I ran into a lowering error "error: 'arith.trunci' op operand type 'i32' and result type 'i32' are cast incompatible" which I worked around by not emitting a truncation when the types are the same. PiperOrigin-RevId: 926542136
copybara-service
Bot
force-pushed
the
test_925836768
branch
from
8c9752d to
9b82b11
Compare
copybara-service Bot
pushed a commit
that referenced
this pull request
Jun 4, 2026
That change was #38154. The problem is that now constants that were 0 or 1 were turned into np.int32 and there are several `is_instance(v, int)` checks throughout Pallas that behave differently than before. I need to rethink how to make 64-bit more work for pipelines. Reverts 9b82b11 PiperOrigin-RevId: 926690252
copybara-service Bot
pushed a commit
that referenced
this pull request
Jun 4, 2026
That change was #38154. The problem is that now constants that were 0 or 1 were turned into np.int32 and there are several `is_instance(v, int)` checks throughout Pallas that behave differently than before. I need to rethink how to make 64-bit more work for pipelines. Reverts 9b82b11 PiperOrigin-RevId: 926690252
copybara-service Bot
pushed a commit
that referenced
this pull request
Jun 4, 2026
That change was #38154. The problem is that now constants that were 0 or 1 were turned into np.int32 and there are several `is_instance(v, int)` checks throughout Pallas that behave differently than before. I need to rethink how to make 64-bit more work for pipelines. Reverts 9b82b11 PiperOrigin-RevId: 926690252
copybara-service Bot
pushed a commit
that referenced
this pull request
Jun 4, 2026
That change was #38154. The problem is that now constants that were 0 or 1 were turned into np.int32 and there are several `is_instance(v, int)` checks throughout Pallas that behave differently than before. I need to rethink how to make 64-bit more work for pipelines. Reverts 9b82b11 PiperOrigin-RevId: 926751271
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
[pallas] Fix Pallas Mosaic 64-bit mode loop indexing
This is another attempt for #37970:
Prior to this change, when running in 64-bit mode, we get
the error
'arith.cmpi' op requires all operands to have the same typebecause in the pipeline loop some of the operations get promoted to 64-bit mode.That change was rolled back due to failure attempting to apply
asarrayto tracers. I do not have a reproducing failure from the user, but an LLM
suggested that this can happen in presence of disable_jit, and in fact
created a failing test case for a pipeline with disable_jit.
We had used
jnp.int32(0)andjnp.int32(num_steps)which createtracers instead of constants, which in turn confuses the disable_jit
machinery. A better approach is to use
np.int32(0),which creates constants of the right type. For
num_stepswe shoulduse
np.int32(num_steps) if isinstance(num_steps, int) else num_steps.While fixing this I ran into a lowering error "error: 'arith.trunci' op operand type 'i32' and result type 'i32' are cast incompatible" which I worked around by not emitting a truncation when the types are the same.