TF Gradient Tape имеет проблемы с перекрестными продуктами?

Question

Я пытаюсь использовать градиентную ленту TF как инструмент автограда для root поиска по методу Ньютона. Но когда я пытаюсь вычислить матрицу Якоби, кажется, что tf.GradientTape.jacobian не может обрабатывать перекрестные продукты:

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))

выдает ошибку ниже:

StagingError : в преобразованном коде: относительно /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'

Принимая во внимание, что если я переключу вызов Jacobian на простой градиент:

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))

дает ожидаемый результат:

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

Это ошибка ?? Или я что-то не так делаю с методом tape.jacobian?

ps python версия 3.7.4; tf версия 2.0.0 Все установлено с conda.