Tensorflow: ZeroDivisionError during concat of singel elements of tensor

Question

I am currently working on a "four wins" actor-critic agent in python. While trying to backpropagate through previously collected action probability distributions, I encountered the following error: ZeroDivisionError: integer division or modulo by zero

I was able to reproduce the error:

import tensorflow as tf

with tf.GradientTape() as tape:
    t = tf.Variable([1.])

    concat = tf.concat(values=[t[0], t[0]], axis=0)
    concat_sum = tf.reduce_sum(concat)

    grads = tape.gradient(concat_sum, t)

I do know the problem may sound trivial in this code example. Why there is an error here is still incomprehensible to me though! If one concatenates the first elements of the tensor and finally adds them should not that be the same as:

with tf.GradientTape() as tape:
    t = tf.Variable([1.])

    result = t + t

    grads = tape.gradient(result, t)

Why does one generate valid gradients while the other does not?

I am running Tensorflow version 2.7.0 on my CPU (Ubuntu 20.04.3 LTS)

Answer 1

This happens when you try to concatenate scalars, which is not supported. Tensorflow does not raise an error in eager mode, which is apparently a bug. The suggestion is to rather use tf.stack:

import tensorflow as tf

with tf.GradientTape() as tape:
    t = tf.Variable([1.])

    result = t + t

    grads = tape.gradient(result, t)

tf.print(grads)

with tf.GradientTape() as tape:
    t = tf.Variable([1.])

    stack = tf.stack(values=[t[0], t[0]], axis=0)
    concat_sum = tf.reduce_sum(stack)
    grads = tape.gradient(concat_sum, t)

tf.print(grads)

[2]
[2]