Allow tf.GradientTape.batch_jacobian to accept containers of tensors as source

[cammyoung]

Click to expand!

Issue Type

Feature Request

Source

binary

Tensorflow Version

tf 2.8.2

Custom Code

No

OS Platform and Distribution

No response

Mobile device

No response

Python version

No response

Bazel version

No response

GCC/Compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current Behaviour?

At the moment, containers of tensors are supported by `tf.GradientTape.jacobian`, but not `tf.GradientTape.batch_jacobian`.

If this behavior was allowed in `batch_jacobian`, I would expect it to return the same container with tensors equal to the `batch_jacobian` of `target` with respect to each tensor in the `source` container.

Example expected output shown in the colab.

Standalone code to reproduce the issue

https://colab.research.google.com/drive/1RTv34VMa-joB45x7xu-9H3pdSTnDAj72?usp=sharing

Relevant log output

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-7-5c1f96355dc5> in <module>()
      6     pred = my_model(x, training=False)[:, :, 0]
      7 
----> 8 j = tape.batch_jacobian(pred, x)

/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/backprop.py in batch_jacobian(self, target, source, unconnected_gradients, parallel_iterations, experimental_use_pfor)
   1287       dim = target_shape.dims[0]
   1288     if not (target_shape.with_rank_at_least(2) and
-> 1289             source.shape.with_rank_at_least(2) and
   1290             dim.is_compatible_with(source.shape[0])):
   1291       raise ValueError(

AttributeError: 'dict' object has no attribute 'shape'