Skip to content

[Pallas] remove jax_export_ignore_forward_compatibility setting#2452

Open
cota wants to merge 2 commits into
pytorch:mainfrom
cota:jax_compat
Open

[Pallas] remove jax_export_ignore_forward_compatibility setting#2452
cota wants to merge 2 commits into
pytorch:mainfrom
cota:jax_compat

Conversation

@cota
Copy link
Copy Markdown
Collaborator

@cota cota commented May 15, 2026

TorchTPU does this for us since 759cfba8e5.

TorchTPU does this for us since 759cfba8e5.
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label May 15, 2026
@cota cota requested a review from AmesingFlank May 15, 2026 21:20
Copy link
Copy Markdown
Contributor

@AmesingFlank AmesingFlank left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One sec, let me confirm that on the latest version, this doesn't regress perf

Copy link
Copy Markdown
Contributor

@AmesingFlank AmesingFlank left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

confirmed, not seeing any attention perf regression. lgtm

@cota
Copy link
Copy Markdown
Collaborator Author

cota commented May 15, 2026

Thanks for testing. I tested it on my side as well when landing 759cfba8e5 in TorchTPU -- the dumped LLO is indeed different without that commit (and without the setting in Helion).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants