Use the supplied gradient in custom_gradient and add a test suite - #47
Open
MarcosAsh wants to merge 3 commits into
Open
Use the supplied gradient in custom_gradient and add a test suite#47MarcosAsh wants to merge 3 commits into
MarcosAsh wants to merge 3 commits into
Conversation
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.
Description
custom_gradientthrew the supplied gradient away. It called the decorated function, discarded the second return value and only warned, sokeras.ops.custom_gradientreturned autodiff of the forward pass rather than the gradient the caller wrote. This wires the gradient intomx.custom_function's vjp. One wrinkle worth knowing: mlx hands the vjp a bare array for a single primal and a tuple for several, so the primals need normalising before they go back to the decorated function.The two in tree users were both getting the wrong answer.
fake_quant_with_min_max_varsnow gets its straight through estimator, gradient 1 inside the range and 0 outside, where before it differentiated the rounding.einsum_dense's int8 path is the other one.Nothing upstream covers this.
keras/src/ops/core_test.py::test_custom_gradientis gated on the backend being tensorflow, jax or torch, and the two other tests that reachcustom_gradientare gated the same way, so this adds the first tests in this repo along with a CI step to run them.keras/src/ops/core_test.py,keras/src/quantizers/andkeras/src/layers/core/einsum_dense_test.pygive 430 passed and 0 failed.