Using make_vjp on functions with multiple return values seems to produce zeroed VJPs

[Ian Williamson]

I am attempting to use autograd.make_vjp on a function that returns multiple values and the VJP function seems to always produced zeroed arrays.

Below is a minimum example:

import numpy as np
import autograd.numpy as npa

A = np.array([
    [1.0, 2.0, 3.0],
    [4.0, 5.0, 6.0],
    [7.0, 8.0, 9.0],
])
B = np.array([
    [10.0, 11.0, 12.0],
    [13.0, 14.0, 15.0],
    [16.0, 17.0, 18.0],
])

def fun(x):
  return npa.matmul(A, x), npa.matmul(B, x)

x = np.array([-1.0, +1.5, 2.25])

vjp_fun, ans = autograd.make_vjp(fun)(x)
vjp_fun(autograd.core.vspace(ans).ones())
# array([0., 0., 0.])

If I modify fun to return npa.matmul(A, x) the end result becomes non-zero.

[Ian Williamson]

For anyone reading this... the issue is that the built-in autograd containers must be used. In this case autograd.builtins.tuple.