Complex result of real-valued function produces an incorrect reshape

[sjperkins]

On

• python 2.7.12
• autograd 1.2

the following:

from __future__ import print_function

import autograd.numpy as np
from autograd import jacobian

lightspeed = 299792458.
minus_two_pi_over_c = -2*np.pi/lightspeed

def complex_phase(lm, uvw, frequency):
    """ Compute complex phase """
    l = lm[:,None,None,0]
    m = lm[:,None,None,1]

    u = uvw[None,:,None,0]
    v = uvw[None,:,None,1]
    w = uvw[None,:,None,2]

    frequency = frequency[None,None,:]

    n = np.sqrt(1.0 - l**2 - m**2) - 1.0
    phase = minus_two_pi_over_c*(l*u + m*v + n*w)*frequency
    return np.exp(phase*1j)

jacob = jacobian(complex_phase, argnum=1)

lm = np.random.random((10,2))*.01
uvw = np.random.random((25,3))
frequency = np.random.random((16,))

print(uvw.dtype)

cp = complex_phase(lm, uvw, frequency)
print(cp.shape, cp.dtype)
print(jacob(lm, uvw, frequency).shape)

results in

$ python test_complex_phase.py 
float64
(10, 25, 16) complex128
Traceback (most recent call last):
  File "test_complex_phase.py", line 34, in <module>
    print(jacob(lm, uvw, frequency).shape)
  File "/home/sperkins/venv/autograd/local/lib/python2.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 "/home/sperkins/venv/autograd/local/lib/python2.7/site-packages/autograd/differential_operators.py", line 55, in jacobian
    return np.reshape(np.stack(grads), jacobian_shape)
  File "/home/sperkins/venv/autograd/local/lib/python2.7/site-packages/autograd/tracer.py", line 48, in f_wrapped
    return f_raw(*args, **kwargs)
  File "/home/sperkins/venv/autograd/local/lib/python2.7/site-packages/numpy/core/fromnumeric.py", line 257, in reshape
    return _wrapfunc(a, 'reshape', newshape, order=order)
  File "/home/sperkins/venv/autograd/local/lib/python2.7/site-packages/numpy/core/fromnumeric.py", line 52, in _wrapfunc
    return getattr(obj, method)(*args, **kwds)
ValueError: cannot reshape array of size 600000 into shape (10,25,16,25,3)

Unfortunately, passing the second argument as complex128 by calling jacob(lm, uvw.astype(np.complex128), frequency) as suggested by https://github.com/HIPS/autograd/issues/27 does not work around this.

[ecat]

Hi, I had a similar issue where I was doing physics simulations that involved complex numbers. I got around this by using the following, which seems to be working okay for me but you should double check for your own scenario.

from autograd.core import make_vjp as _make_vjp, make_jvp as _make_jvp
from autograd.extend import primitive, defvjp_argnum, vspace
from autograd.wrap_util import unary_to_nary
from autograd.builtins import tuple as atuple

@unary_to_nary
def jacobian_pkl(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())

    grads_out = np.stack(grads)
    if(np.prod(jacobian_shape) == np.prod(grads_out.shape)):
        return np.reshape(grads_out, jacobian_shape)
    else:       
        my_jacobian_shape = ans_vspace.shape + vspace(x).shape + (2,) # 2 to support real/im        
        re_im_grads = np.squeeze(np.reshape(grads_out, my_jacobian_shape))
        out = re_im_grads[..., 0] + 1j * re_im_grads[..., 1]         

        return out

Edit: I tested the above hack a bit more and it seems to not work in some cases that I did not carefully debug. Separating real/im below works fine.

In my case I also had some affine transformations with complex numbers. If that is what you have, you can modify your function to take in a vector with the real and imaginary components stacked, and do the complex math within the function - that way the jacobian never sees a complex number.