Inefficient jacobian computation for embedding layers.

[mohamad-amin]

Hello,

When I try to compute the NTK of a model with an embedding layer, I get the following warning:

/usr/local/lib/python3.10/dist-packages/neural_tangents/_src/empirical.py:2215: UserWarning: No Jacobian rule found for gather.
  warnings.warn(f'No Jacobian rule found for {primitive}.')

And ntk computation fails, due to OOM errors. This is a reproduction:
https://colab.research.google.com/drive/1Z8ClXo85VjNEoKmWYHsS5dNccZ-Xf_JS?usp=sharing

[romanngg]

Thanks for pointing this out and the repro! Yes structured derivatives don't have structure annotation / jacobian implementation for scatter/gather primitives, and would be very inefficient currently (so I recommend using methods 1/2); will take a look and see if it can be improved.