KEMBAR78
Requesting Functorch Batching Rule for aten::linalg_eigh · Issue #108481 · pytorch/pytorch · GitHub
Skip to content

Requesting Functorch Batching Rule for aten::linalg_eigh #108481

@Mark-Lowell

Description

@Mark-Lowell

🚀 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

Metadata

Metadata

Assignees

No one assigned

    Labels

    actionablefeatureA request for a proper, new feature.module: functorchPertaining to torch.func or pytorch/functorchmodule: linear algebraIssues related to specialized linear algebra operations in PyTorch; includes matrix multiply matmultriagedThis issue has been looked at a team member, and triaged and prioritized into an appropriate module

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions