[TF2.0] Broadcasting

[awav]

Hello everyone,

I would like resuscitate very old issue. Actually, it is so old that even github's autocompletion doesn't offer it after typing "#" - #216. This request was raised several times, but still it hasn't been resolved.

In short, broadcasting interface is not "good enough" in TensorFlow :)

Lets first check how broadcasting works in numpy:

In [1]: import numpy as np
In [2]: a = np.random.rand(2, 3, 4)
In [3]: b = np.random.rand(4, 5)
In [4]: a @ b
Out[4]:
array([[[1.42709275, 1.40630067, 0.46525725, 0.68734581, 0.65227036],
        [2.01336504, 1.59980866, 0.93739699, 0.63190484, 0.92472892],
        [1.82979902, 1.46193243, 0.85498406, 0.5994646 , 0.77767957]],

       [[1.83010035, 1.49088728, 0.76694665, 0.65568003, 0.89110954],
        [2.12214864, 1.41728107, 1.04566743, 0.60652825, 0.97115822],
        [2.32478779, 2.06297214, 1.02016205, 0.81821249, 1.02604722]]])

Now, let's check what the TF is offering:

In [25]: a = tf.random.normal((2, 3, 4))
In [26]: b = tf.random.normal((4, 5))
In [27]: a @ b
... InvalidArgumentError: In[0] is not a matrix. Instead it has shape [2,3,4] [Op:MatMul] name: matmul/

Ouch! The "correct" way of doing it in the TF (of course there are other) is:

In [26]: a = tf.random.normal((2, 3, 4))
In [27]: b = tf.random.normal((4, 5))
In [28]: a @ tf.broadcast_to(b, tf.concat([a.shape[:-2], b.shape], axis=0))
<tf.Tensor: id=87, shape=(2, 3, 5), dtype=float32, numpy=
array([[[ 1.1977772 , -1.363074  ,  1.8021748 ,  0.1448586 , -0.6269997 ],
        [ 1.2322128 , -2.1586194 ,  0.09486479,  0.02937585, 0.9694344 ],
        [ 0.5580032 ,  6.11664   , -0.24535722,  0.16691092, -2.2263217 ]],
       [[-0.7386743 ,  1.2142425 ,  1.1371945 , -1.2736351 , -2.971829  ],
        [-1.9222848 , -0.7198772 , -0.9807504 ,  0.02805561, 1.0210879 ],
        [ 1.8334148 ,  0.80895233,  1.2308785 , -0.23910654, -1.5128168 ]]], dtype=float32)>

You can see how much effort it requires to make operations broadcastable for two distinct tensors: extract leading shape from the left tensor, extract shape from the right tensors, concatenate these shapes with correct axis, call tf.broadcast_to...

The same applies to cholesky, triangular solve and other operations. That is very upsetting that such a crucial feature isn't available out of the box.

Another concern is the performance of these "solutions". E.g. memory consumption for tiling and broadcast_to operations, as they simply copy the tensor to match leading dimensions. Of course, native TensorFlow broadcasting implementation would be preferable in this case.

Kind regards,
Artem Artemev

[awav]

@rmlarsen , @jvishnuvardhan ping

[alextp]

@awav thankfully when we do add support for this this it will not be an API removal; it just takes code which currently does not work and make it work.

So I don't think we need to block the tf2.0 release on fixing this, though it'd be very nice.

@rmlarsen now that we have a proper batch_matmul can we make it broadcast correctly on leading indices? I feel like this shouldn't be hard but I'm not familiar enough with eigen intricacies to make this happen.

[bloops]

This should be now fixed! (pending 3 week compatibility window)
47ab68d

[awav]

@bloops , That is fantastic news. Will broadcasting be fixed for all linear algebra operations or only for matmul?

[bloops]

Only for matmul. What else do you have in mind?
We are considering adding broadcasting support for matrix solve as well.

[awav]

@bloops, ideally, all linear algebra operations, as far as I know some of them already have broadcasting support. The triangular_solve is one important missing bit.

[awav]

Hello @bloops , any updates on triangular_solve?

[bloops]

Sorry, I am not working on it currently.

[awav]

@bloops , does anyone else is going to work on that? Do you know what should be changed in the code to make it work?

[awav]

@bloops , ping? :)

[awav]

One year has passed!!! Thanks to @srvasude.

[rmlarsen]

Closing since triangular solve now has broadcasting.