Add broadcasting support for `tf.where` #15982

yongtang · 2018-01-09T22:05:59Z

Adds where_v2 (which will be where in TF 2.0), which has numpy's broadcasting semantics.

This fix fixes #9284.

Signed-off-by: Yong Tang yong.tang.github@outlook.com

drpngx · 2018-01-09T22:37:40Z

ebrevdo · 2018-01-11T05:40:33Z

tensorflow/python/ops/array_ops.py

@@ -2515,13 +2539,24 @@ def where(condition, x=None, y=None, name=None):
  has the same shape as `x` and `y`, then it chooses which element to copy from
  `x` and `y`.

+  If `broadcast` is True, then values of `x`, `y` and `condition` are


It sounds like the new behavior is backwards compatible? Why hide it behind a flag if it doesn't break existing usage?

Thanks @ebrevdo. Initially I though the behaviors were different with respect to broadcast. Now think again it might be possible to not break the old behavior, while at the same time extend the broadcast 👍 . Let me take a look and update the PR.

yongtang · 2018-01-12T01:10:47Z

@ebrevdo The PR has been updated with broadcast attribute removed. Please take a look.

ebrevdo · 2018-02-06T16:50:19Z

tensorflow/python/kernel_tests/where_op_test.py

+    y = np.ones((7, 11))
+    np_val = np.where(f < 0, x, y)
+    with self.test_session(use_gpu=True):
+      tf_val = array_ops.where(constant_op.constant(f) < 0,


You should be able to use:

self.evaluate(f < 0, x, y)

same below.

ebrevdo · 2018-02-06T16:55:28Z

tensorflow/core/kernels/cwise_op_select.cc

+    BCast bcast(BCast::FromShape(cond->shape()),
+                BCast::FromShape(then->shape()));
+
+    if (bcast.IsValid()) {


add a comment about where this kicks in. we currently have scalar broadcasting and primary-dimension vector broadcasting and i'd like to know if this will kick in in one of those existing cases -- because it may affect performance.

yongtang · 2018-03-20T00:21:47Z

@ebrevdo The PR has been rebase to merge conflict.

However, after reviewing the code again, I realized that there is still one scenario where exiting tf.where is not compatible with np.where (below is the output before this PR):

$ python
Python 2.7.12 (default, Dec  4 2017, 14:50:18)
[GCC 5.4.0 20160609] on linux2
Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorflow as tf
>>> import numpy as np
>>>
>>> x = np.arange(4)
>>> y = np.zeros((4, 4))
>>> z = np.ones((4, 4))
>>>
>>> np.where(x > 1, y, z)
array([[1., 1., 0., 0.],
       [1., 1., 0., 0.],
       [1., 1., 0., 0.],
       [1., 1., 0., 0.]])
>>> v = tf.where(x > 1, y, z)
>>> tf.Session().run(v)
array([[1., 1., 1., 1.],
       [1., 1., 1., 1.],
       [0., 0., 0., 0.],
       [0., 0., 0., 0.]])
>>>

As is shown from the above example, when x has the shape (4) and y/z has the shape (4, 4), the broadcast orientation is different.

Because of that, the current PR will fail several test cases.

Unfortunately, I couldn't think of a way to make the proposed broadcasting changes in tf.where backward-compatible.

I am wondering maybe it would be better to name a new op (e.g., tf.where_v2) so that it is compatible with numpy and not breaking existing users?

carlthome · 2018-04-09T08:26:45Z

Status?

velicue · 2018-04-19T21:35:44Z

Any updates? Also as in the issue someone mentioned:

There is also not full support for broadcasting in the other way, when condition is smaller than x or y; as indicated in the docs it only works if condition is a vector, but it could be extended to work with condition matching an arbitrary number of first dimensions of x and y (I'm not sure if this is considered also "broadcasting") and/or broadcasting singleton dimensions.

Would this also be implemented?

ebrevdo · 2018-04-19T21:38:56Z

I think if you'd like the interface to exactly match np.where, it does make sense to create a new op + kernel, v2. you would want to expose it in tf.contrib somewhere (not in core).

shoyer · 2018-04-20T01:40:26Z

"If x and y are vectors of higher rank, then condition must be either a vector with size matching the first dimension of x, or must have the same shape as x."

Could we deprecate this behavior from tf.where (start issuing FutureWarning) and change it in some future API breaking release?

To ease the transition, we could safely add broadcasting support when the total number of dimensions match between all arguments, or add a function to contrib (temporarily) with the appropriate behavior.

Deviating from NumPy's broadcasting rules feels like a design mistake to me, and I suspect this will be a repeated source of confusion in the future.

yongtang · 2018-05-29T13:02:24Z

Sorry for the delay. The PR has been updated. Now a new op tf.contrib.framework.where is exposed so that the broadcast rule follows numpy conventions. The original tf.where remains intact. Please take a look.

yongtang · 2018-07-03T12:37:33Z

I rebased the PR to resolve the merge conflict though it looks like there are some build failures after that. Will take a look and update the PR shortly to fix the build.

charan223 · 2018-07-12T04:50:08Z

Status guys?

yongtang · 2018-07-12T20:44:02Z

The PR has been rebased with build error fixed. All test passed now. Sorry for the long wait.

rmlarsen · 2019-05-01T23:26:15Z

Looking now.

rmlarsen · 2019-05-01T23:28:22Z

tensorflow/core/kernels/cwise_op_select.cc

+    // 2-ary broadcast.
+
+    // Combine `then` and `else`.
+    BCast elem_bcast(BCast::FromShape(then->shape()),


maybe call this then_else_bcast?

Thanks @rmlarsen, the name has been changed.

rmlarsen · 2019-05-01T23:38:35Z

tensorflow/core/kernels/cwise_op_select.cc

@@ -324,10 +444,43 @@ struct BatchSelectFunctor<CPUDevice, T> {
  }
 };

+template <typename T, int NDIMS>
+struct BCastSelectFunctor<CPUDevice, T, NDIMS> {


Why is this defined in multiple places? Can you just define this once in the header file and template it on device type as well?

Thanks @rmlarsen, the PR has been updated with two definitions (one for CPU and one for SYCL) consolidated into one.

…pecification based on review comment. Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

yongtang · 2019-05-03T16:05:08Z

@rmlarsen Thanks for the review. The PR has been updated. Please take a look and let me know if there are any issues.

martinwicke · 2019-05-03T23:18:27Z

Since this has been going on for so long, I'd favor a cherry pick for this.

…

On Fri, May 3, 2019, 16:01 Alexandre Passos ***@***.***> wrote: ***@***.**** commented on this pull request. ------------------------------ In tensorflow/python/ops/array_ops.py <#15982 (comment)> : > @@ -3177,6 +3181,46 @@ def where(condition, x=None, y=None, name=None): raise ValueError("x and y must both be non-None or both be None.") ***@***.***_export("where", v1=["where_v2"]) +def where_v2(condition, x=None, y=None, name=None): Since I wrote this message the tf v2 API has been frozen for the 1.14 release. This means we'll need to export this symbol as where_v2 both in tf1 and in tf2 :-/ — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#15982 (comment)>, or mute the thread <https://github.com/notifications/unsubscribe-auth/AAEM57N3MKTMXHZBAVTHJQDPTS737ANCNFSM4ELCBNEQ> .

alextp · 2019-05-07T16:18:49Z

We had to roll this back because the tests had no adequate coverage and missed things such as the broadcasting selectv2 op having no gradient defined for it.

PiperOrigin-RevId: 247032095

yongtang · 2019-05-07T18:22:08Z

@alextp sorry about that. I will take a look to add grad and resubmit the PR later.

brianwa84 · 2019-05-07T19:38:13Z

I have a suggestion for the gradient, haven't tested it but maybe it gets you started.

@ops.RegisterGradient("SelectV2")
def _SelectGrad(op, grad):
  c = op.inputs[0]
  x = op.inputs[1]
  y = op.inputs[2]
  zeros = array_ops.zeros([], dtype=grad.dtype.base_dtype)
  gx = array_ops.where_v2(c, grad, zeros)
  gx_shape = array_ops.shape(gx)
  x_shape = array_ops.shape(x)
  rankdiff_x = array_ops.rank(gx) - array_ops.rank(x)
  # Reduce away broadcasted leading dims.
  gx = math_ops.reduce_sum(gx, axis=math_ops.range(rankdiff_x))
  # Reduce but keep x's 1-valued dims which were broadcast.
  gx = math_ops.reduce_sum(
      gx, keepdims=1, axis=array_ops.where(grad_shape[rankdiff_x:] > x_shape))

  gy = array_ops.where_v2(c, zeros, grad)
  gy_shape = array_ops.shape(gy)
  y_shape = array_ops.shape(y)
  rankdiff_y = array_ops.rank(gy) - array_ops.rank(y)
  # Reduce away broadcasted leading dims.
  gy = math_ops.reduce_sum(gy, axis=math_ops.range(rankdiff_y))
  # Reduce but keep y's 1-valued dims which were broadcast.
  gy = math_ops.reduce_sum(
      gy, keepdims=1, axis=array_ops.where(grad_shape[rankdiff_y:] > y_shape))

  return (None, gx, gy)

@brianwa84

Credit to @brianwa84: tensorflow#15982 (comment) Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

tejaslodaya · 2019-12-11T12:42:57Z

@yongtang any update on this? This was rolled back

brianwa84 · 2019-12-11T12:51:03Z

It's supported in tf.compat.v2.where

…

On Wed, Dec 11, 2019, 7:43 AM Tejas Lodaya ***@***.***> wrote: @yongtang <https://github.com/yongtang> any update on this? This was rolled back — You are receiving this because you commented. Reply to this email directly, view it on GitHub <#15982?email_source=notifications&email_token=AFJFSIY2TJ6NWWAGUSFZJDLQYDN5VA5CNFSM4ELCBNE2YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOEGS672A#issuecomment-564522984>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AFJFSI3P5IYRBTGAJK6GOSLQYDN5VANCNFSM4ELCBNEQ> .

yongtang · 2019-12-11T13:44:11Z

@tejaslodaya It was rolled back, then rolled forward (with the help from @brianwa84 for providing the gradient op 👍 ❤️ ). It is now available in 2.0.

tejaslodaya · 2019-12-12T05:47:57Z

I tried and it works! For anyone coming to this PR, here's how you do it

Before (TF 1.x)-

with tf.Session() as sess:  
    col = tf.convert_to_tensor([1,2,3,4,5,6,7,8,9,10,11,12])    
    print(tf.where(tf.math.greater(col, 10),
                  tf.zeros_like(col),
                  tf.ones_like(col)).eval())

[1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 0. 0.]

After (TF 2.x)-

import tensorflow as tf
col = [1,2,3,4,5,6,7,8,9,10,11,12]
print(tf.where(tf.math.greater(col, 10),
              tf.zeros([1]),
               tf.ones([1])))

tf.Tensor([1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 0. 0.], shape=(12,), dtype=float32)

Notice, I had to do zeros_like and broadcast it to the shape of the column to make it work in 1.x

Thanks @yongtang , great work!

googlebot added the cla: yes label Jan 9, 2018

yongtang mentioned this pull request Jan 9, 2018

Broadcasting support in tf.where #9284

Closed

zheng-xq assigned ebrevdo Jan 11, 2018

ebrevdo reviewed Jan 11, 2018

View reviewed changes

yongtang force-pushed the 9284-tf.where-broadcasting branch from 758804d to c06f5f3 Compare January 12, 2018 01:09

yongtang mentioned this pull request Jan 17, 2018

Feature request: extend tf.select to broadcast a scalar condition #3945

Closed

rmlarsen requested a review from aselle January 22, 2018 23:43

rmlarsen added the awaiting review Pull request awaiting review label Jan 22, 2018

ebrevdo reviewed Feb 6, 2018

View reviewed changes

yongtang force-pushed the 9284-tf.where-broadcasting branch from c06f5f3 to 2cceb0f Compare February 6, 2018 20:19

drpngx added the kokoro:force-run Tests on submitted change label Feb 11, 2018

kokoro-team removed the kokoro:force-run Tests on submitted change label Feb 11, 2018

yongtang force-pushed the 9284-tf.where-broadcasting branch from 2cceb0f to 07782e7 Compare March 20, 2018 00:14

yongtang force-pushed the 9284-tf.where-broadcasting branch 2 times, most recently from 67ddf92 to f97061b Compare May 29, 2018 00:03

yongtang force-pushed the 9284-tf.where-broadcasting branch from 73aa8da to 06037ac Compare July 2, 2018 18:33

yongtang force-pushed the 9284-tf.where-broadcasting branch 2 times, most recently from 5a80bd1 to c62c8ff Compare July 12, 2018 17:32

rmlarsen requested changes May 1, 2019

View reviewed changes

Rename elem_bcast to then_else_bcast, and remove duplicate template s…

e75409c

…pecification based on review comment. Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

rthadur requested a review from rmlarsen May 2, 2019 22:00

yongtang added the kokoro:force-run Tests on submitted change label May 3, 2019

kokoro-team removed the kokoro:force-run Tests on submitted change label May 3, 2019

Fix GPU build failure

33cd7b8

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

yongtang force-pushed the 9284-tf.where-broadcasting branch from 4d710d9 to 33cd7b8 Compare May 3, 2019 02:13

yongtang added the kokoro:force-run Tests on submitted change label May 3, 2019

kokoro-team removed the kokoro:force-run Tests on submitted change label May 3, 2019

rmlarsen approved these changes May 3, 2019

View reviewed changes

tensorflow-bot bot added kokoro:force-run Tests on submitted change ready to pull PR ready for merge process labels May 3, 2019

kokoro-team removed the kokoro:force-run Tests on submitted change label May 3, 2019

tensorflow-copybara merged commit 33cd7b8 into tensorflow:master May 7, 2019

PR Queue automation moved this from Reviewer Requested Changes to Merged May 7, 2019

yongtang deleted the 9284-tf.where-broadcasting branch May 7, 2019 05:02

tensorflow-copybara pushed a commit that referenced this pull request May 7, 2019

Automated rollback of commit 60524c4. Revert #15982.

14ee900

PiperOrigin-RevId: 247032095

seanpmorgan mentioned this pull request May 8, 2019

Nightly Tests have stopped or are lagging tensorflow/addons#237

Closed

yongtang added a commit to yongtang/tensorflow that referenced this pull request May 10, 2019

Add gradient of SelectV2

db27152

Credit to @brianwa84: tensorflow#15982 (comment) Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

yongtang mentioned this pull request May 10, 2019

Add broadcasting support for tf.where #28616

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add broadcasting support for `tf.where` #15982

Add broadcasting support for `tf.where` #15982

yongtang commented Jan 9, 2018 •

edited by martinwicke

drpngx commented Jan 9, 2018

ebrevdo Jan 11, 2018

yongtang Jan 11, 2018

yongtang commented Jan 12, 2018

ebrevdo Feb 6, 2018

ebrevdo Feb 6, 2018

yongtang commented Mar 20, 2018

carlthome commented Apr 9, 2018

velicue commented Apr 19, 2018

ebrevdo commented Apr 19, 2018

shoyer commented Apr 20, 2018 •

edited

yongtang commented May 29, 2018

yongtang commented Jul 3, 2018

charan223 commented Jul 12, 2018

yongtang commented Jul 12, 2018

rmlarsen commented May 1, 2019

rmlarsen May 1, 2019

yongtang May 3, 2019

rmlarsen May 1, 2019 •

edited

yongtang May 3, 2019

yongtang commented May 3, 2019

martinwicke commented May 3, 2019 via email

alextp commented May 7, 2019

yongtang commented May 7, 2019

brianwa84 commented May 7, 2019

tejaslodaya commented Dec 11, 2019

brianwa84 commented Dec 11, 2019 via email

yongtang commented Dec 11, 2019

tejaslodaya commented Dec 12, 2019

Add broadcasting support for tf.where #15982

Add broadcasting support for tf.where #15982

Conversation

yongtang commented Jan 9, 2018 • edited by martinwicke

drpngx commented Jan 9, 2018

ebrevdo Jan 11, 2018

Choose a reason for hiding this comment

yongtang Jan 11, 2018

Choose a reason for hiding this comment

yongtang commented Jan 12, 2018

ebrevdo Feb 6, 2018

Choose a reason for hiding this comment

ebrevdo Feb 6, 2018

Choose a reason for hiding this comment

yongtang commented Mar 20, 2018

carlthome commented Apr 9, 2018

velicue commented Apr 19, 2018

ebrevdo commented Apr 19, 2018

shoyer commented Apr 20, 2018 • edited

yongtang commented May 29, 2018

yongtang commented Jul 3, 2018

charan223 commented Jul 12, 2018

yongtang commented Jul 12, 2018

rmlarsen commented May 1, 2019

rmlarsen May 1, 2019

Choose a reason for hiding this comment

yongtang May 3, 2019

Choose a reason for hiding this comment

rmlarsen May 1, 2019 • edited

Choose a reason for hiding this comment

yongtang May 3, 2019

Choose a reason for hiding this comment

yongtang commented May 3, 2019

martinwicke commented May 3, 2019 via email

alextp commented May 7, 2019

yongtang commented May 7, 2019

brianwa84 commented May 7, 2019

tejaslodaya commented Dec 11, 2019

brianwa84 commented Dec 11, 2019 via email

yongtang commented Dec 11, 2019

tejaslodaya commented Dec 12, 2019

Add broadcasting support for `tf.where` #15982

Add broadcasting support for `tf.where` #15982

yongtang commented Jan 9, 2018 •

edited by martinwicke

shoyer commented Apr 20, 2018 •

edited

rmlarsen May 1, 2019 •

edited