TF Gradient Tape has issues with cross products?

Question

I'm trying to use TF gradient tape as an autograd tool for root finding via Newton's method. But when I'm trying to compute the Jacobian matrix, it seems that tf.GradientTape.jacobian can't handle cross products:

x = tf.convert_to_tensor(np.array([1., 2., 3.]))
Wx = np.ones((3))
with tf.GradientTape() as tape:
    tape.watch(x)
    y = tf.linalg.cross(x, Wx)
print(tape.jacobian(y, x))

gives below error:

StagingError: in converted code: relative to /Users/xinzhang/anaconda3/lib/python3.7/site-packages:

tensorflow_core/python/ops/parallel_for/control_flow_ops.py:184 f  *
    return _pfor_impl(loop_fn, iters, parallel_iterations=parallel_iterations)
tensorflow_core/python/ops/parallel_for/control_flow_ops.py:257 _pfor_impl
    outputs.append(converter.convert(loop_fn_output))
tensorflow_core/python/ops/parallel_for/pfor.py:1231 convert
    output = self._convert_helper(y)
tensorflow_core/python/ops/parallel_for/pfor.py:1395 _convert_helper
    if flags.FLAGS.op_conversion_fallback_to_while_loop:
tensorflow_core/python/platform/flags.py:84 __getattr__
    wrapped(_sys.argv)
absl/flags/_flagvalues.py:633 __call__
    name, value, suggestions=suggestions)

UnrecognizedFlagError: Unknown command line flag 'f'

Whereas if I switch out the call to jacobian to a simple gradient:

x = tf.convert_to_tensor(np.array([1., 2., 3.]))
Wx = np.ones((3))
with tf.GradientTape() as tape:
    tape.watch(x)
    y = tf.linalg.cross(x, Wx)
print(tape.gradient(y, x))

gives the expected result:

tf.Tensor([0. 0. 0.], shape=(3,), dtype=float64)

Is this a bug?? Or am I doing something wrong with the tape.jacobian method?

ps python version 3.7.4; tf version 2.0.0 Everything installed with conda.

Answer 1

This might be a bug in Tensorflow Version 2.0 but it is fixed in Tensorflow Version 2.1 .

So, please upgrade your Tensorflow Version to either 2.1 or 2.2 and the issue will be resolved.

Working code is mentioned below:

!pip install tensorflow==2.2

import tensorflow as tf
import numpy as np

print(tf.__version__)

x = tf.convert_to_tensor(np.array([1., 2., 3.]))
Wx = np.ones((3))
with tf.GradientTape() as tape:
    tape.watch(x)
    y = tf.linalg.cross(x, Wx)
print(tape.jacobian(y, x))

Output is shown below:

2.2.0

tf.Tensor(
[[ 0.  1. -1.]
 [-1.  0.  1.]
 [ 1. -1.  0.]], shape=(3, 3), dtype=float64)