Numpy 1.25 breaks a few linalg functions

[fjosw]

With the upcoming release of numpy 1.25 the return types of the functions eig(), eigh(), qr(), slogdet(), and svd() change from tuple to a corresponding sublass of namedtuple (for example numpy.linalg.linalg.EighResult for eigh). autograd in its current form seems to be unable to deal with this. Here is a minimal example:

import autograd.numpy as anp
from autograd import jacobian

mat = anp.random.random((4, 4))

def func(arg):
    w, v = anp.linalg.eigh(arg)
    return w

print(jacobian(func)(mat))

which leads to the following error message:

Traceback (most recent call last):
  File "/home/fjosw/miniconda3/envs/np1.25/lib/python3.11/site-packages/autograd/tracer.py", line 118, in new_box
    return box_type_mappings[type(value)](value, trace, node)
           ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^
KeyError: <class 'numpy.linalg.linalg.EighResult'>

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/fjosw/tmp/test_eigh.py", line 12, in <module>
    print(jacobian(func)(mat))
          ^^^^^^^^^^^^^^^^^^^
  File "/home/fjosw/miniconda3/envs/np1.25/lib/python3.11/site-packages/autograd/wrap_util.py", line 20, in nary_f
    return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/fjosw/miniconda3/envs/np1.25/lib/python3.11/site-packages/autograd/differential_operators.py", line 60, in jacobian
    vjp, ans = _make_vjp(fun, x)
               ^^^^^^^^^^^^^^^^^
  File "/home/fjosw/miniconda3/envs/np1.25/lib/python3.11/site-packages/autograd/core.py", line 10, in make_vjp
    end_value, end_node =  trace(start_node, fun, x)
                           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/fjosw/miniconda3/envs/np1.25/lib/python3.11/site-packages/autograd/tracer.py", line 10, in trace
    end_box = fun(start_box)
              ^^^^^^^^^^^^^^
  File "/home/fjosw/miniconda3/envs/np1.25/lib/python3.11/site-packages/autograd/wrap_util.py", line 15, in unary_f
    return fun(*subargs, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/fjosw/tmp/test_eigh.py", line 9, in func
    w, v = anp.linalg.eigh(arg)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/fjosw/miniconda3/envs/np1.25/lib/python3.11/site-packages/autograd/tracer.py", line 46, in f_wrapped
    return new_box(ans, trace, node)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/fjosw/miniconda3/envs/np1.25/lib/python3.11/site-packages/autograd/tracer.py", line 120, in new_box
    raise TypeError("Can't differentiate w.r.t. type {}".format(type(value)))
TypeError: Can't differentiate w.r.t. type <class 'numpy.linalg.linalg.EighResult'>

Any idea how to fix this?

[asmeurer]

The issue is likely from this line in the traceback

    return box_type_mappings[type(value)](value, trace, node)

There is a mapping of types, but ideally the mapping should also accept subtypes. Perhaps the mapping object should be replaced with one that accepts subtypes. It should be possible to use functools.singledispatch for this, which properly handles subclasses.