more effective cond

2020-10-28 22:53:24 +09:00 · 2020-10-28 22:53:24 +09:00 · 09fc26999c
commit 09fc26999c
parent d26b3348df
1 changed files with 5 additions and 4 deletions
--- a/tensorflow/python/ops/numpy_ops/np_math_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py
@ -264,10 +264,11 @@ def cross(a, b, axisa=-1, axisb=-1, axisc=-1, axis=None):  # pylint: disable=mis

  def f(a, b):  # pylint: disable=missing-docstring
    # We can't assign to captured variable `axisa`, so make a new variable
-    axis_a = axisa
-    axis_b = axisb
-    axis_c = axisc
-    if axis is not None:
+    if axis is None:
+      axis_a = axisa
+      axis_b = axisb
+      axis_c = axisc
+    else:
      axis_a = axis
      axis_b = axis
      axis_c = axis