Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[jax2tf] Fix native lowering for modules that don't use some inputs. #13603

Merged
merged 1 commit into from Dec 12, 2022

Conversation

gnecula
Copy link
Collaborator

@gnecula gnecula commented Dec 11, 2022

We must drop the corresponding actual arguments when invoking the function and we must also be careful to compute the dim_args_specs based only on the kept inputs.

As a side benefit, we now allow dimension variables that occur in more complex polynomials, as long as they also occur as trivial monomials somewhere in the input shapes for the kept arguments.

@gnecula gnecula self-assigned this Dec 11, 2022
@gnecula gnecula added the pull ready Ready for copybara import and testing label Dec 11, 2022
@gnecula
Copy link
Collaborator Author

gnecula commented Dec 11, 2022

@maxwillzq @rxsang PTAL

@gnecula gnecula force-pushed the native_unused branch 2 times, most recently from 0571784 to 49d0346 Compare December 11, 2022 12:51
jax/experimental/jax2tf/README.md Outdated Show resolved Hide resolved
jax/experimental/jax2tf/README.md Outdated Show resolved Hide resolved
We must drop the corresponding actual arguments when invoking the function
and we must also be careful to compute the dim_args_specs based only
on the kept inputs.

As a side benefit, we now allow dimension variables that occur in
more complex polynomials, as long as they also occur as trivial
monomials somewhere in the input shapes for the kept arguments.
@gnecula
Copy link
Collaborator Author

gnecula commented Dec 12, 2022

@marcvanzee We are merging this to unblock another CL.

@copybara-service copybara-service bot merged commit 23001ae into google:main Dec 12, 2022
@marcvanzee
Copy link
Collaborator

@marcvanzee We are merging this to unblock another CL.

Sure, I approved it already in the first review so I was good to go from my end!

@gnecula gnecula deleted the native_unused branch December 16, 2022 15:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
pull ready Ready for copybara import and testing
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants