Compute jacobian for tf.image.rot90 throws error "Encountered an exception while vectorizing the jacobian computation"

[MThalberg]

Click to expand!

Issue Type

Bug

Source

source

Tensorflow Version

tf-nightly

Custom Code

Yes

OS Platform and Distribution

Linux Ubuntu 20.04

Mobile device

No response

Python version

3.9

Bazel version

No response

GCC/Compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current Behaviour?

Compute jacobian for `tf.image.rot90` fails and throws error "Encountered an exception while vectorizing the jacobian computation". This only happens when I set jit_compile=True.

Standalone code to reproduce the issue

import tensorflow as tf
image = tf.random.uniform([2, 2, 1, 1], minval=0, maxval=1, dtype=tf.float64)
@tf.function(jit_compile=True)
def rot_func(image):
    return tf.image.rot90(image, k=3)
with tf.GradientTape(persistent=True) as tape:
    tape.watch(image)
    outputs = rot_func(image)
print(outputs.shape)
gradient = tape.jacobian(outputs, [image])
print(gradient)

Relevant log output

(2, 1, 2, 1)

......

ValueError: in user code: 
ValueError: Dimension 2 in both shapes must be equal, but are 1 and 2. Shapes are [4,2,1,2,1] and [?,2,2,1,1].
Encountered an exception while vectorizing the jacobian computation. Vectorization can be disabled by setting experimental_use_pfor to False.

[ganeshbagal072]

Can I work on this issue.

[timshen91]

Hi @ganeshbagal072 , I unassigned myself since I'm not working on it. Feel free to take it, but I don't know whether you need to get someone to assign it to you. I'm not familiar with the contribution workflow.