ponder-lab / Hybridize-Functions-Refactoring / #3245

87%

main: 88%

push

github

Hybridize a custom-gradient function above its tf.custom_gradient (#999)

`convertToHybrid` inserted `@tf.function` at the `def` line, so it
always went innermost. Beneath `@tf.custom_gradient`, a Keras layer
that calls the function can no longer infer its output shape, and the
layer's build fails with "We could not automatically infer the static
shape of the layer's output". Above it, the shape, the call and the
gradient all work.

The decorator now goes directly above the outermost
`@tf.custom_gradient`, so decorators above that one, such as
`@staticmethod`, stay outermost. `tf.recompute_grad` composes with
`tf.function` in either order, so it keeps the old placement.

Fixes #996.


Claude-Session: https://claude.ai/code/session_01NTnpjPjJj2aZy1ZkUD1eAn

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

10 of 12 new or added lines in 1 file covered. (83.33%)

4093 of 4724 relevant lines covered (86.64%)

0.87 hits per line