Tfjs: Overriding ReLU gradients using gradient registry

[ludwigschubert]

TensorFlow.js version

1.5.0

Describe the problem or feature request

I'm attempting to override the gradient computation of certain Ops, say ReLUs, in a loaded GraphModel. The gradient registry seems like the right place, but currently tf.getKernelsForBackend doesn't seem to return a ReLU kernel. It seems like most core Ops currently don't yet use the gradient registry.

Does my description here seem correct? Is there currently already a different way to override Op gradients of loaded models? If not, could and should I try porting the specific ops I'd like to be able to override to the new kernel and gradient registry system?

Thank you for all your work on TFJS!

[rthadur]

@ludwigschubert thank you , can you please provide your use case.

[ludwigschubert]

Hm, my use case? As in, why do I want to register custom gradients? LMK if I misunderstand your question!

I'm looking to implement various techniques from the explainability and interpretability literature, which often use TF's Graph's gradient_override_map (example 1, example 2).

[tafsiri]

Hi Ludwig, you are correct, not much is currently using the gradient registry (or even kernel registry). We've been doing some of the groundwork for using a registry approach for kernels, gradients etc. We are going to be doing a bunch more work in this direction this quarter (so there will be quite a few changes there soon).

All new ops will use the registry approach and old ones will be ported to use it. So PRs that do that will be welcome, but there will be some more moving parts there. Is it a lot of ops you need?

[ludwigschubert]

Exciting to hear! I mainly need ReLU and MaxPool atm. I'm experimenting on a fork right now to see if I can port those two ops to use the registry. I'm not sure I'll know enough about the inner workings of TFJS to make it PR-ready, but I'll try. If there's a commit/PR/branch that ports an existing op, I'd be glad about any pointers. (But no pressure, I'll explore on my own, too. :D)

Thanks for giving me extra context! :-)

[rthadur]

@ludwigschubert can we close this issue ?