Optimize: Handle gradient wrt scalar inputs and guard against unsupported types #1784
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
ricardoV94
changed the title
ricardoV94
force-pushed
the
minimize_grad_scalar
branch
3 times, most recently
from
f26580f to
32dc544
Compare
ricardoV94
commented
Dec 10, 2025
| x: TensorVariable, | ||
| *args: Variable, | ||
| objective: Variable, | ||
| method: str = "BFGS", |
Member
Author
There was a problem hiding this comment.
I removed the default method from the Ops, since the helpers already have them. Helps avoiding mismatches
ricardoV94
force-pushed
the
minimize_grad_scalar
branch
from
32dc544 to
2a3adf0
Compare
ricardoV94
commented
Dec 10, 2025
| outputs: Sequence[Variable], | ||
| output_grads: Sequence[Variable], | ||
| ) -> list[Variable]: | ||
| def L_op(self, inputs, outputs, output_grads): |
Member
Author
There was a problem hiding this comment.
This was the only L_op with type-hints, I removed for consistency
ricardoV94
commented
Dec 10, 2025
Comment on lines
+166
to
+169
| grad_wrt_arg = dot(output_grad, arg_grad) | ||
| if isinstance(arg.type, ScalarType): | ||
| grad_wrt_arg = scalar_from_tensor(grad_wrt_arg) | ||
| grad_wrt_args.append(grad_wrt_arg) |
Member
Author
There was a problem hiding this comment.
An an LLM would have put it: "🥳 The actual bugfix"
ricardoV94
force-pushed
the
minimize_grad_scalar
branch
from
2a3adf0 to
5c29ce8
Compare
ricardoV94
force-pushed
the
minimize_grad_scalar
branch
from
5c29ce8 to
fbaa2fe
Compare
Member
Author
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.
Closes #1744
Closes #1743
MinimizeOps were very liberal with assuming inputs would all be tensors and so would gradients wrt to them.
This PR makes sure we only work with support types: Tensor and ScalarTypes. We could also try to support SparseVariables.
Right now the Op may still fail in the gradient pass if args are of other types (most I can think about are slices in index operations, and maybe rngs if they make sense). I added a test with a custom string type. A robust implementation of Minimize gradient would need to:
connection_patternmethod to tellgradabout this in advance (it complains otherwise).That can be left for a future PR
Gave up on mypy
Mypy didn't like my bugfix / more precise type-hints. Given the empty track record of it helping find bugs in our codebase and my dislike for it, I removed all typing.cast and added it to mypy-failing. Now we have both: more correct type-hints and more readable code. If someone wants to revert it be my guest, but I won't be helping