tf.custom_gradient silently fails when using all keyword arguments about tensorflow

[tilakrayal]

Triage Notes:
I was able to reproduce the issue on tensorflow v2.15, v2.16 and tf-nightly. Kindly find the gist of it here.

The tf.custom_gradients are working as expected when z = bar(x, y) & z = bar(x=x, y=y) and it was failing with z = bar(x, y=y) by throwing the error.

[ctargon]

@tilakrayal when running your gist, the gradient is NOT as expected when z = bar(x=x, y=y). the expected output should be the custom gradient (which is [[-2000], [-4000], [-6000]] for dz/dx), given the function is wrapped in tf.custom_gradient, but instead returns the gradient produced by autograd (which is [[-0.66], [-1.33], [-2.]]).

The issue is that the custom gradient is silently ignored when using keyword arguments for every parameter.

[cantonios]

The decorator @tf.custom_gradient uses the input to deduce which arguments are differentiable inputs, and which arguments are non-differentiable parameters to the function. Positional arguments are considered differentiable, keyword arguments non-differentiable.

When passing all keyword arguments, the custom gradient is silently ignored and defaults to the autograd.

It thinks your function has no differentiable inputs, so simply forwards the upstream gradient and continues the computation.

When passing a mix of positional and keyword arguments, an exception is thrown.

Because you're returning multiple gradients when only the one non-keyword input is considered a differentiable input.

---

As for what to do: we will not change the behavior of the custom_gradient decorator. You can update the documentation with a PR if you like to reflect this behavior.

[github-actions]

This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No