Function for *Full* Forward Mode Jacobian

[twhughes]

Following Issue #439, I'm interested in implementing a version of the jacobian function from differential_operators.py but using forward mode differentiation.

From this comment from @j-towns, it seems to make this work one would need To get the full Jacobian of f you just need to write a loop to evaluate make_jvp(f)(x)(v) for each v in the standard basis of f's domain.

I'm trying to implement this below but running into some issues and would really appreciate help.

The reverse mode jacobian function defined in differential_operators.py is copied below:

@unary_to_nary
def jacobian_backward(fun, x):
    vjp, ans = _make_vjp(fun, x)
    ans_vspace = vspace(ans)
    jacobian_shape = ans_vspace.shape + vspace(x).shape
    grads = map(vjp, ans_vspace.standard_basis())
    return np.reshape(np.stack(grads), jacobian_shape)

I am trying to implement a forward mode version of this function using make_jvp instead of make_vjp, which is pasted below

@unary_to_nary
def jacobian_forward(fun, x):
    jvp, ans = _make_jvp(fun, x)
    ans_vspace = vspace(ans)
    jacobian_shape = ans_vspace.shape + vspace(x).shape
    grads = map(jvp, ans_vspace.standard_basis())
    return np.reshape(np.stack(grads), jacobian_shape)

However, this gives errors when run. Here is a very simple code sample to reproduce, using these two functions:

N = 2
A = np.random.random((N,N))
print('A = \n', A)

def fn(x):
   # partial of this fn w.r.t. `x` should be `A`
    return A @ x

x0 = np.random.random((N,))
print('Jac_rev = \n', jacobian_backward(fn)(x0))
print('Jac_for = \n', jacobian_forward(fn)(x0))

Which gives the following output

A =
 [[0.24700088 0.60585769]
 [0.02701674 0.71704424]]
Jac_rev =
 [[0.24700088 0.60585769]
 [0.02701674 0.71704424]]
Traceback (most recent call last):
  File "forward_test.py", line 67, in <module>
    print('Jac_for = \n', jacobian_forward(fn)(x0))
  File ".../anaconda3/lib/python3.7/site-packages/autograd/wrap_util.py", line 20, in nary_f
    return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)
  File "forward_test.py", line 34, in jacobian_forward
    jvp, ans = _make_jvp(fun, x)
TypeError: cannot unpack non-iterable function object

As you can see, the reverse mode works fine but the forward mode gives an error when trying to make the jvp.

Do you have any ideas on how to construct a forward version of the standard jacobian function?

Thanks for your help

[twhughes]

update: I managed to get something working but it isn't very robust, the code is below

@unary_to_nary
def jacobian_forward(fun, x):
    """ Compute jacobian of fun with respect to x using forward mode differentiation"""
    jvp = make_jvp(fun, x)
    ans = fun(x)
    grads = map(lambda b: jvp(b)[1], vspace(x).standard_basis())
    return np.reshape(np.stack(grads), jac_shape(x, ans)) # jac_shape just uses x and ans to determine the jacobian shape