Stuck on an issue?

Lightrun Answers was designed to reduce the constant googling that comes with debugging 3rd party libraries. It collects links to all the places you might be looking at while hunting down a tough bug.

And, if you’re still stuck at the end, we’re happy to hop on a call to see how we can help out.

autodiff support for jax.numpy.linalg.eig

See original GitHub issue

Note that eigh is already taken care of.

Issue Analytics

State:
Created 3 years ago
Comments:48 (27 by maintainers)

Top GitHub Comments

2reactions

Randlcommented, Dec 13, 2020

There is a good post on gauge problem: https://re-ra.xyz/Gauge-Problem-in-Automatic-Differentiation/ Also related discussion in tensorflow https://github.com/tensorflow/tensorflow/pull/33808

1reaction

j-townscommented, May 17, 2021

Please make sure that you’re using a recent version of JAX. When I run

import jax
from jax import grad

def matrixEigs(a1):
    arr = jax.numpy.array([[a1, 1], [1, 2]])
    eigs = jax.numpy.linalg.eigvals(arr)
    return eigs

eigGrad = grad(matrixEigs)

gradients = eigGrad(1.)

I get

TypeError: Gradient only defined for scalar-output functions. Output had shape: (2,).

Top Results From Across the Web

numpy.linalg.eig() - JAX documentation

jax.numpy.linalg.eig# ... Compute the eigenvalues and right eigenvectors of a square array. LAX-backend implementation of numpy.linalg.eig() . This differs from ...

jax.lax.linalg.eig - JAX documentation

jax.lax.linalg.eig# ... Eigendecomposition of a general matrix. Nonsymmetric eigendecomposition is at present only implemented on CPU. Parameters. x ( Union [ ...

numpy.linalg.eigh() - JAX documentation - Read the Docs

(conjugate symmetric) or a real symmetric matrix. Returns two objects, a 1-D array containing the eigenvalues of a , and a 2-D square...

numpy.linalg.eigvals() - JAX documentation - Read the Docs

LAX-backend implementation of numpy.linalg.eigvals() . Original docstring below. Main difference between eigvals and eig : the eigenvectors aren't returned.

numpy.linalg.eigvalsh() - JAX documentation

jax.numpy.linalg.eigvalsh# · a ((..., M, M) array_like) – A complex- or real-valued matrix whose eigenvalues are to be computed. · UPLO ({'L', 'U'},...