-
Notifications
You must be signed in to change notification settings - Fork 25.7k
Description
🚀 The feature, motivation and pitch
I have a neural network that outputs (among other things) a 3-dimensional batch of matrices, representing the variance of a Gaussian distribution. I am enforcing that these matrices be symmetric positive-definite by applying a batched matrix exponent. As part of this calculation, I am performing a batch calculation of the eigenvalues and eigenvectors of the matrices, using functorch's excellent vmap operator. This triggers an internal warning asking me to file an issue to request a batching rule, so I'm doing that.
Alternatives
Continue with the current implementation. That will work for me; this is not a performance bottleneck. I'm only submitting the issue because the warning asked me to.
Additional context
Here is the warning message I received:
/content/utils.py:9: UserWarning: There is a performance drop because we have not yet implemented the batching rule for aten::linalg_eigh. Please file us an issue on GitHub so that we can prioritize its implementation. (Triggered internally at ../aten/src/ATen/functorch/BatchedFallback.cpp:82.)
q, eigvals = la.eigh(x)
cc @jianyuh @nikitaved @pearu @mruberry @walterddr @IvanYashchuk @xwang233 @lezcano @zou3519 @Chillee @samdow @kshitij12345 @janeyx99