jacobians and hessians for an application involving complex numbers

The imaginary components of the derivatives of the following complex-valued functions are always zero in jax. In contrast, tensorflow returns a non zero imaginary component. I think this is a bug on the JAX side of things.

from jax.config import config; config.update("jax_enable_x64", True)
import jax
import jax.numpy as np
import numpy as onp

import tensorflow as tf

zs = 0.5j * np.arange(5) + np.arange(5)

print("input", zs)


def fn(z):
    return np.cos(np.linalg.norm(z*2))

grad = jax.jacfwd(fn)
print("jax", fn(zs), grad(zs))

def tf_fn(z):
    return tf.cos(tf.norm(z*2))

tf_zs = tf.convert_to_tensor(0.5j * onp.arange(5) + onp.arange(5))
tf_res = tf_fn(tf_zs)

sess = tf.Session()


grad_ys = tf.ones_like(tf_res)
grad_op = tf.gradients(tf_res, tf_zs, grad_ys=grad_ys)
print("tf", sess.run([tf_res, grad_op, grad_ys]))

input [0.+0.j  1.+0.5j 2.+1.j  3.+1.5j 4.+2.j ]
jax 0.9495740004388323 [0.        +0.j 0.10240272+0.j 0.20480544+0.j 0.30720815+0.j
 0.40961087+0.j]
tf [(0.9495740004388323+0j), [array([0.        +0.j        , 0.10240272+0.05120136j,
       0.20480544+0.10240272j, 0.30720815+0.15360408j,
       0.40961087+0.20480544j])], (1+0j)]

[proteneer]

Again, thanks for all the help on this issue!

[proteneer]

Confirmed bug in tensorflow:

import numpy as onp
import autograd as ag
import autograd.numpy as anp
import numpy as onp
import tensorflow as tf

inp = anp.array(2.0)

print("input", inp)

def ag_fn(x):
    real = anp.cos(x+2)
    imag = anp.sin(x-1)
    return anp.abs(real+1j*imag)

ag_hess = ag.hessian(ag_fn)

print("ag val:", ag_fn(inp))
print("ag hess:", ag_hess(inp))

def tf_fn(x):
    real = tf.cos(x+2)
    imag = tf.sin(x-1)
    return tf.abs(tf.complex(real, imag))

# tf_inp = tf.convert_to_tensor(inp)
tf_inp = tf.placeholder(shape=tuple(), dtype=onp.float64)

out_op = tf_fn(tf_inp)

tf_grad = tf.gradients(out_op, tf_inp)[0]
tf_hess = tf.hessians(out_op, tf_inp)[0]

sess = tf.Session()
delta = 1e-7

_, d0, tf_ad = sess.run([out_op, tf_grad, tf_hess], feed_dict={tf_inp: inp})
_, d1, _ = sess.run([out_op, tf_grad, tf_hess], feed_dict={tf_inp: inp+delta})

print("tf_numerical derivative:", (d1-d0)/delta)
print("tf_autodiff derivative:", tf_ad)

input 2.0
ag val: 1.0655155566059393
ag hess: -0.25533014019223726
2019-04-14 22:55:43.481283: I tensorflow/core/platform/cpu_feature_guard.cc:141] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
tf_numerical derivative: -0.25533013481293665
tf_autodiff derivative: -1.0655155566059389