chore: import upstream snapshot with attribution

tensorflow/python/ops/accumulate_n_benchmark.py

@@ -0,0 +1,145 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Benchmark for accumulate_n() in math_ops."""
+
+import random
+import time
+
+
+from tensorflow.python.client import session
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import data_flow_ops
+from tensorflow.python.ops import gen_control_flow_ops
+from tensorflow.python.ops import gen_state_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.platform import test
+
+
+class AccumulateNBenchmark(test.Benchmark):
+
+  def _AccumulateNTemplate(self, inputs, init, shape, validate_shape):
+    var = gen_state_ops.temporary_variable(
+        shape=shape, dtype=inputs[0].dtype.base_dtype)
+    ref = state_ops.assign(var, init, validate_shape=validate_shape)
+    update_ops = [
+        state_ops.assign_add(
+            ref, tensor, use_locking=True).op for tensor in inputs
+    ]
+    with ops.control_dependencies(update_ops):
+      return gen_state_ops.destroy_temporary_variable(ref, var_name=var.op.name)
+
+  def _AccumulateNInitializedWithFirst(self, inputs):
+    return self._AccumulateNTemplate(
+        inputs,
+        init=array_ops.zeros_like(inputs[0]),
+        shape=inputs[0].get_shape(),
+        validate_shape=True)
+
+  def _AccumulateNInitializedWithMerge(self, inputs):
+    return self._AccumulateNTemplate(
+        inputs,
+        init=array_ops.zeros_like(gen_control_flow_ops.merge(inputs)[0]),
+        shape=tensor_shape.TensorShape([0]),
+        validate_shape=False)
+
+  def _AccumulateNInitializedWithShape(self, inputs):
+    return self._AccumulateNTemplate(
+        inputs,
+        init=array_ops.zeros(
+            shape=inputs[0].get_shape(), dtype=inputs[0].dtype.base_dtype),
+        shape=inputs[0].get_shape(),
+        validate_shape=True)
+
+  def _GenerateUnorderedInputs(self, size, n):
+    inputs = [random_ops.random_uniform(shape=[size]) for _ in range(n)]
+    random.shuffle(inputs)
+    return inputs
+
+  def _GenerateReplicatedInputs(self, size, n):
+    return n * self._GenerateUnorderedInputs(size, 1)
+
+  def _GenerateOrderedInputs(self, size, n):
+    inputs = self._GenerateUnorderedInputs(size, 1)
+    queue = data_flow_ops.FIFOQueue(
+        capacity=1, dtypes=[inputs[0].dtype], shapes=[inputs[0].get_shape()])
+    for _ in range(n - 1):
+      op = queue.enqueue(inputs[-1])
+      with ops.control_dependencies([op]):
+        inputs.append(math_ops.tanh(1.0 + queue.dequeue()))
+    return inputs
+
+  def _GenerateReversedInputs(self, size, n):
+    inputs = self._GenerateOrderedInputs(size, n)
+    inputs.reverse()
+    return inputs
+
+  def _SetupAndRunBenchmark(self, graph, inputs, repeats, format_args):
+    with graph.as_default():
+      add_n = math_ops.add_n(inputs)
+      acc_n_first = self._AccumulateNInitializedWithFirst(inputs)
+      acc_n_merge = self._AccumulateNInitializedWithMerge(inputs)
+      acc_n_shape = self._AccumulateNInitializedWithShape(inputs)
+
+    test_ops = (("AddN", add_n.op),
+                ("AccNFirst", acc_n_first.op),
+                ("AccNMerge", acc_n_merge.op),
+                ("AccNShape", acc_n_shape.op))
+
+    with session.Session(graph=graph):
+      for tag, op in test_ops:
+        for _ in range(100):
+          op.run()  # Run for warm up.
+        start = time.time()
+        for _ in range(repeats):
+          op.run()
+        duration = time.time() - start
+        args = format_args + (tag, duration)
+        print(self._template.format(*args))
+
+  def _RunBenchmark(self, tag, input_fn, sizes, ninputs, repeats):
+    for size in sizes:
+      for ninput in ninputs:
+        graph = ops.Graph()
+        with graph.as_default():
+          inputs = input_fn(size, ninput)
+
+        format_args = (tag, size, ninput, repeats)
+        self._SetupAndRunBenchmark(graph, inputs, repeats, format_args)
+
+  def benchmarkAccumulateN(self):
+    self._template = "{:<15}" * 6
+    args = {
+        "sizes": (128, 128**2),
+        "ninputs": (1, 10, 100, 300),
+        "repeats": 100
+    }
+    benchmarks = (("Replicated", self._GenerateReplicatedInputs),
+                  ("Unordered", self._GenerateUnorderedInputs),
+                  ("Ordered", self._GenerateOrderedInputs),
+                  ("Reversed", self._GenerateReversedInputs))
+
+    print(self._template.format("", "Size", "#Inputs", "#Repeat", "Method",
+                                "Duration"))
+    print("-" * 90)
+    for benchmark in benchmarks:
+      self._RunBenchmark(*benchmark, **args)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/ops/array_grad_test.py

@@ -0,0 +1,187 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for array_grad."""
+
+from tensorflow.python.eager import backprop
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gradient_checker_v2
+from tensorflow.python.platform import test
+
+
+@test_util.with_eager_op_as_function
+@test_util.run_all_in_graph_and_eager_modes
+class ArrayGradTest(test.TestCase):
+
+  def _testGrad(self, f, x):
+    max_error = gradient_checker_v2.max_error(
+        *gradient_checker_v2.compute_gradient(f, [x]))
+    self.assertLess(max_error, 1e-4)
+
+  def test_gather_v2_simple(self):
+    x = constant_op.constant([1., 2., 3., 4., 5.], dtype=dtypes.float64)
+
+    def f(x):
+      return array_ops.gather_v2(
+          x, constant_op.constant([2, 0, 2, 4], dtype=dtypes.int32))
+
+    self._testGrad(f, x)
+
+  def test_gather_v2_more_index_dims(self):
+    x = constant_op.constant([1., 2., 3., 4., 5.], dtype=dtypes.float64)
+
+    def f(x):
+      return array_ops.gather_v2(
+          x, constant_op.constant([[2, 0], [2, 4]], dtype=dtypes.int32))
+
+    self._testGrad(f, x)
+
+  def test_gather_v2_more_param_dims(self):
+    x = constant_op.constant([[1., 2.], [3., 4.]], dtype=dtypes.float64)
+
+    def f(x):
+      return array_ops.gather_v2(
+          x, constant_op.constant([1, 0], dtype=dtypes.int32))
+
+    self._testGrad(f, x)
+
+  def test_gather_v2_axis(self):
+    x = constant_op.constant([[1., 2.], [3., 4.]], dtype=dtypes.float64)
+
+    def f(x):
+      return array_ops.gather_v2(
+          x, constant_op.constant([1, 0], dtype=dtypes.int32), axis=1)
+
+    self._testGrad(f, x)
+
+  def test_gather_v2_batch_dims(self):
+    x = constant_op.constant([[1., 2.], [3., 4.]], dtype=dtypes.float64)
+
+    def f(x):
+      return array_ops.gather_v2(
+          x,
+          constant_op.constant([[1, 0], [0, 0]], dtype=dtypes.int32),
+          axis=1,
+          batch_dims=1)
+
+    self._testGrad(f, x)
+
+  def test_gather_v2_2batch_dims(self):
+    x = constant_op.constant([[[1., 2.], [3., 4.]]], dtype=dtypes.float64)
+
+    def f(x):
+      return array_ops.gather_v2(
+          x,
+          constant_op.constant([[[1, 0], [0, 0]]], dtype=dtypes.int32),
+          axis=2,
+          batch_dims=2)
+
+    self._testGrad(f, x)
+
+  def test_gather_v2_batch_dims_with_axis(self):
+    x = constant_op.constant([[[1., 2.]], [[3., 4.]]], dtype=dtypes.float64)
+
+    def f(x):
+      return array_ops.gather_v2(
+          x,
+          constant_op.constant([[0], [0]], dtype=dtypes.int32),
+          axis=2,
+          batch_dims=1)
+
+    self._testGrad(f, x)
+
+  def test_gather_v2_zero_bsize_grad_has_matching_shapes(self):
+    params = array_ops.zeros(shape=[0, 1, 8, 16], dtype=dtypes.float64)
+    indices = array_ops.zeros(shape=[0, 1, 3], dtype=dtypes.int32)
+
+    def f(params):
+      return array_ops.gather_v2(params, indices, axis=2, batch_dims=2)
+
+    grads = backprop.gradients_function(f)(params)
+    self.assertLen(grads, 1)
+    self.assertAllEqual(params.shape, grads[0].shape)
+
+  def test_broadcast_to(self):
+    x = constant_op.constant([1., 2., 3.], dtype=dtypes.float64)
+    y = constant_op.constant([2, 3], dtype=dtypes.int32)
+
+    def f(x):
+      return array_ops.broadcast_to(
+          x,
+          y)
+
+    self._testGrad(f, x)
+
+  def test_broadcast_to_int64(self):
+    x = constant_op.constant([1., 2., 3.], dtype=dtypes.float64)
+    y = constant_op.constant([2, 3], dtype=dtypes.int64)
+
+    def f(x):
+      return array_ops.broadcast_to(
+          x,
+          y)
+
+    self._testGrad(f, x)
+
+  def test_slice_int64(self):
+    x = constant_op.constant([1., 2., 3.], dtype=dtypes.float64)
+    begin = constant_op.constant([1], dtype=dtypes.int64)
+    size = constant_op.constant([1], dtype=dtypes.int64)
+
+    def f(x):
+      return array_ops.slice(x, begin, size)
+
+    self._testGrad(f, x)
+
+  def test_reshape_simple(self):
+    x = constant_op.constant([1., 2., 3.], dtype=dtypes.float64)
+    y = constant_op.constant([3, 1], dtype=dtypes.int64)
+
+    def f(x):
+      return array_ops.reshape(x, y)
+
+    self._testGrad(f, x)
+
+  def test_reshape_one_unknown_dim(self):
+    def f(x):
+      x_without_shape = array_ops.placeholder_with_default(x, shape=[None, 2])
+      return array_ops.reshape(x_without_shape, [3, 2])
+
+    x = constant_op.constant([[1., 2.], [3., 4.], [5., 6.]],
+                             dtype=dtypes.float64)
+    self._testGrad(f, x)
+
+  def test_reshape_two_unknown_dims(self):
+    def f(x):
+      x_without_shape = array_ops.placeholder_with_default(x,
+                                                           shape=[None, None])
+      return array_ops.reshape(x_without_shape, [6])
+
+    x = constant_op.constant([[1., 2., 3.], [4., 5., 6.]], dtype=dtypes.float64)
+    self._testGrad(f, x)
+
+  def test_reshape_one_unknown_dim_and_zero_elements(self):
+    def f(x):
+      x_without_shape = array_ops.placeholder_with_default(x, shape=[None, 0])
+      return array_ops.reshape(x_without_shape, [0])
+
+    x = constant_op.constant([], shape=[3, 0], dtype=dtypes.float64)
+    self._testGrad(f, x)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/ops/array_ops_shape_test.py

@@ -0,0 +1,51 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for shape op int64 output."""
+
+from tensorflow.core.config import flags
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+class ArrayOpShapeSizeTest(test.TestCase):
+
+  def testShapeInt64Flag(self):
+    # The tf_shape_default_int64 flag should be set when this test runs
+    self.assertTrue(flags.config().tf_shape_default_int64.value())
+    s1 = array_ops.shape_v2(array_ops.zeros([1, 2]))
+    self.assertEqual(s1.dtype, dtypes.int64)
+
+  def testShapeInt64FlagTf1(self):
+    # The tf_shape_default_int64 flag should be set when this test runs
+    self.assertTrue(flags.config().tf_shape_default_int64.value())
+    s1 = array_ops.shape(array_ops.zeros([1, 2]))
+    self.assertEqual(s1.dtype, dtypes.int64)
+
+  def testSizeInt64Flag(self):
+    # The tf_shape_default_int64 flag should be set when this test runs
+    self.assertTrue(flags.config().tf_shape_default_int64.value())
+    s1 = array_ops.size_v2(array_ops.zeros([1, 2]))
+    self.assertEqual(s1.dtype, dtypes.int64)
+
+  def testSizeInt64FlagTf1(self):
+    # The tf_shape_default_int64 flag should be set when this test runs
+    self.assertTrue(flags.config().tf_shape_default_int64.value())
+    s1 = array_ops.size(array_ops.zeros([1, 2]))
+    self.assertEqual(s1.dtype, dtypes.int64)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/ops/array_ops_stack.py

@@ -0,0 +1,214 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Tests for this file live in python/kernel_tests/array_ops_test.py
+"""Operations to stack and unstack tensors."""
+
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.util import dispatch
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("stack")
+@dispatch.add_dispatch_support
+def stack(values, axis=0, name="stack"):
+  """Stacks a list of rank-`R` tensors into one rank-`(R+1)` tensor.
+
+  See also `tf.concat`, `tf.tile`, `tf.repeat`.
+
+  Packs the list of tensors in `values` into a tensor with rank one higher than
+  each tensor in `values`, by packing them along the `axis` dimension.
+  Given a list of length `N` of tensors of shape `(A, B, C)`;
+
+  If `axis == 0` then the `output` tensor will have the shape `(N, A, B, C)`.
+  If `axis == 1` then the `output` tensor will have the shape `(A, N, B, C)`.
+  Etc.
+
+  For example:
+
+  >>> x = tf.constant([1, 4])
+  >>> y = tf.constant([2, 5])
+  >>> z = tf.constant([3, 6])
+  >>> tf.stack([x, y, z])
+  <tf.Tensor: shape=(3, 2), dtype=int32, numpy=
+  array([[1, 4],
+         [2, 5],
+         [3, 6]], dtype=int32)>
+  >>> tf.stack([x, y, z], axis=1)
+  <tf.Tensor: shape=(2, 3), dtype=int32, numpy=
+  array([[1, 2, 3],
+         [4, 5, 6]], dtype=int32)>
+
+  This is the opposite of unstack.  The numpy equivalent is `np.stack`
+
+  >>> np.array_equal(np.stack([x, y, z]), tf.stack([x, y, z]))
+  True
+
+  Args:
+    values: A list of `Tensor` objects with the same shape and type.
+    axis: An `int`. The axis to stack along. Defaults to the first dimension.
+      Negative values wrap around, so the valid range is `[-(R+1), R+1)`.
+    name: A name for this operation (optional).
+
+  Returns:
+    output: A stacked `Tensor` with the same type as `values`.
+
+  Raises:
+    ValueError: If `axis` is out of the range [-(R+1), R+1).
+  """
+  if axis == 0:
+    try:
+      # If the input is a constant list, it can be converted to a constant op
+      return ops.convert_to_tensor(values, name=name)
+    except (TypeError, ValueError, NotImplementedError):
+      pass  # Input list contains non-constant tensors
+
+  value_shape = ops.convert_to_tensor(values[0], name=name)._shape_tuple()  # pylint: disable=protected-access
+  if value_shape is not None:
+    expanded_num_dims = len(value_shape) + 1
+    if axis < -expanded_num_dims or axis >= expanded_num_dims:
+      raise ValueError(f"Argument `axis` = {axis} not in range "
+                       f"[{-expanded_num_dims}, {expanded_num_dims})")
+
+  return gen_array_ops.pack(values, axis=axis, name=name)
+
+
+@tf_export("unstack")
+@dispatch.add_dispatch_support
+def unstack(value, num=None, axis=0, name="unstack"):
+  """Unpacks the given dimension of a rank-`R` tensor into rank-`(R-1)` tensors.
+
+  Unpacks tensors from `value` by chipping it along the `axis` dimension.
+
+  >>> x = tf.reshape(tf.range(12), (3,4))
+  >>>
+  >>> p, q, r = tf.unstack(x)
+  >>> p.shape.as_list()
+  [4]
+
+  >>> i, j, k, l = tf.unstack(x, axis=1)
+  >>> i.shape.as_list()
+  [3]
+
+  This is the opposite of stack.
+
+  >>> x = tf.stack([i, j, k, l], axis=1)
+
+  More generally if you have a tensor of shape `(A, B, C, D)`:
+
+  >>> A, B, C, D = [2, 3, 4, 5]
+  >>> t = tf.random.normal(shape=[A, B, C, D])
+
+  The number of tensor returned is equal to the length of the target `axis`:
+
+  >>> axis = 2
+  >>> items = tf.unstack(t, axis=axis)
+  >>> len(items) == t.shape[axis]
+  True
+
+  The shape of each result tensor is equal to the shape of the input tensor,
+  with the target `axis` removed.
+
+  >>> items[0].shape.as_list()  # [A, B, D]
+  [2, 3, 5]
+
+  The value of each tensor `items[i]` is equal to the slice of `input` across
+  `axis` at index `i`:
+
+  >>> for i in range(len(items)):
+  ...   slice = t[:,:,i,:]
+  ...   assert tf.reduce_all(slice == items[i])
+
+  #### Python iterable unpacking
+
+  With eager execution you _can_ unstack the 0th axis of a tensor using python's
+  iterable unpacking:
+
+  >>> t = tf.constant([1,2,3])
+  >>> a,b,c = t
+
+  `unstack` is still necessary because Iterable unpacking doesn't work in
+  a `@tf.function`: Symbolic tensors are not iterable.
+
+  You need to use `tf.unstack` here:
+
+  >>> @tf.function
+  ... def bad(t):
+  ...   a,b,c = t
+  ...   return a
+  >>>
+  >>> bad(t)
+  Traceback (most recent call last):
+  ...
+  OperatorNotAllowedInGraphError: ...
+
+  >>> @tf.function
+  ... def good(t):
+  ...   a,b,c = tf.unstack(t)
+  ...   return a
+  >>>
+  >>> print(good(t).numpy())
+  1
+
+  #### Unknown shapes
+
+  Eager tensors have concrete values, so their shape is always known.
+  Inside a `tf.function` the symbolic tensors may have unknown shapes.
+  If the length of `axis` is unknown `tf.unstack` will fail because it cannot
+  handle an unknown number of tensors:
+
+  >>> @tf.function(input_signature=[tf.TensorSpec([None], tf.float32)])
+  ... def bad(t):
+  ...   tensors = tf.unstack(t)
+  ...   return tensors[0]
+  >>>
+  >>> bad(tf.constant([1.0, 2.0, 3.0]))
+  Traceback (most recent call last):
+  ...
+  ValueError: Cannot infer argument `num` from shape (None,)
+
+  If you know the `axis` length you can pass it as the `num` argument. But this
+  must be a constant value.
+
+  If you actually need a variable number of tensors in a single `tf.function`
+  trace, you will need to use explicit loops and a `tf.TensorArray` instead.
+
+  Args:
+    value: A rank `R > 0` `Tensor` to be unstacked.
+    num: An `int`. The length of the dimension `axis`. Automatically inferred if
+      `None` (the default).
+    axis: An `int`. The axis to unstack along. Defaults to the first dimension.
+      Negative values wrap around, so the valid range is `[-R, R)`.
+    name: A name for the operation (optional).
+
+  Returns:
+    The list of `Tensor` objects unstacked from `value`.
+
+  Raises:
+    ValueError: If `axis` is out of the range `[-R, R)`.
+    ValueError: If `num` is unspecified and cannot be inferred.
+    InvalidArgumentError: If `num` does not match the shape of `value`.
+  """
+  if num is None:
+    value = ops.convert_to_tensor(value)
+    value_shape = value.get_shape()
+    if value_shape.ndims is not None:
+      if axis < -value_shape.ndims or axis >= value_shape.ndims:
+        raise ValueError(f"Argument `axis` = {axis} not in range "
+                         f"[{-value_shape.ndims}, {value_shape.ndims})")
+      num = value_shape.dims[axis].value
+    if num is None:
+      raise ValueError(f"Cannot infer argument `num` from shape {value_shape}")
+  return gen_array_ops.unpack(value, num=num, axis=axis, name=name)

tensorflow/python/ops/array_ops_test.py

@@ -0,0 +1,154 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for array operations."""
+
+from tensorflow.core.config import flags
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.platform import test
+
+
+class ArrayOpTest(test.TestCase):
+
+  def testGatherGradHasPartialStaticShape(self):
+    # Create a tensor with an unknown dim 1.
+    x = random_ops.random_normal([4, 10, 10])
+    x = array_ops.gather(
+        x, array_ops.reshape(array_ops.where_v2(x[0, :, 0] > 0.5), [-1]), axis=1
+    )
+    x.shape.assert_is_compatible_with([4, None, 10])
+
+    with backprop.GradientTape() as tape:
+      tape.watch(x)
+      a = array_ops.gather(array_ops.gather(x, [0, 1]), [0, 1])
+    grad_a = tape.gradient(a, x)
+    with backprop.GradientTape() as tape:
+      tape.watch(x)
+      b = array_ops.gather(array_ops.gather(x, [2, 3], axis=2), [0, 1])
+    grad_b = tape.gradient(b, x)
+
+    # We make sure that the representation of the shapes are correct; the shape
+    # equality check will always eval to false due to the shapes being partial.
+    grad_a.shape.assert_is_compatible_with([None, None, 10])
+    grad_b.shape.assert_is_compatible_with([4, None, 10])
+
+  def testReshapeShapeInference(self):
+    # Create a tensor with an unknown dim 1.
+    x = random_ops.random_normal([4, 10, 10])
+    x = array_ops.gather(
+        x, array_ops.reshape(array_ops.where_v2(x[0, :, 0] > 0.5), [-1]), axis=1
+    )
+    x.shape.assert_is_compatible_with([4, None, 10])
+    a = array_ops.reshape(x, array_ops.shape(x))
+    a.shape.assert_is_compatible_with([4, None, 10])
+    b = array_ops.reshape(x, math_ops.cast(array_ops.shape(x), dtypes.int64))
+    b.shape.assert_is_compatible_with([4, None, 10])
+
+    # We do not shape-infer across a tf.cast into anything that's not tf.int32
+    # or tf.int64, since they might end up mangling the shape.
+    c = array_ops.reshape(
+        x,
+        math_ops.cast(
+            math_ops.cast(array_ops.shape(x), dtypes.float32), dtypes.int32
+        ),
+    )
+    c.shape.assert_is_compatible_with([None, None, None])
+
+  def testEmptyMeshgrid(self):
+    self.assertEqual(array_ops.meshgrid(), [])
+
+  def testSlicedPartialShapeInference(self):
+    @def_function.function(autograph=False)
+    def g(x):
+      return array_ops.zeros([array_ops.shape(x)[0]])
+
+    conc = g.get_concrete_function(tensor_spec.TensorSpec([10, None]))
+    self.assertAllEqual(conc.output_shapes.as_list(), [10])
+
+  def testIdentityOnSlicedPartialShapeInference(self):
+    @def_function.function(autograph=False)
+    def g(x):
+      return array_ops.zeros([array_ops.identity(array_ops.shape(x)[0])])
+
+    conc = g.get_concrete_function(tensor_spec.TensorSpec([10, None]))
+    self.assertAllEqual(conc.output_shapes.as_list(), [10])
+
+  @test_util.run_in_graph_and_eager_modes
+  def testParallelConcatFailsWithRankZeroShape(self):
+    op = array_ops.ParallelConcat
+    para = {"shape": 0, "values": [1]}
+
+    def func():
+      y = op(**para)
+      return y
+
+    with self.assertRaisesRegex(
+        Exception, "(rank|dimension) of .* must be greater than .* 0"
+    ):
+      func()
+
+  @test_util.run_in_graph_and_eager_modes
+  def testUpperBoundValuesWrongRank(self):
+    # Used to cause a segfault, b/266336058
+    arg0 = array_ops.zeros([2, 3], dtype=dtypes.float32)
+    arg1 = array_ops.zeros([2, 1, 0], dtype=dtypes.float32)
+    with self.assertRaisesRegex(
+        Exception, "Shape must be rank 2 but is rank 3"
+    ):
+      gen_array_ops.upper_bound(arg0, arg1)
+
+  def testLowerBoundValuesWrongRank(self):
+    # Used to cause a segfault, b/266336058
+    arg0 = array_ops.zeros([2, 3], dtype=dtypes.float32)
+    arg1 = array_ops.zeros([2, 1, 0], dtype=dtypes.float32)
+    with self.assertRaisesRegex(
+        Exception, "Shape must be rank 2 but is rank 3"
+    ):
+      gen_array_ops.lower_bound(arg0, arg1)
+
+  def testUpperBoundInputsWrongRank(self):
+    # Used to cause a segfault, b/266336058
+    arg0 = array_ops.zeros([2, 1, 0], dtype=dtypes.float32)
+    arg1 = array_ops.zeros([2, 3], dtype=dtypes.float32)
+    with self.assertRaisesRegex(
+        Exception, "Shape must be rank 2 but is rank 3"
+    ):
+      gen_array_ops.upper_bound(arg0, arg1)
+
+  def testLowerBoundInputsWrongRank(self):
+    # Used to cause a segfault, b/266336058
+    arg0 = array_ops.zeros([2, 1, 0], dtype=dtypes.float32)
+    arg1 = array_ops.zeros([2, 3], dtype=dtypes.float32)
+    with self.assertRaisesRegex(
+        Exception, "Shape must be rank 2 but is rank 3"
+    ):
+      gen_array_ops.lower_bound(arg0, arg1)
+
+  def testShapeDefaultIn32(self):
+    # The tf_shape_default_int64 flag should NOT be set when this test runs
+    self.assertFalse(flags.config().tf_shape_default_int64.value())
+    s1 = array_ops.shape_v2(array_ops.zeros([1, 2]))
+    self.assertEqual(s1.dtype, dtypes.int32)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/ops/autograph_ops.py

@@ -0,0 +1,146 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Autograph specific overrides for objects covered by tensor_util.is_tf_type."""
+
+from tensorflow.python.autograph.operators import py_builtins
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import script_ops
+from tensorflow.python.ops import sort_ops
+from tensorflow.python.ops.parallel_for import control_flow_ops as parallel_ops
+
+
+def wrap_py_func(f, args, kwargs=None):
+  """Helper that wraps a callable to py_func.
+
+  The helper passes tensor arguments through the py_func interface. Non-tensor
+  arguments are allowed, and will be passed to f directly. Note that non-tensor
+  arguments are captured by f will not update every time the wrapper is
+  called (this is consistent with its argument list, which only includes
+  the tensor arguments). In general, it's safest not to reuse this wrapper.
+
+  Args:
+    f: Callable
+    args: Positional arguments for f, as list or tuple.
+    kwargs: Keyword arguments for f, as dict with string keys. May be None.
+
+  Returns:
+    The return values of f converted to tensor.
+  Raises:
+    ValueError: if any of the arguments are incorrect.
+  """
+  tensor_args = []
+  tensor_args_idx = {}
+
+  # Of the positional arguments, only grab the tensor ones to be passed through
+  # the py_func.
+  n_args = len(args)
+  arg_is_tensor = tuple(map(tensor_util.is_tf_type, args))
+  for i in range(n_args):
+    if arg_is_tensor[i]:
+      tensor_args_idx[i] = len(tensor_args)
+      tensor_args.append(args[i])
+
+  # We essentially take the tensor kwargs, if any, and add them to the list of
+  # positional arguments. The kwargs are then reconstructed inside the py_func.
+  #
+  # For example, if
+  #
+  #     args = [Tensor(1), 'foo']
+  #     kwargs = {'a': Tensor(2), 'b': 'bar'}
+  #
+  # Then
+  #
+  #     tensor_args = (Tensor(1), Tensor(2))
+  #     kwarg_keys = ('a', 'b')
+  if kwargs:
+    kwarg_keys = tuple(kwargs.keys())
+    kwarg_is_tensor = {k: tensor_util.is_tf_type(kwargs[k]) for k in kwarg_keys}
+    for k in kwarg_keys:
+      if kwarg_is_tensor[k]:
+        tensor_args_idx[k] = len(tensor_args)
+        tensor_args.append(kwargs[k])
+  else:
+    kwarg_keys = ()
+
+  def f_wrapper(*tensor_args):
+    f_args = tuple(
+        tensor_args[tensor_args_idx[i]] if arg_is_tensor[i] else a
+        for i, a in enumerate(args)
+    )
+    f_kwargs = {
+        k: tensor_args[tensor_args_idx[k]] if kwarg_is_tensor[k] else kwargs[k]
+        for i, k in enumerate(kwarg_keys)
+    }
+    f(*f_args, **f_kwargs)
+    return 1
+
+  return script_ops.eager_py_func(f_wrapper, tensor_args, dtypes.int32)
+
+
+def _tf_py_func_print(*objects, **kwargs):
+  """Overload of print_ as a py_func implementation."""
+  override_kwargs = {
+      k: v for k, v in kwargs.items() if v is not py_builtins.UNSPECIFIED
+  }
+  if 'flush' not in override_kwargs:
+    # Defaulting to flushing the console in graph mode, which helps reduce
+    # garbled output in IPython.
+    override_kwargs['flush'] = True
+
+  def print_wrapper(*vals, **kwargs):
+    vals = tuple(v.numpy() if tensor_util.is_tf_type(v) else v for v in vals)
+    # TensorFlow doesn't seem to generate Unicode when passing strings to
+    # py_func. This causes the print to add a "b'" wrapper to the output,
+    # which is probably never what you want.
+    vals = tuple(v.decode('utf-8') if isinstance(v, bytes) else v for v in vals)
+    print(*vals, **kwargs)
+
+  return wrap_py_func(print_wrapper, objects, override_kwargs)
+
+
+def _tf_sorted(iterable, key, reverse):
+  """Overload of sorted_ for Tensor iterable."""
+  if reverse is py_builtins.UNSPECIFIED:
+    direction = 'ASCENDING'
+  else:
+    direction = 'DESCENDING'
+  if key is not py_builtins.UNSPECIFIED:
+    mapped = parallel_ops.vectorized_map(key, iterable)
+    if mapped.shape.rank is not None and mapped.shape.rank != 1:
+      raise ValueError('sort only supports only 1D tensors')
+    with ops.control_dependencies([
+        check_ops.assert_rank_v2(mapped, 1,
+                                 'sort only supports only 1D tensors')
+    ]):
+      order = sort_ops.argsort(mapped, direction=direction)
+      return array_ops.gather_v2(iterable, order)
+  if iterable.shape.rank is not None and iterable.shape.rank != 1:
+    raise ValueError('sort only supports only 1D tensors')
+  with ops.control_dependencies([
+      check_ops.assert_rank_v2(iterable, 1,
+                               'sort only supports only 1D tensors')
+  ]):
+    return sort_ops.sort(iterable, direction=direction)
+
+py_builtins.print_registry.register(
+    tensor_util.tf_type_classes, _tf_py_func_print
+)
+py_builtins.sorted_registry.register(
+    tensor_util.tf_type_classes, _tf_sorted
+)

tensorflow/python/ops/autograph_ops_test.py

@@ -0,0 +1,40 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for autograph_ops module."""
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import autograph_ops
+from tensorflow.python.platform import test
+
+
+class AutographOpsTest(test.TestCase):
+
+  def test_wrap_py_func_dummy_return(self):
+    side_counter = [0]
+
+    def test_fn(_):
+      side_counter[0] += 1
+
+    with self.cached_session():
+      result = autograph_ops.wrap_py_func(test_fn, (5,))
+      self.assertEqual(1, self.evaluate(result))
+      self.assertEqual([1], side_counter)
+      result = autograph_ops.wrap_py_func(test_fn, (constant_op.constant(5),))
+      self.assertEqual(1, self.evaluate(result))
+      self.assertEqual([2], side_counter)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/ops/batch_norm_benchmark.py

@@ -0,0 +1,259 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""End-to-end benchmark for batch normalization."""
+
+import argparse
+import sys
+import time
+
+from tensorflow.python.client import session as session_lib
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_nn_ops
+from tensorflow.python.ops import gradients_impl
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_impl
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import variables
+import tensorflow.python.ops.nn_grad  # pylint: disable=unused-import
+from tensorflow.python.platform import test
+
+
+def batch_norm_op(tensor, mean, variance, beta, gamma, scale):
+  """Fused kernel for batch normalization."""
+  # _batch_norm_with_global_normalization is deprecated in v9
+  test_util.set_producer_version(ops.get_default_graph(), 8)
+  # pylint: disable=protected-access
+  return gen_nn_ops._batch_norm_with_global_normalization(
+      tensor, mean, variance, beta, gamma, 0.001, scale)
+  # pylint: enable=protected-access
+
+
+# Note that the naive implementation is much slower:
+# batch_norm = (tensor - mean) * tf.math.rsqrt(variance + 0.001)
+# if scale:
+#   batch_norm *= gamma
+# return batch_norm + beta
+def batch_norm_py(tensor, mean, variance, beta, gamma, scale):
+  """Python implementation of batch normalization."""
+  return nn_impl.batch_normalization(tensor, mean, variance, beta, gamma if
+                                     scale else None, 0.001)
+
+
+def batch_norm_slow(tensor, mean, variance, beta, gamma, scale):
+  batch_norm = (tensor - mean) * math_ops.rsqrt(variance + 0.001)
+  if scale:
+    batch_norm *= gamma
+  return batch_norm + beta
+
+
+def build_graph(device, input_shape, axes, num_layers, mode, scale, train):
+  """Build a graph containing a sequence of batch normalizations.
+
+  Args:
+    device: string, the device to run on.
+    input_shape: shape of the input tensor.
+    axes: axes that are to be normalized across.
+    num_layers: number of batch normalization layers in the graph.
+    mode: "op", "py" or "slow" depending on the implementation.
+    scale: scale after normalization.
+    train: if true, also run backprop.
+
+  Returns:
+    An array of tensors to run()
+  """
+  moment_shape = []
+  keep_dims = mode == "py" or mode == "slow"
+  if keep_dims:
+    for axis in range(len(input_shape)):
+      if axis in axes:
+        moment_shape.append(1)
+      else:
+        moment_shape.append(input_shape[axis])
+  else:
+    for axis in range(len(input_shape)):
+      if axis not in axes:
+        moment_shape.append(input_shape[axis])
+  with ops.device("/%s:0" % device):
+    tensor = variables.Variable(random_ops.truncated_normal(input_shape))
+    for _ in range(num_layers):
+      if train:
+        mean, variance = nn_impl.moments(tensor, axes, keep_dims=keep_dims)
+      else:
+        mean = array_ops.zeros(moment_shape)
+        variance = array_ops.ones(moment_shape)
+      beta = variables.Variable(array_ops.zeros(moment_shape))
+      gamma = variables.Variable(constant_op.constant(1.0, shape=moment_shape))
+      if mode == "py":
+        tensor = batch_norm_py(tensor, mean, variance, beta, gamma, scale)
+      elif mode == "op":
+        tensor = batch_norm_op(tensor, mean, variance, beta, gamma, scale)
+      elif mode == "slow":
+        tensor = batch_norm_slow(tensor, mean, variance, beta, gamma, scale)
+    if train:
+      return gradients_impl.gradients([tensor], variables.trainable_variables())
+    else:
+      return [tensor]
+
+
+def print_difference(mode, t1, t2):
+  """Print the difference in timing between two runs."""
+  difference = (t2 - t1) / t1 * 100.0
+  print("=== %s: %.1f%% ===" % (mode, difference))
+
+
+class BatchNormBenchmark(test.Benchmark):
+  """Benchmark batch normalization."""
+
+  def _run_graph(self, device, input_shape, axes, num_layers, mode, scale,
+                 train, num_iters):
+    """Run the graph and print its execution time.
+
+    Args:
+      device: string, the device to run on.
+      input_shape: shape of the input tensor.
+      axes: axes that are to be normalized across.
+      num_layers: number of batch normalization layers in the graph.
+      mode: "op", "py" or "slow" depending on the implementation.
+      scale: scale after normalization.
+      train: if true, also run backprop.
+      num_iters: number of steps to run.
+
+    Returns:
+      The duration of the run in seconds.
+    """
+    graph = ops.Graph()
+    with graph.as_default():
+      outputs = build_graph(device, input_shape, axes, num_layers, mode, scale,
+                            train)
+    with session_lib.Session(graph=graph) as session:
+      variables.global_variables_initializer().run()
+      _ = session.run([out.op for out in outputs])  # warm up.
+      start_time = time.time()
+      for _ in range(num_iters):
+        _ = session.run([out.op for out in outputs])
+      duration = time.time() - start_time
+    print("%s shape:%d/%d #layers:%d mode:%s scale:%r train:%r - %f secs" %
+          (device, len(input_shape), len(axes), num_layers, mode, scale, train,
+           duration / num_iters))
+
+    name_template = (
+        "batch_norm_{device}_input_shape_{shape}_axes_{axes}_mode_{mode}_"
+        "layers_{num_layers}_scale_{scale}_"
+        "train_{train}")
+
+    self.report_benchmark(
+        name=name_template.format(
+            device=device,
+            mode=mode,
+            num_layers=num_layers,
+            scale=scale,
+            train=train,
+            shape=str(input_shape).replace(" ", ""),
+            axes=str(axes)).replace(" ", ""),
+        iters=num_iters,
+        wall_time=duration / num_iters)
+
+    return duration
+
+  def benchmark_batch_norm(self):
+    print("Forward convolution (lower layers).")
+    shape = [8, 128, 128, 32]
+    axes = [0, 1, 2]
+    t1 = self._run_graph("cpu", shape, axes, 10, "op", True, False, 5)
+    t2 = self._run_graph("cpu", shape, axes, 10, "py", True, False, 5)
+    t3 = self._run_graph("cpu", shape, axes, 10, "slow", True, False, 5)
+    print_difference("op vs py", t1, t2)
+    print_difference("py vs slow", t2, t3)
+    if FLAGS.use_gpu:
+      t1 = self._run_graph("gpu", shape, axes, 10, "op", True, False, 50)
+      t2 = self._run_graph("gpu", shape, axes, 10, "py", True, False, 50)
+      t3 = self._run_graph("gpu", shape, axes, 10, "slow", True, False, 50)
+      print_difference("op vs py", t1, t2)
+      print_difference("py vs slow", t2, t3)
+    print("Forward/backward convolution (lower layers).")
+    t1 = self._run_graph("cpu", shape, axes, 10, "op", True, True, 5)
+    t2 = self._run_graph("cpu", shape, axes, 10, "py", True, True, 5)
+    t3 = self._run_graph("cpu", shape, axes, 10, "slow", True, True, 5)
+    print_difference("op vs py", t1, t2)
+    print_difference("py vs slow", t2, t3)
+    if FLAGS.use_gpu:
+      t1 = self._run_graph("gpu", shape, axes, 10, "op", True, True, 50)
+      t2 = self._run_graph("gpu", shape, axes, 10, "py", True, True, 50)
+      t3 = self._run_graph("gpu", shape, axes, 10, "slow", True, True, 50)
+      print_difference("op vs py", t1, t2)
+      print_difference("py vs slow", t2, t3)
+    print("Forward convolution (higher layers).")
+    shape = [256, 17, 17, 32]
+    axes = [0, 1, 2]
+    t1 = self._run_graph("cpu", shape, axes, 10, "op", True, False, 5)
+    t2 = self._run_graph("cpu", shape, axes, 10, "py", True, False, 5)
+    t3 = self._run_graph("cpu", shape, axes, 10, "slow", True, False, 5)
+    print_difference("op vs py", t1, t2)
+    print_difference("py vs slow", t2, t3)
+    if FLAGS.use_gpu:
+      t1 = self._run_graph("gpu", shape, axes, 10, "op", True, False, 50)
+      t2 = self._run_graph("gpu", shape, axes, 10, "py", True, False, 50)
+      t3 = self._run_graph("gpu", shape, axes, 10, "slow", True, False, 50)
+      print_difference("op vs py", t1, t2)
+      print_difference("py vs slow", t2, t3)
+    print("Forward/backward convolution (higher layers).")
+    t1 = self._run_graph("cpu", shape, axes, 10, "op", True, True, 5)
+    t2 = self._run_graph("cpu", shape, axes, 10, "py", True, True, 5)
+    t3 = self._run_graph("cpu", shape, axes, 10, "slow", True, True, 5)
+    print_difference("op vs py", t1, t2)
+    print_difference("py vs slow", t2, t3)
+    if FLAGS.use_gpu:
+      t1 = self._run_graph("gpu", shape, axes, 10, "op", True, True, 50)
+      t2 = self._run_graph("gpu", shape, axes, 10, "py", True, True, 50)
+      t3 = self._run_graph("gpu", shape, axes, 10, "slow", True, True, 50)
+      print_difference("op vs py", t1, t2)
+      print_difference("py vs slow", t2, t3)
+    print("Forward fully-connected.")
+    shape = [1024, 32]
+    axes = [0]
+    t1 = self._run_graph("cpu", shape, axes, 10, "py", True, False, 5)
+    t2 = self._run_graph("cpu", shape, axes, 10, "slow", True, False, 5)
+    print_difference("py vs slow", t1, t2)
+    if FLAGS.use_gpu:
+      t1 = self._run_graph("gpu", shape, axes, 10, "py", True, False, 50)
+      t2 = self._run_graph("gpu", shape, axes, 10, "slow", True, False, 50)
+      print_difference("py vs slow", t1, t2)
+    print("Forward/backward fully-connected.")
+    t1 = self._run_graph("cpu", shape, axes, 10, "py", True, True, 50)
+    t2 = self._run_graph("cpu", shape, axes, 10, "slow", True, True, 50)
+    print_difference("py vs slow", t1, t2)
+    if FLAGS.use_gpu:
+      t1 = self._run_graph("gpu", shape, axes, 10, "py", True, True, 5)
+      t2 = self._run_graph("gpu", shape, axes, 10, "slow", True, True, 5)
+      print_difference("py vs slow", t1, t2)
+
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.register("type", "bool", lambda v: v.lower() == "true")
+  parser.add_argument(
+      "--use_gpu",
+      type="bool",
+      nargs="?",
+      const=True,
+      default=True,
+      help="Run GPU benchmarks."
+  )
+  global FLAGS  # pylint:disable=global-at-module-level
+  FLAGS, unparsed = parser.parse_known_args()
+  test.main(argv=[sys.argv[0]] + unparsed)

tensorflow/python/ops/batch_ops.py

@@ -0,0 +1,123 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Operations for automatic batching and unbatching."""
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import gen_batch_ops
+# pylint: disable=wildcard-import
+from tensorflow.python.ops.gen_batch_ops import *
+# pylint: enable=wildcard-import
+from tensorflow.python.util import nest
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("nondifferentiable_batch_function")
+def batch_function(num_batch_threads,
+                   max_batch_size,
+                   batch_timeout_micros,
+                   allowed_batch_sizes=None,
+                   max_enqueued_batches=10,
+                   autograph=True,
+                   enable_large_batch_splitting=True):
+  """Batches the computation done by the decorated function.
+
+  So, for example, in the following code
+
+  ```python
+  @batch_function(1, 2, 3)
+  def layer(a):
+    return tf.matmul(a, a)
+
+  b = layer(w)
+  ```
+
+  if more than one session.run call is simultaneously trying to compute `b`
+  the values of `w` will be gathered, non-deterministically concatenated
+  along the first axis, and only one thread will run the computation. See the
+  documentation of the `Batch` op for more details.
+
+  Assumes that all arguments of the decorated function are Tensors which will
+  be batched along their first dimension.
+
+  SparseTensor is not supported. The return value of the decorated function
+  must be a Tensor or a list/tuple of Tensors.
+
+  Args:
+    num_batch_threads: Number of scheduling threads for processing batches
+     of work. Determines the number of batches processed in parallel.
+    max_batch_size: Batch sizes will never be bigger than this.
+    batch_timeout_micros: Maximum number of microseconds to wait before
+     outputting an incomplete batch.
+    allowed_batch_sizes: Optional list of allowed batch sizes. If left empty,
+     does nothing. Otherwise, supplies a list of batch sizes, causing the op
+     to pad batches up to one of those sizes. The entries must increase
+     monotonically, and the final entry must equal max_batch_size.
+    max_enqueued_batches: The maximum depth of the batch queue. Defaults to 10.
+    autograph: Whether to use autograph to compile python and eager style code
+     for efficient graph-mode execution.
+    enable_large_batch_splitting: The value of this option doesn't affect
+     processing output given the same input; it affects implementation details
+     as stated below: 1. Improve batching efficiency by eliminating unnecessary
+     adding. 2.`max_batch_size` specifies the limit of input and
+     `allowed_batch_sizes` specifies the limit of a task to be processed. API
+     user can give an input of size 128 when 'max_execution_batch_size'
+     is 32 -> implementation can split input of 128 into 4 x 32, schedule
+     concurrent processing, and then return concatenated results corresponding
+     to 128.
+
+  Returns:
+    The decorated function will return the unbatched computation output Tensors.
+  """
+
+  def decorator(fn):  # pylint: disable=missing-docstring
+
+    def decorated(*args):  # pylint: disable=missing-docstring
+
+      @def_function.function(autograph=autograph)
+      def computation(*computation_args):
+        return fn(*computation_args)
+
+      computation = computation.get_concrete_function(*[
+          tensor.TensorSpec(
+              dtype=x.dtype, shape=x.shape, name="batch_" + str(i))
+          for i, x in enumerate(args)
+      ])
+
+      with ops.name_scope("batch") as name:
+        for a in args:
+          if not isinstance(a, tensor.Tensor):
+            raise ValueError("All arguments to functions decorated with "
+                             "`batch_function`  are supposed to be Tensors; "
+                             f"found {a!r}.")
+        outputs = gen_batch_ops.batch_function(
+            num_batch_threads=num_batch_threads,
+            max_batch_size=max_batch_size,
+            batch_timeout_micros=batch_timeout_micros,
+            allowed_batch_sizes=allowed_batch_sizes,
+            max_enqueued_batches=max_enqueued_batches,
+            shared_name=name,
+            enable_large_batch_splitting=enable_large_batch_splitting,
+            f=computation,
+            in_tensors=list(args),
+            captured_tensors=computation.captured_inputs,
+            Tout=[o.dtype for o in computation.outputs])
+        return nest.pack_sequence_as(
+            computation.structured_outputs, outputs, expand_composites=True)
+
+    return decorated
+
+  return decorator

tensorflow/python/ops/batch_ops_test.py

@@ -0,0 +1,635 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Tests for the currently experimental in-graph batch ops."""
+import threading
+import time
+import numpy as np
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import function
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.framework.errors import InvalidArgumentError
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import batch_ops
+from tensorflow.python.ops import gen_batch_ops
+from tensorflow.python.ops import gen_functional_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import script_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+def delayed_plus1(x):
+  """Sleeps for 100ms then returns x+1."""
+  time.sleep(0.1)
+  return x + 1
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class BatchOpsTest(test.TestCase):
+  """Tests for batch_ops.{un,}batch."""
+
+  # Test for only non eager mode as batching in eager context as a functionality
+  # is TBD.
+  def testBasicBatch(self):
+    """Tests that a single batched tensor executes together and only once."""
+    if context.executing_eagerly():
+      return
+    with self.cached_session() as sess:
+      inp = array_ops.placeholder(dtype=dtypes.int32, shape=[1])
+      batched, index, _ = batch_ops.batch(
+          [inp], num_batch_threads=1, max_batch_size=2,
+          batch_timeout_micros=36000000, grad_timeout_micros=0,
+          batching_queue="")
+      thread_results = []
+
+      def worker():
+        thread_results.extend(
+            sess.run([batched, index], feed_dict={inp: [1]}))
+
+      worker_thread = threading.Thread(target=worker)
+      worker_thread.start()
+      main_results = sess.run([batched, index], feed_dict={inp: [2]})
+      worker_thread.join()
+
+      # At this point either the thread or the main did the batch and the other
+      # should have empty results.
+      if list(thread_results[0][0]):
+        batch_t = thread_results[0][0]
+        index_t = thread_results[1]
+        empty_b = main_results[0][0]
+        empty_m = main_results[1]
+      else:
+        batch_t = main_results[0][0]
+        index_t = main_results[1]
+        empty_b = thread_results[0][0]
+        empty_m = thread_results[1]
+
+      # Check that both the inputs made it out exactly once.
+      self.assertAllEqual(sorted(batch_t), (1, 2))
+      # Check that we get 2 rows in the index tensor.
+      self.assertEqual(len(index_t), 2)
+      # Check that the other ones are empty.
+      self.assertEqual(len(empty_b), 0)
+      self.assertEqual(len(empty_m), 0)
+
+  def testBatchWithPadding(self):
+    """Test that batching with padding up to an allowed batch size works."""
+    if context.executing_eagerly():
+      return
+    with self.cached_session() as sess:
+      inp = array_ops.placeholder(dtype=dtypes.int32, shape=[2])
+      batched, index, _ = batch_ops.batch(
+          [inp], num_batch_threads=1, max_batch_size=10,
+          batch_timeout_micros=100000,  # 100ms
+          allowed_batch_sizes=[5, 10],
+          grad_timeout_micros=0, batching_queue="")
+      thread_results = []
+
+      def worker():
+        thread_results.extend(
+            sess.run([batched, index], feed_dict={inp: [1, 3]}))
+
+      worker_thread = threading.Thread(target=worker)
+      worker_thread.start()
+      main_results = sess.run([batched, index], feed_dict={inp: [2, 4]})
+      worker_thread.join()
+
+      # At this point either the thread or the main did the batch and the other
+      # should have empty results.
+      if list(thread_results[0][0]):
+        batch_t = thread_results[0][0]
+      else:
+        batch_t = main_results[0][0]
+
+      # Check that the batch tensor incorporates the padding.
+      self.assertEqual(len(batch_t), 5)
+
+  def testMultipleBatch(self):
+    """Tests that multiple batched tensors execute together."""
+    if context.executing_eagerly():
+      return
+    with self.cached_session() as sess:
+      inp0 = array_ops.placeholder(dtype=dtypes.int32, shape=[1])
+      inp1 = array_ops.placeholder(dtype=dtypes.int32, shape=[1])
+      batched, _, _ = batch_ops.batch(
+          [inp0, inp1],
+          num_batch_threads=1,
+          max_batch_size=2,
+          batch_timeout_micros=36000000,
+          grad_timeout_micros=0,
+          batching_queue="")
+      thread_results = []
+
+      def worker():
+        thread_results.extend(
+            sess.run([batched], feed_dict={inp0: [1],
+                                           inp1: [2]}))
+
+      worker_thread = threading.Thread(target=worker)
+      worker_thread.start()
+      main_results = sess.run([batched], feed_dict={inp0: [2], inp1: [3]})
+      worker_thread.join()
+
+      # At this point either the thread or the main did the batch and the other
+      # should have empty results.
+      if list(thread_results[0][0]):
+        batch_t = thread_results[0]
+        empty_t = main_results[0]
+      else:
+        batch_t = main_results[0]
+        empty_t = thread_results[0]
+
+      # Assert that the tensors were batched together.
+      self.assertAllEqual(sorted(batch_t[0]), [1, 2])
+      self.assertAllEqual(sorted(batch_t[1]), [2, 3])
+      self.assertAllEqual(empty_t[0], [])
+      self.assertAllEqual(empty_t[1], [])
+
+  def testIllegalBatchDifferentDim0Sizes(self):
+    """Tests illegally feeding tensors with different dim0 sizes."""
+    if context.executing_eagerly():
+      return
+    with self.cached_session() as sess:
+      inp0 = array_ops.placeholder(dtype=dtypes.int32, shape=[1])
+      inp1 = array_ops.placeholder(dtype=dtypes.int32, shape=[2])
+      batched, index, _ = batch_ops.batch(
+          [inp0, inp1], num_batch_threads=1, max_batch_size=2,
+          batch_timeout_micros=0, grad_timeout_micros=0, batching_queue="")
+      with self.assertRaises(Exception) as raised:
+        _ = sess.run([batched, index], feed_dict={inp0: [0], inp1: [1, 2]})
+      self.assertGreater(
+          raised.exception.message.find("must have equal 0th-dimension size"),
+          0)
+
+  def testBasicUnbatch(self):
+    """Tests that batch and unbatch work together."""
+    if context.executing_eagerly():
+      return
+    with self.cached_session() as sess:
+      inp = array_ops.placeholder(dtype=dtypes.int32, shape=[1])
+      batched, index, id_t = batch_ops.batch(
+          [inp], num_batch_threads=1, max_batch_size=10,
+          batch_timeout_micros=100000,  # 100ms
+          allowed_batch_sizes=[3, 10],
+          grad_timeout_micros=0, batching_queue="")
+      computation = batched[0] + 1
+      result = batch_ops.unbatch(computation, index, id_t,
+                                 timeout_micros=1000000, shared_name="unbatch")
+      thread_results = []
+
+      def worker():
+        thread_results.extend(sess.run([result], feed_dict={inp: [1]}))
+
+      worker_thread = threading.Thread(target=worker)
+      worker_thread.start()
+      main_results = sess.run([result], feed_dict={inp: [2]})
+      worker_thread.join()
+      self.assertEqual(thread_results[0], [2])
+      self.assertEqual(main_results[0], [3])
+
+  def testBasicUnbatchDecorated(self):
+    """Tests that the batch_function decorator works."""
+    if context.executing_eagerly():
+      return
+    with self.cached_session() as sess:
+      # TODO(apassos): Removing this line causes test flakiness! Ideally should
+      # be investigated.
+      default_inp = array_ops.placeholder_with_default(2, shape=[])  # pylint: disable=unused-variable
+
+      @batch_ops.batch_function(1, 10, 100000)
+      def computation(in_t):
+        self.assertTrue(in_t.shape is not None)
+        return in_t + 1
+
+      inp = array_ops.placeholder(dtype=dtypes.int32, shape=[1])
+      result = computation(inp)
+      thread_results = []
+
+      def worker():
+        thread_results.extend(sess.run([result], feed_dict={inp: [1]}))
+
+      worker_thread = threading.Thread(target=worker)
+      worker_thread.start()
+      main_results = sess.run([result], feed_dict={inp: [2]})
+      worker_thread.join()
+      self.assertEqual(thread_results[0], [2])
+      self.assertEqual(main_results[0], [3])
+
+  def testUnbatchInvalidIdArg(self):
+    """Tests that unbatch work together."""
+    if context.executing_eagerly():
+      batched_tensor = constant_op.constant(
+          value=np.random.random(size=(3, 3, 1)), dtype=dtypes.float64)
+      batched_index = constant_op.constant(
+          value=np.random.randint(0, 100, size=(3, 3, 1)), dtype=dtypes.int64)
+      arg_id = constant_op.constant(
+          value=np.random.randint(0, 100, size=(3, 3, 1)), dtype=dtypes.int64)
+
+      with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                  "Input id should be scalar;"):
+        batch_ops.unbatch(
+            batched_tensor=batched_tensor,
+            batch_index=batched_index,
+            id=arg_id,
+            timeout_micros=50,
+            container="",
+            shared_name="")
+
+  def testBatchDecoratedWithCapturedInput(self):
+    """Tests that the batch_function decorator works."""
+    if context.executing_eagerly():
+      return
+    with self.cached_session() as sess:
+      captured_inp0 = array_ops.placeholder_with_default(2., shape=[])
+      captured_inp1 = resource_variable_ops.ResourceVariable(3.)
+      with ops.device("/cpu:0"):
+        captured_inp2 = resource_variable_ops.ResourceVariable(4.)
+
+      @batch_ops.batch_function(1, 10, 100000)
+      def computation(in_t):
+        return in_t + captured_inp0 + captured_inp1 + captured_inp2
+
+      inp = array_ops.placeholder(dtype=dtypes.float32, shape=[1])
+      result = computation(inp)
+      thread_results = []
+
+      def worker():
+        thread_results.extend(sess.run([result], feed_dict={inp: [1]}))
+
+      sess.run(variables.global_variables_initializer())
+      worker_thread = threading.Thread(target=worker)
+      worker_thread.start()
+      main_results = sess.run([result], feed_dict={inp: [2]})
+      worker_thread.join()
+      self.assertEqual(thread_results[0], [10])
+      self.assertEqual(main_results[0], [11])
+
+  @test_util.disable_xla("DeviceIndex returns sentinel value with XLA")
+  def testBatchDecoratedGpu(self):
+    if context.executing_eagerly():
+      return
+    with self.cached_session() as sess:
+
+      @batch_ops.batch_function(1, 10, 100000)
+      def computation(in_t):
+        # index is 0 on CPU and 1 on GPU
+        index = gen_functional_ops.DeviceIndex(device_names=["CPU", "GPU"])
+        return in_t + math_ops.cast(index, dtypes.float32)
+
+      inp = array_ops.placeholder(dtype=dtypes.float32, shape=[1])
+      result = computation(inp)
+      thread_results = []
+
+      def worker():
+        thread_results.extend(sess.run([result], feed_dict={inp: [10.]}))
+
+      worker_thread = threading.Thread(target=worker)
+      worker_thread.start()
+      main_results = sess.run([result], feed_dict={inp: [20.]})
+      worker_thread.join()
+      self.assertEqual(thread_results[0], [10 + test_util.is_gpu_available()])
+      self.assertEqual(main_results[0], [20 + test_util.is_gpu_available()])
+
+  def testParallelRunsWithCpuAndGpu(self):
+    # Run multiple instances of a batch function in parallel. This is a
+    # regression test: this used to fail because _Send nodes for one call would
+    # send the tensor to the _Recv node for a different call.
+    if context.executing_eagerly():
+      return
+    @batch_ops.batch_function(1, 2, 1)
+    def f(x):
+      with ops.device("/GPU:0"):
+        x = x + 1.
+      with ops.device("/CPU:0"):
+        return x + 1
+    num_calls = 10
+    placeholders = [array_ops.placeholder(dtypes.float32, shape=(1,))
+                    for _ in range(num_calls)]
+    results = []
+    for p in placeholders:
+      result = f(p)
+      results.append(result)
+    inputs = [[float(i)] for i in range(num_calls)]
+    expected = [[float(i + 2)] for i in range(num_calls)]
+    with self.session() as sess:
+      outputs = sess.run(results, feed_dict=dict(zip(placeholders, inputs)))
+      self.assertAllEqual(outputs, expected)
+
+  def testSoftPlacement(self):
+    if context.executing_eagerly():
+      return
+
+    @batch_ops.batch_function(1, 10, 100000)
+    def computation(in_t):
+      with ops.device("/GPU:0"):
+        return in_t + 1.
+
+    inp = array_ops.placeholder(dtype=dtypes.float32, shape=[1])
+    result = computation(inp)
+
+    # With soft placement, the function will run even without a GPU
+    config = config_pb2.ConfigProto(allow_soft_placement=True)
+    with self.session(config=config) as sess:
+      sess.run([result], feed_dict={inp: [20.]})
+
+    # Without soft placement, the function fails without a GPU due to the
+    # addition explicitly being placed on the GPU
+    config.allow_soft_placement = False
+    with self.session(config=config) as sess:
+      if test_util.is_gpu_available():
+        sess.run([result], feed_dict={inp: [20.]})
+      else:
+        with self.assertRaisesRegex(InvalidArgumentError,
+                                    "Cannot assign a device for operation"):
+          sess.run([result], feed_dict={inp: [20.]})
+
+  def testBatchFunctionOp(self):
+    """Tests that the batch_function op works."""
+    if context.executing_eagerly():
+      return
+    with self.cached_session() as sess:
+
+      @function.Defun(dtypes.int32)
+      def computation(in_t):
+        return in_t + 1
+
+      inp = array_ops.placeholder(dtype=dtypes.int32, shape=[1])
+      result = gen_batch_ops.batch_function(
+          [inp],
+          num_batch_threads=1,
+          max_batch_size=10,
+          batch_timeout_micros=100000,
+          Tout=[dtypes.int32],
+          f=computation,
+          captured_tensors=computation.captured_inputs)
+      thread_results = []
+
+      def worker():
+        thread_results.extend(sess.run([result], feed_dict={inp: [1]}))
+
+      worker_thread = threading.Thread(target=worker)
+      worker_thread.start()
+      main_results = sess.run([result], feed_dict={inp: [2]})
+      worker_thread.join()
+      self.assertEqual(thread_results[0], [2])
+      self.assertEqual(main_results[0], [3])
+
+  def testBatchFunctionOpWithCapturedInput(self):
+    """Tests that batch_function op works with captured input."""
+    if context.executing_eagerly():
+      return
+    with self.cached_session() as sess:
+      captured_inp0 = array_ops.placeholder_with_default(2, shape=[])
+      captured_inp1 = array_ops.placeholder_with_default(1, shape=[])
+      inp = array_ops.placeholder(dtype=dtypes.int32, shape=[1])
+
+      @function.Defun(dtypes.int32)
+      def computation(inp):
+        return inp + captured_inp0 - captured_inp1
+
+      result = gen_batch_ops.batch_function(
+          num_batch_threads=1,
+          max_batch_size=10,
+          batch_timeout_micros=100000,  # 100ms
+          allowed_batch_sizes=[3, 10],
+          batching_queue="",
+          f=computation,
+          in_tensors=[inp],
+          captured_tensors=computation.captured_inputs,
+          Tout=[o.type for o in computation.definition.signature.output_arg])
+
+      thread_results = []
+
+      def worker():
+        thread_results.extend(sess.run([result], feed_dict={inp: [1]}))
+
+      worker_thread = threading.Thread(target=worker)
+      worker_thread.start()
+      main_results = sess.run([result], feed_dict={inp: [2]})
+      worker_thread.join()
+      self.assertEqual(thread_results[0], [2])
+      self.assertEqual(main_results[0], [3])
+
+  def testBatchFunctionOpWithInputError(self):
+    """Tests that batch_function op works with error in the inputs."""
+    if context.executing_eagerly():
+      return
+    with self.cached_session() as sess:
+      inp = array_ops.placeholder(dtype=dtypes.int32, shape=[1])
+
+      @function.Defun(dtypes.int32, dtypes.int32)
+      def computation(in0, in1):
+        return in0 + in1
+
+      result = gen_batch_ops.batch_function(
+          [inp],  # computation actually expects 2 inputs.
+          num_batch_threads=1,
+          max_batch_size=10,
+          batch_timeout_micros=100000,  # 100ms
+          batching_queue="",
+          f=computation,
+          captured_tensors=computation.captured_inputs,
+          Tout=[o.type for o in computation.definition.signature.output_arg])
+
+      with self.assertRaisesRegex(
+          InvalidArgumentError,
+          r"Function takes 2 argument\(s\) but 1 argument\(s\) were passed"):
+        sess.run([result], feed_dict={inp: [2]})
+
+  def testBatchFunctionOpWithLargeBatchSplitted(self):
+    """Tests that the batch_function op works with large batch splitted."""
+    if context.executing_eagerly():
+      return
+
+    with self.cached_session() as sess:
+
+      @function.Defun(dtypes.int32)
+      def computation(in_t):
+        return in_t + 3
+
+      inp = array_ops.placeholder(dtype=dtypes.int32)
+      result = gen_batch_ops.batch_function(
+          [inp],
+          num_batch_threads=2,
+          # enable_large_batch_splitting is True, so it's valid as long as
+          # max('allowed_batch_sizes') <= 'max_batch_size'.
+          allowed_batch_sizes=[1, 2],
+          max_batch_size=5,
+          batch_timeout_micros=100000,  # 100ms
+          Tout=[dtypes.int32],
+          enable_large_batch_splitting=True,
+          f=computation,
+          captured_tensors=computation.captured_inputs)
+      thread1_results = []
+      thread2_results = []
+
+      # Input sizes of worker1 and main thread are larger than
+      # max(allowed_batch_sizes), while input size of worker2 is smaller.
+      def worker1():
+        thread1_results.extend(
+            sess.run([result], feed_dict={inp: [5, 6, 7, 8, 9]}))
+
+      worker_thread1 = threading.Thread(target=worker1)
+      worker_thread1.start()
+
+      def worker2():
+        thread2_results.extend(sess.run([result], feed_dict={inp: [10]}))
+
+      worker_thread2 = threading.Thread(target=worker2)
+      worker_thread2.start()
+
+      main_results = sess.run([result], feed_dict={inp: [2, 3, 4]})
+      worker_thread1.join()
+      worker_thread2.join()
+      self.assertTrue(
+          np.all(np.equal(thread2_results[0], np.array([13], dtype=np.int32))))
+      self.assertTrue(
+          np.all(
+              np.equal(thread1_results[0],
+                       np.array([8, 9, 10, 11, 12], dtype=np.int32))))
+      self.assertTrue(
+          np.all(
+              np.equal(main_results[0], np.array([5, 6, 7], dtype=np.int32))))
+
+  def testBasicUnbatchDecoratedWithReshape(self):
+    """Tests that the batch_function decorator works."""
+    if context.executing_eagerly():
+      return
+    with self.cached_session() as sess:
+
+      @batch_ops.batch_function(1, 10, 100000)
+      def computation(in_t):
+        return array_ops.reshape(in_t, [-1]) + 1
+
+      inp = array_ops.placeholder(dtype=dtypes.int32, shape=[1, 1])
+      result = computation(inp)
+      thread_results = []
+
+      def worker():
+        thread_results.extend(sess.run([result], feed_dict={inp: [[1]]}))
+
+      worker_thread = threading.Thread(target=worker)
+      worker_thread.start()
+      main_results = sess.run([result], feed_dict={inp: [[2]]})
+      worker_thread.join()
+      self.assertEqual(thread_results[0], [2])
+      self.assertEqual(main_results[0], [3])
+
+  def testUnbatchTimeout(self):
+    """Tests that the unbatch timeout works."""
+    if context.executing_eagerly():
+      return
+    with self.cached_session() as sess:
+      inp = array_ops.placeholder(dtype=dtypes.int32, shape=[1])
+      batched, index, id_t = batch_ops.batch(
+          [inp], num_batch_threads=1, max_batch_size=2,
+          batch_timeout_micros=36000000, grad_timeout_micros=0,
+          batching_queue="")
+      computation = batched[0] + 1
+      timeout_micros = 10
+      result = batch_ops.unbatch(computation, index, id_t, timeout_micros,
+                                 shared_name="shared_unbatch")
+      # Set up a parallel pipeline that delays the computation, but uses the
+      # same unbatch resource object as the non-delayed pipeline.
+      computation_delayed = script_ops.py_func(delayed_plus1,
+                                               [batched[0]],
+                                               dtypes.int32)
+      result_delayed = batch_ops.unbatch(computation_delayed,
+                                         index,
+                                         id_t,
+                                         timeout_micros,
+                                         shared_name="shared_unbatch")
+
+      thread_results = []
+      def worker():
+        # A first call using the non-delayed pipeline. The batcher will send an
+        # empty tensor along the non-delayed pipeline.
+        thread_results.extend(sess.run([result], feed_dict={inp: [1]}))
+      worker_thread = threading.Thread(target=worker)
+      worker_thread.start()
+      time.sleep(0.1)  # Ensure the thread's call starts first.
+      # A second call using the delayed pipeline.  The batcher will send the
+      # batched tensor along the delayed pipeline, thus delaying the arrival of
+      # the batched tensor at the unbatch op, relative to the empty tensor.
+      #
+      # TODO(olston, apassos): Avoid relying on the order in which the batch op
+      # emits the empty tensor versus the batched one.
+      _ = sess.run([result_delayed], feed_dict={inp: [2]})
+      worker_thread.join()
+      # The thread's call should hit the timeout, and thus get 0 results.
+      self.assertEqual(len(thread_results), 0)
+
+  def testUnbatchGradInvalidId(self):
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.evaluate(
+          gen_batch_ops.unbatch_grad(
+              original_input=constant_op.constant([1]),
+              batch_index=constant_op.constant([
+                  [0, 0, 0],
+              ], dtype=dtypes.int64),
+              grad=constant_op.constant([
+                  1,
+              ]),
+              id=constant_op.constant([
+                  1,
+                  1,
+              ], dtype=dtypes.int64)))
+
+  def testUnbatchGradInvalidBatchId(self):
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.evaluate(
+          gen_batch_ops.unbatch_grad(
+              original_input=constant_op.constant([1]),
+              batch_index=constant_op.constant([
+                  [0, 0],
+              ], dtype=dtypes.int64),
+              grad=constant_op.constant([
+                  1,
+              ]),
+              id=constant_op.constant([
+                  1,
+              ], dtype=dtypes.int64)))
+
+  def testUnbatchGradInvalidArgs(self):
+    original_input = random_ops.random_uniform(
+        shape=(3, 1), dtype=dtypes.float64, maxval=None)
+    batch_index = random_ops.random_uniform(
+        shape=(3, 1), dtype=dtypes.int64, maxval=65536)
+    grad = random_ops.random_uniform(
+        shape=(3, 1), dtype=dtypes.float64, maxval=None)
+    batch_id = random_ops.random_uniform(
+        shape=(3, 1), dtype=dtypes.int64, maxval=65536)
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.evaluate(
+          gen_batch_ops.unbatch_grad(
+              original_input=original_input,
+              batch_index=batch_index,
+              grad=grad,
+              id=batch_id,
+              container="",
+              shared_name="",
+              name=""))
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/ops/bincount_ops.py

@@ -0,0 +1,235 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# maxlengthations under the License.
+# ==============================================================================
+"""bincount ops."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import tensor_conversion
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_math_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.util import deprecation
+from tensorflow.python.util import dispatch
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("math.bincount", v1=[])
+@dispatch.add_dispatch_support
+def bincount(arr,
+             weights=None,
+             minlength=None,
+             maxlength=None,
+             dtype=dtypes.int32,
+             name=None,
+             axis=None,
+             binary_output=False):
+  """Counts the number of occurrences of each value in an integer array.
+
+  If `minlength` and `maxlength` are not given, returns a vector with length
+  `tf.reduce_max(arr) + 1` if `arr` is non-empty, and length 0 otherwise.
+
+  >>> values = tf.constant([1,1,2,3,2,4,4,5])
+  >>> tf.math.bincount(values)
+  <tf.Tensor: ... numpy=array([0, 2, 2, 1, 2, 1], dtype=int32)>
+
+  Vector length = Maximum element in vector `values` is 5. Adding 1, which is 6
+                  will be the vector length.
+
+  Each bin value in the output indicates number of occurrences of the particular
+  index. Here, index 1 in output has a value 2. This indicates value 1 occurs
+  two times in `values`.
+
+  **Bin-counting with weights**
+
+  >>> values = tf.constant([1,1,2,3,2,4,4,5])
+  >>> weights = tf.constant([1,5,0,1,0,5,4,5])
+  >>> tf.math.bincount(values, weights=weights)
+  <tf.Tensor: ... numpy=array([0, 6, 0, 1, 9, 5], dtype=int32)>
+
+  When `weights` is specified, bins will be incremented by the corresponding
+  weight instead of 1. Here, index 1 in output has a value 6. This is the
+  summation of `weights` corresponding to the value in `values` (i.e. for index
+  1, the first two values are 1 so the first two weights, 1 and 5, are
+  summed).
+
+  There is an equivilance between bin-counting with weights and
+  `unsorted_segement_sum` where `data` is the weights and `segment_ids` are the
+  values.
+
+  >>> values = tf.constant([1,1,2,3,2,4,4,5])
+  >>> weights = tf.constant([1,5,0,1,0,5,4,5])
+  >>> tf.math.unsorted_segment_sum(weights, values, num_segments=6).numpy()
+  array([0, 6, 0, 1, 9, 5], dtype=int32)
+
+  On GPU, `bincount` with weights is only supported when XLA is enabled
+  (typically when a function decorated with `@tf.function(jit_compile=True)`).
+  `unsorted_segment_sum` can be used as a workaround for the non-XLA case on
+  GPU.
+
+  **Bin-counting matrix rows independently**
+
+  This example uses `axis=-1` with a 2 dimensional input and returns a
+  `Tensor` with bincounting where axis 0 is **not** flattened, i.e. an
+  independent bincount for each matrix row.
+
+  >>> data = np.array([[1, 2, 3, 0], [0, 0, 1, 2]], dtype=np.int32)
+  >>> tf.math.bincount(data, axis=-1)
+  <tf.Tensor: shape=(2, 4), dtype=int32, numpy=
+    array([[1, 1, 1, 1],
+           [2, 1, 1, 0]], dtype=int32)>
+
+  **Bin-counting with binary_output**
+
+  This example gives binary output instead of counting the occurrence.
+
+  >>> data = np.array([[1, 2, 3, 0], [0, 0, 1, 2]], dtype=np.int32)
+  >>> tf.math.bincount(data, axis=-1, binary_output=True)
+  <tf.Tensor: shape=(2, 4), dtype=int32, numpy=
+    array([[1, 1, 1, 1],
+           [1, 1, 1, 0]], dtype=int32)>
+
+  **Missing zeros in SparseTensor**
+
+  Note that missing zeros (implict zeros) in SparseTensor are **NOT** counted.
+  This supports cases such as `0` in the values tensor indicates that index/id
+  `0`is present and a missing zero indicates that no index/id is present.
+
+  If counting missing zeros is desired, there are workarounds.
+  For the `axis=0` case, the number of missing zeros can computed by subtracting
+  the number of elements in the SparseTensor's `values` tensor from the
+  number of elements in the dense shape, and this difference can be added to the
+  first element of the output of `bincount`. For all cases, the SparseTensor
+  can be converted to a dense Tensor with `tf.sparse.to_dense` before calling
+  `tf.math.bincount`.
+
+  Args:
+    arr: A Tensor, RaggedTensor, or SparseTensor whose values should be counted.
+      These tensors must have a rank of 2 if `axis=-1`.
+    weights: If non-None, must be the same shape as arr. For each value in
+      `arr`, the bin will be incremented by the corresponding weight instead of
+      1. If non-None, `binary_output` must be False.
+    minlength: If given, ensures the output has length at least `minlength`,
+      padding with zeros at the end if necessary.
+    maxlength: If given, skips values in `arr` that are equal or greater than
+      `maxlength`, ensuring that the output has length at most `maxlength`.
+    dtype: If `weights` is None, determines the type of the output bins.
+    name: A name scope for the associated operations (optional).
+    axis: The axis to slice over. Axes at and below `axis` will be flattened
+      before bin counting. Currently, only `0`, and `-1` are supported. If None,
+      all axes will be flattened (identical to passing `0`).
+    binary_output: If True, this op will output 1 instead of the number of times
+      a token appears (equivalent to one_hot + reduce_any instead of one_hot +
+      reduce_add). Defaults to False.
+
+  Returns:
+    A vector with the same dtype as `weights` or the given `dtype` containing
+    the bincount values.
+
+  Raises:
+    `InvalidArgumentError` if negative values are provided as an input.
+
+  """
+  name = "bincount" if name is None else name
+  with ops.name_scope(name):
+    arr = tensor_conversion.convert_to_tensor_v2_with_dispatch(arr, name="arr")
+    if weights is not None:
+      weights = tensor_conversion.convert_to_tensor_v2_with_dispatch(
+          weights, name="weights")
+
+    if weights is not None and binary_output:
+      raise ValueError("Arguments `binary_output` and `weights` are mutually "
+                       "exclusive. Please specify only one.")
+
+    if not arr.dtype.is_integer:
+      arr = math_ops.cast(arr, dtypes.int32)
+    if axis is None:
+      axis = 0
+
+    if axis not in [0, -1]:
+      raise ValueError(f"Unsupported value for argument axis={axis}. Only 0 and"
+                       " -1 are currently supported.")
+
+    array_is_nonempty = array_ops.size(arr) > 0
+    output_size = math_ops.cast(array_is_nonempty, arr.dtype) * (
+        math_ops.reduce_max(arr) + 1)
+    if minlength is not None:
+      minlength = ops.convert_to_tensor(
+          minlength, name="minlength", dtype=arr.dtype)
+      output_size = gen_math_ops.maximum(minlength, output_size)
+    if maxlength is not None:
+      maxlength = ops.convert_to_tensor(
+          maxlength, name="maxlength", dtype=arr.dtype)
+      output_size = gen_math_ops.minimum(maxlength, output_size)
+
+    if axis == 0:
+      if weights is not None:
+        weights = array_ops.reshape(weights, [-1])
+      arr = array_ops.reshape(arr, [-1])
+
+    weights = validate_dense_weights(arr, weights, dtype)
+    return gen_math_ops.dense_bincount(
+        input=arr,
+        size=output_size,
+        weights=weights,
+        binary_output=binary_output)
+
+
+@tf_export(v1=["math.bincount", "bincount"])
+@deprecation.deprecated_endpoints("bincount")
+def bincount_v1(arr,
+                weights=None,
+                minlength=None,
+                maxlength=None,
+                dtype=dtypes.int32):
+  """Counts the number of occurrences of each value in an integer array.
+
+  If `minlength` and `maxlength` are not given, returns a vector with length
+  `tf.reduce_max(arr) + 1` if `arr` is non-empty, and length 0 otherwise.
+  If `weights` are non-None, then index `i` of the output stores the sum of the
+  value in `weights` at each index where the corresponding value in `arr` is
+  `i`.
+
+  Args:
+    arr: An int32 tensor of non-negative values.
+    weights: If non-None, must be the same shape as arr. For each value in
+      `arr`, the bin will be incremented by the corresponding weight instead of
+      1.
+    minlength: If given, ensures the output has length at least `minlength`,
+      padding with zeros at the end if necessary.
+    maxlength: If given, skips values in `arr` that are equal or greater than
+      `maxlength`, ensuring that the output has length at most `maxlength`.
+    dtype: If `weights` is None, determines the type of the output bins.
+
+  Returns:
+    A vector with the same dtype as `weights` or the given `dtype`. The bin
+    values.
+  """
+  return bincount(arr, weights, minlength, maxlength, dtype)
+
+
+def validate_dense_weights(values, weights, dtype=None):
+  """Validates the passed weight tensor or creates an empty one."""
+  if weights is None:
+    if dtype:
+      return array_ops.constant([], dtype=dtype)
+    return array_ops.constant([], dtype=values.dtype)
+
+  if not isinstance(weights, tensor.Tensor):
+    raise ValueError(
+        "Argument `weights` must be a tf.Tensor if `values` is a tf.Tensor. "
+        f"Received weights={weights} of type: {type(weights).__name__}")
+
+  return weights

tensorflow/python/ops/bincount_ops_test.py

@@ -0,0 +1,569 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# maxlengthations under the License.
+# ==============================================================================
+"""Tests for bincount ops."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.framework import config as tf_config
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import bincount_ops
+from tensorflow.python.ops import gen_count_ops
+from tensorflow.python.ops import sparse_ops
+from tensorflow.python.platform import test
+
+
+def _adjust_expected_rank1(x, minlength, maxlength):
+  """Trim or pad an expected result based on minlength and maxlength."""
+  n = len(x)
+  if (minlength is not None) and (n < minlength):
+    x = x + [0] * (minlength - n)
+  if (maxlength is not None) and (n > maxlength):
+    x = x[:maxlength]
+  return x
+
+
+def _adjust_expected_rank2(x, minlength, maxlength):
+  return [_adjust_expected_rank1(i, minlength, maxlength) for i in x]
+
+
+class TestDenseBincount(test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters([{
+      "dtype": np.int32,
+  }, {
+      "dtype": np.int64,
+  }])
+  def test_sparse_input_all_count(self, dtype):
+    np.random.seed(42)
+    num_rows = 4096
+    size = 1000
+    n_elems = 128
+    inp_indices = np.random.randint(0, num_rows, (n_elems, 1))
+    inp_indices = np.concatenate([inp_indices, np.zeros((n_elems, 1))], axis=1)
+    inp_vals = np.random.randint(0, size, (n_elems,), dtype=dtype)
+    sparse_inp = sparse_tensor.SparseTensor(inp_indices, inp_vals,
+                                            [num_rows, 1])
+
+    # Note that the default for sparse tensors is to not count implicit zeros.
+    np_out = np.bincount(inp_vals, minlength=size)
+    self.assertAllEqual(
+        np_out,
+        self.evaluate(
+            bincount_ops.bincount(sparse_inp, axis=0, minlength=size)
+        ),
+    )
+
+  @parameterized.parameters([{
+      "dtype": np.int32,
+  }, {
+      "dtype": np.int64,
+  }])
+  def test_sparse_input_all_count_with_weights(self, dtype):
+    np.random.seed(42)
+    num_rows = 4096
+    size = 1000
+    n_elems = 128
+    inp_indices = np.random.randint(0, num_rows, (n_elems, 1))
+    inp_indices = np.concatenate([inp_indices, np.zeros((n_elems, 1))], axis=1)
+    inp_vals = np.random.randint(0, size, (n_elems-1,), dtype=dtype)
+    # Add an element with value `size-1` to input so bincount output has `size`
+    # elements.
+    inp_vals = np.concatenate([inp_vals, [size-1]], axis=0)
+    sparse_inp = sparse_tensor.SparseTensor(inp_indices, inp_vals,
+                                            [num_rows, 1])
+    weight_vals = np.random.random((n_elems,))
+    sparse_weights = sparse_tensor.SparseTensor(inp_indices, weight_vals,
+                                                [num_rows, 1])
+
+    np_out = np.bincount(inp_vals, minlength=size, weights=weight_vals)
+    self.assertAllEqual(
+        np_out,
+        self.evaluate(bincount_ops.bincount(
+            sparse_inp, sparse_weights, axis=0)))
+
+  @parameterized.parameters([{
+      "dtype": np.int32,
+  }, {
+      "dtype": np.int64,
+  }])
+  def test_sparse_input_all_binary(self, dtype):
+    np.random.seed(42)
+    num_rows = 4096
+    size = 10
+    n_elems = 128
+    inp_indices = np.random.randint(0, num_rows, (n_elems, 1))
+    inp_indices = np.concatenate([inp_indices, np.zeros((n_elems, 1))], axis=1)
+    inp_vals = np.random.randint(0, size, (n_elems,), dtype=dtype)
+    sparse_inp = sparse_tensor.SparseTensor(inp_indices, inp_vals,
+                                            [num_rows, 1])
+
+    np_out = np.ones((size,))
+    self.assertAllEqual(
+        np_out,
+        self.evaluate(bincount_ops.bincount(sparse_inp, binary_output=True)))
+
+  @parameterized.parameters([{
+      "dtype": np.int32,
+  }, {
+      "dtype": np.int64,
+  }])
+  def test_sparse_input_col_reduce_count(self, dtype):
+    num_rows = 128
+    num_cols = 27
+    size = 100
+    np.random.seed(42)
+    inp = np.random.randint(0, size, (num_rows, num_cols), dtype=dtype)
+    np_out = np.reshape(
+        np.concatenate(
+            [np.bincount(inp[j, :], minlength=size) for j in range(num_rows)],
+            axis=0), (num_rows, size))
+    # from_dense will filter out 0s.
+    inp = inp + 1
+    # from_dense will cause OOM in GPU.
+    with ops.device("/CPU:0"):
+      inp_sparse = sparse_ops.from_dense(inp)
+      inp_sparse = sparse_tensor.SparseTensor(inp_sparse.indices,
+                                              inp_sparse.values - 1,
+                                              inp_sparse.dense_shape)
+    self.assertAllEqual(
+        np_out, self.evaluate(bincount_ops.bincount(arr=inp_sparse, axis=-1)))
+
+  @parameterized.parameters([{
+      "dtype": np.int32,
+  }, {
+      "dtype": np.int64,
+  }])
+  def test_sparse_input_col_reduce_binary(self, dtype):
+    num_rows = 128
+    num_cols = 27
+    size = 100
+    np.random.seed(42)
+    inp = np.random.randint(0, size, (num_rows, num_cols), dtype=dtype)
+    np_out = np.reshape(
+        np.concatenate([
+            np.where(np.bincount(inp[j, :], minlength=size) > 0, 1, 0)
+            for j in range(num_rows)
+        ],
+                       axis=0), (num_rows, size))
+    # from_dense will filter out 0s.
+    inp = inp + 1
+    # from_dense will cause OOM in GPU.
+    with ops.device("/CPU:0"):
+      inp_sparse = sparse_ops.from_dense(inp)
+      inp_sparse = sparse_tensor.SparseTensor(inp_sparse.indices,
+                                              inp_sparse.values - 1,
+                                              inp_sparse.dense_shape)
+    self.assertAllEqual(
+        np_out,
+        self.evaluate(
+            bincount_ops.bincount(arr=inp_sparse, axis=-1, binary_output=True)))
+
+  @parameterized.product(
+      (
+          dict(
+              tid="_d1",
+              x=[1, 2, 2, 3, 3, 3],
+              expected=[0, 1, 2, 3],
+          ),
+          dict(
+              tid="_d2",
+              x=[[0, 0, 0], [0, 1, 0], [2, 0, 2], [3, 3, 3]],
+              expected=[6, 1, 2, 3],
+          ),
+          dict(
+              tid="_d3",
+              x=[[[0, 0, 0], [0, 1, 0]], [[2, 0, 2], [3, 3, 3]]],
+              expected=[6, 1, 2, 3],
+          ),
+      ),
+      (
+          dict(minlength=None, maxlength=None),
+          dict(minlength=3, maxlength=None),
+          dict(minlength=5, maxlength=None),
+          dict(minlength=None, maxlength=3),
+          dict(minlength=None, maxlength=5),
+          dict(minlength=2, maxlength=3),
+          dict(minlength=3, maxlength=5),
+          dict(minlength=5, maxlength=10),
+          dict(minlength=None, maxlength=0),
+      ),
+  )
+  def test_default(
+      self,
+      x,
+      minlength,
+      maxlength,
+      expected,
+      tid=None,
+  ):
+    expected = _adjust_expected_rank1(expected, minlength, maxlength)
+    self.assertAllEqual(
+        expected,
+        self.evaluate(
+            bincount_ops.bincount(x, minlength=minlength, maxlength=maxlength)
+        ),
+    )
+    self.assertAllEqual(
+        expected,
+        self.evaluate(
+            bincount_ops.bincount(
+                x, minlength=minlength, maxlength=maxlength, axis=0
+            )
+        ),
+    )
+
+  @parameterized.product(
+      (
+          dict(
+              tid="_d2",
+              x=[[0, 0, 0], [0, 1, 0], [2, 0, 2], [3, 3, 3]],
+              expected=[[3, 0, 0, 0], [2, 1, 0, 0], [1, 0, 2, 0], [0, 0, 0, 3]],
+          ),
+      ),
+      (
+          dict(minlength=None, maxlength=None),
+          dict(minlength=3, maxlength=None),
+          dict(minlength=5, maxlength=None),
+          dict(minlength=None, maxlength=3),
+          dict(minlength=None, maxlength=5),
+          dict(minlength=2, maxlength=3),
+          dict(minlength=3, maxlength=5),
+          dict(minlength=5, maxlength=10),
+          dict(minlength=None, maxlength=0),
+      ),
+  )
+  def test_axis_neg_1(
+      self, tid, x, minlength, maxlength, expected
+  ):
+    expected = _adjust_expected_rank2(expected, minlength, maxlength)
+    self.assertAllEqual(
+        expected,
+        self.evaluate(
+            bincount_ops.bincount(
+                x, minlength=minlength, maxlength=maxlength, axis=-1
+            )
+        ),
+    )
+
+  @parameterized.product(
+      (
+          dict(
+              tid="_d1",
+              x=[1, 2, 2, 3, 3, 3],
+              weights=[1, 2, 3, 4, 5, 6],
+              axis=None,
+              expected=[0, 1, 5, 15],
+          ),
+          dict(
+              tid="_d2",
+              x=[[0, 0, 0], [0, 1, 0], [2, 0, 2], [3, 3, 3]],
+              weights=[[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]],
+              axis=None,
+              expected=[24, 5, 16, 33],
+          ),
+          dict(
+              tid="_d3",
+              x=[[[0, 0, 0], [0, 1, 0]], [[2, 0, 2], [3, 3, 3]]],
+              weights=[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]],
+              axis=None,
+              expected=[24, 5, 16, 33],
+          ),
+          dict(
+              tid="_d2_axis_neg_1",
+              x=[[0, 0, 0], [0, 1, 0], [2, 0, 2], [3, 3, 3]],
+              weights=[[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]],
+              axis=-1,
+              expected=[
+                  [6, 0, 0, 0],
+                  [10, 5, 0, 0],
+                  [8, 0, 16, 0],
+                  [0, 0, 0, 33],
+              ],
+          ),
+      ),
+      (
+          dict(minlength=None, maxlength=None),
+          dict(minlength=3, maxlength=None),
+          dict(minlength=5, maxlength=None),
+          dict(minlength=None, maxlength=3),
+          dict(minlength=None, maxlength=5),
+          dict(minlength=2, maxlength=3),
+          dict(minlength=3, maxlength=5),
+          dict(minlength=5, maxlength=10),
+          dict(minlength=None, maxlength=0),
+      ),
+  )
+  def test_weights(
+      self,
+      tid,
+      x,
+      weights,
+      minlength,
+      maxlength,
+      expected,
+      axis=None,
+  ):
+    device_set = set([d.device_type for d in tf_config.list_physical_devices()])
+    if "GPU" in device_set and not test_util.is_xla_enabled():
+      self.skipTest(
+          "b/263004039 The DenseBincount GPU kernel does not support weights."
+          " unsorted_segment_sum should be used instead on GPU."
+      )
+    if axis == -1:
+      expected = _adjust_expected_rank2(expected, minlength, maxlength)
+    else:
+      expected = _adjust_expected_rank1(expected, minlength, maxlength)
+    self.assertAllEqual(
+        expected,
+        self.evaluate(
+            bincount_ops.bincount(
+                x,
+                weights=weights,
+                minlength=minlength,
+                maxlength=maxlength,
+                axis=axis,
+            )
+        ),
+    )
+
+  @parameterized.product(
+      (
+          dict(
+              tid="_d1",
+              x=[1, 2, 2, 3, 3, 3],
+              expected=[0, 1, 1, 1],
+              axis=None,
+          ),
+          dict(
+              tid="_d2",
+              x=[[0, 0, 0], [0, 1, 0], [2, 0, 2], [3, 3, 3]],
+              expected=[1, 1, 1, 1],
+              axis=None,
+          ),
+          dict(
+              tid="_d3",
+              x=[[[0, 0, 0], [0, 1, 0]], [[2, 0, 2], [3, 3, 3]]],
+              expected=[1, 1, 1, 1],
+              axis=None,
+          ),
+          dict(
+              tid="_d2_axis_neg_1",
+              x=[[0, 0, 0], [0, 1, 0], [2, 0, 2], [3, 3, 3]],
+              expected=[[1, 0, 0, 0], [1, 1, 0, 0], [1, 0, 1, 0], [0, 0, 0, 1]],
+              axis=-1,
+          ),
+      ),
+      (
+          dict(minlength=None, maxlength=None),
+          dict(minlength=3, maxlength=None),
+          dict(minlength=5, maxlength=None),
+          dict(minlength=None, maxlength=3),
+          dict(minlength=None, maxlength=5),
+          dict(minlength=2, maxlength=3),
+          dict(minlength=3, maxlength=5),
+          dict(minlength=5, maxlength=10),
+          dict(minlength=None, maxlength=0),
+      ),
+  )
+  def test_binary_output(
+      self,
+      tid,
+      x,
+      minlength,
+      maxlength,
+      expected,
+      axis=None,
+  ):
+    if axis == -1:
+      expected = _adjust_expected_rank2(expected, minlength, maxlength)
+    else:
+      expected = _adjust_expected_rank1(expected, minlength, maxlength)
+    self.assertAllEqual(
+        expected,
+        self.evaluate(
+            bincount_ops.bincount(
+                x,
+                minlength=minlength,
+                maxlength=maxlength,
+                binary_output=True,
+                axis=axis,
+            )
+        ),
+    )
+
+
+class RawOpsHeapOobTest(test.TestCase, parameterized.TestCase):
+
+  @test_util.run_v1_only("Test security error")
+  def testSparseCountSparseOutputBadIndicesShapeTooSmall(self):
+    indices = [1]
+    values = [[1]]
+    weights = []
+    dense_shape = [10]
+    with self.assertRaisesRegex(ValueError,
+                                "Shape must be rank 2 but is rank 1 for"):
+      self.evaluate(
+          gen_count_ops.SparseCountSparseOutput(
+              indices=indices,
+              values=values,
+              dense_shape=dense_shape,
+              weights=weights,
+              binary_output=True))
+
+
+@test_util.run_all_in_graph_and_eager_modes
+@test_util.disable_tfrt
+class RawOpsTest(test.TestCase, parameterized.TestCase):
+
+  def testSparseCountSparseOutputBadIndicesShape(self):
+    indices = [[[0], [0]], [[0], [1]], [[1], [0]], [[1], [2]]]
+    values = [1, 1, 1, 10]
+    weights = [1, 2, 4, 6]
+    dense_shape = [2, 3]
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "Input indices must be a 2-dimensional tensor"):
+      self.evaluate(
+          gen_count_ops.SparseCountSparseOutput(
+              indices=indices,
+              values=values,
+              dense_shape=dense_shape,
+              weights=weights,
+              binary_output=False))
+
+  def testSparseCountSparseOutputBadWeightsShape(self):
+    indices = [[0, 0], [0, 1], [1, 0], [1, 2]]
+    values = [1, 1, 1, 10]
+    weights = [1, 2, 4]
+    dense_shape = [2, 3]
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "Weights and values must have the same shape"):
+      self.evaluate(
+          gen_count_ops.SparseCountSparseOutput(
+              indices=indices,
+              values=values,
+              dense_shape=dense_shape,
+              weights=weights,
+              binary_output=False))
+
+  def testSparseCountSparseOutputBadNumberOfValues(self):
+    indices = [[0, 0], [0, 1], [1, 0]]
+    values = [1, 1, 1, 10]
+    weights = [1, 2, 4, 6]
+    dense_shape = [2, 3]
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Number of values must match first dimension of indices"):
+      self.evaluate(
+          gen_count_ops.SparseCountSparseOutput(
+              indices=indices,
+              values=values,
+              dense_shape=dense_shape,
+              weights=weights,
+              binary_output=False))
+
+  def testSparseCountSparseOutputNegativeValue(self):
+    indices = [[0, 0], [0, 1], [1, 0], [1, 2]]
+    values = [1, 1, -1, 10]
+    dense_shape = [2, 3]
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "Input values must all be non-negative"):
+      self.evaluate(
+          gen_count_ops.SparseCountSparseOutput(
+              indices=indices,
+              values=values,
+              dense_shape=dense_shape,
+              binary_output=False))
+
+  def testRaggedCountSparseOutput(self):
+    splits = [0, 4, 7]
+    values = [1, 1, 2, 1, 2, 10, 5]
+    weights = [1, 2, 3, 4, 5, 6, 7]
+    output_indices, output_values, output_shape = self.evaluate(
+        gen_count_ops.RaggedCountSparseOutput(
+            splits=splits, values=values, weights=weights, binary_output=False))
+    self.assertAllEqual([[0, 1], [0, 2], [1, 2], [1, 5], [1, 10]],
+                        output_indices)
+    self.assertAllEqual([7, 3, 5, 7, 6], output_values)
+    self.assertAllEqual([2, 11], output_shape)
+
+  def testRaggedCountSparseOutputBadWeightsShape(self):
+    splits = [0, 4, 7]
+    values = [1, 1, 2, 1, 2, 10, 5]
+    weights = [1, 2, 3, 4, 5, 6]
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "Weights and values must have the same shape"):
+      self.evaluate(
+          gen_count_ops.RaggedCountSparseOutput(
+              splits=splits,
+              values=values,
+              weights=weights,
+              binary_output=False))
+
+  def testRaggedCountSparseOutputEmptySplits(self):
+    splits = []
+    values = [1, 1, 2, 1, 2, 10, 5]
+    weights = [1, 2, 3, 4, 5, 6, 7]
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Must provide at least 2 elements for the splits argument"):
+      self.evaluate(
+          gen_count_ops.RaggedCountSparseOutput(
+              splits=splits,
+              values=values,
+              weights=weights,
+              binary_output=False))
+
+  def testRaggedCountSparseOutputBadSplitsStart(self):
+    splits = [1, 7]
+    values = [1, 1, 2, 1, 2, 10, 5]
+    weights = [1, 2, 3, 4, 5, 6, 7]
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "Splits must start with 0"):
+      self.evaluate(
+          gen_count_ops.RaggedCountSparseOutput(
+              splits=splits,
+              values=values,
+              weights=weights,
+              binary_output=False))
+
+  def testRaggedCountSparseOutputBadSplitsEnd(self):
+    splits = [0, 5]
+    values = [1, 1, 2, 1, 2, 10, 5]
+    weights = [1, 2, 3, 4, 5, 6, 7]
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "Splits must end with the number of values"):
+      self.evaluate(
+          gen_count_ops.RaggedCountSparseOutput(
+              splits=splits,
+              values=values,
+              weights=weights,
+              binary_output=False))
+
+  def testRaggedCountSparseOutputNegativeValue(self):
+    splits = [0, 4, 7]
+    values = [1, 1, 2, 1, -2, 10, 5]
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "Input values must all be non-negative"):
+      self.evaluate(
+          gen_count_ops.RaggedCountSparseOutput(
+              splits=splits, values=values, binary_output=False))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/ops/bitwise_ops.py

@@ -0,0 +1,33 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Operations for manipulating the binary representations of integers.
+
+API docstring: tensorflow.bitwise
+"""
+
+from tensorflow.python.framework import ops
+# go/tf-wildcard-import
+# pylint: disable=wildcard-import
+from tensorflow.python.ops.gen_bitwise_ops import *
+# pylint: enable=wildcard-import
+
+ops.NotDifferentiable("BitwiseAnd")
+ops.NotDifferentiable("BitwiseOr")
+ops.NotDifferentiable("BitwiseXor")
+ops.NotDifferentiable("Invert")
+ops.NotDifferentiable("PopulationCount")
+ops.NotDifferentiable("LeftShift")
+ops.NotDifferentiable("RightShift")

tensorflow/python/ops/bitwise_ops_test.py

@@ -0,0 +1,185 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for bitwise operations."""
+
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import bitwise_ops
+from tensorflow.python.ops import gen_bitwise_ops
+from tensorflow.python.platform import googletest
+
+
+class BitwiseOpTest(test_util.TensorFlowTestCase):
+
+  def __init__(self, method_name="runTest"):
+    super(BitwiseOpTest, self).__init__(method_name)
+
+  @test_util.run_deprecated_v1
+  def testBinaryOps(self):
+    dtype_list = [dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64,
+                  dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64]
+
+    with self.session() as sess:
+      for dtype in dtype_list:
+        lhs = constant_op.constant([0, 5, 3, 14], dtype=dtype)
+        rhs = constant_op.constant([5, 0, 7, 11], dtype=dtype)
+        and_result, or_result, xor_result = sess.run(
+            [bitwise_ops.bitwise_and(lhs, rhs),
+             bitwise_ops.bitwise_or(lhs, rhs),
+             bitwise_ops.bitwise_xor(lhs, rhs)])
+        self.assertAllEqual(and_result, [0, 0, 3, 10])
+        self.assertAllEqual(or_result, [5, 5, 7, 15])
+        self.assertAllEqual(xor_result, [5, 5, 4, 5])
+
+  def testPopulationCountOp(self):
+    dtype_list = [
+        dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64, dtypes.uint8,
+        dtypes.uint16, dtypes.uint32, dtypes.uint64
+    ]
+    raw_inputs = [0, 1, -1, 3, -3, 5, -5, 14, -14,
+                  127, 128, 255, 256, 65535, 65536,
+                  2**31 - 1, 2**31, 2**32 - 1, 2**32, -2**32 + 1, -2**32,
+                  -2**63 + 1, 2**63 - 1]
+    def count_bits(x):
+      return sum(bin(z).count("1") for z in x.tobytes())
+    for dtype in dtype_list:
+      with self.cached_session():
+        print("PopulationCount test: ", dtype)
+        inputs = np.array(raw_inputs).astype(dtype.as_numpy_dtype)
+        truth = [count_bits(x) for x in inputs]
+        input_tensor = constant_op.constant(inputs, dtype=dtype)
+        popcnt_result = self.evaluate(
+            gen_bitwise_ops.population_count(input_tensor))
+        self.assertAllEqual(truth, popcnt_result)
+
+  def testPopulationCountOpEmptyInput(self):
+    with self.cached_session():
+      popcnt_result = self.evaluate(
+          gen_bitwise_ops.population_count(
+              constant_op.constant(
+                  [],
+                  shape=[0],
+                  dtype=dtypes.int64,
+              ),
+          )
+      )
+      self.assertAllEqual(popcnt_result, [])
+
+  @test_util.run_deprecated_v1
+  def testInvertOp(self):
+    dtype_list = [dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64,
+                  dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64]
+    inputs = [0, 5, 3, 14]
+    with self.session() as sess:
+      for dtype in dtype_list:
+        # Because of issues with negative numbers, let's test this indirectly.
+        # 1. invert(a) and a = 0
+        # 2. invert(a) or a = invert(0)
+        input_tensor = constant_op.constant(inputs, dtype=dtype)
+        not_a_and_a, not_a_or_a, not_0 = sess.run(
+            [bitwise_ops.bitwise_and(
+                input_tensor, bitwise_ops.invert(input_tensor)),
+             bitwise_ops.bitwise_or(
+                 input_tensor, bitwise_ops.invert(input_tensor)),
+             bitwise_ops.invert(constant_op.constant(0, dtype=dtype))])
+        self.assertAllEqual(not_a_and_a, [0, 0, 0, 0])
+        self.assertAllEqual(not_a_or_a, [not_0] * 4)
+        # For unsigned dtypes let's also check the result directly.
+        if dtype.is_unsigned:
+          inverted = self.evaluate(bitwise_ops.invert(input_tensor))
+          expected = [dtype.max - x for x in inputs]
+          self.assertAllEqual(inverted, expected)
+
+  @test_util.run_deprecated_v1
+  def testShiftsWithPositiveLHS(self):
+    dtype_list = [np.int8, np.int16, np.int32, np.int64,
+                  np.uint8, np.uint16, np.uint32, np.uint64]
+
+    with self.session() as sess:
+      for dtype in dtype_list:
+        lhs = np.array([0, 5, 3, 14], dtype=dtype)
+        rhs = np.array([5, 0, 7, 3], dtype=dtype)
+        left_shift_result, right_shift_result = sess.run(
+            [bitwise_ops.left_shift(lhs, rhs),
+             bitwise_ops.right_shift(lhs, rhs)])
+        self.assertAllEqual(left_shift_result, np.left_shift(lhs, rhs))
+        self.assertAllEqual(right_shift_result, np.right_shift(lhs, rhs))
+
+  @test_util.run_deprecated_v1
+  def testShiftsWithNegativeLHS(self):
+    dtype_list = [np.int8, np.int16, np.int32, np.int64]
+
+    with self.session() as sess:
+      for dtype in dtype_list:
+        lhs = np.array([-1, -5, -3, -14], dtype=dtype)
+        rhs = np.array([5, 0, 7, 11], dtype=dtype)
+        left_shift_result, right_shift_result = sess.run(
+            [bitwise_ops.left_shift(lhs, rhs),
+             bitwise_ops.right_shift(lhs, rhs)])
+        self.assertAllEqual(left_shift_result, np.left_shift(lhs, rhs))
+        self.assertAllEqual(right_shift_result, np.right_shift(lhs, rhs))
+
+  @test_util.run_deprecated_v1
+  def testImplementationDefinedShiftsDoNotCrash(self):
+    dtype_list = [np.int8, np.int16, np.int32, np.int64]
+
+    with self.session() as sess:
+      for dtype in dtype_list:
+        lhs = np.array([-1, -5, -3, -14], dtype=dtype)
+        rhs = np.array([-2, 64, 101, 32], dtype=dtype)
+        # We intentionally do not test for specific values here since the exact
+        # outputs are implementation-defined. However, we should not crash or
+        # trigger an undefined-behavior error from tools such as
+        # AddressSanitizer.
+        sess.run([bitwise_ops.left_shift(lhs, rhs),
+                  bitwise_ops.right_shift(lhs, rhs)])
+
+  @test_util.run_deprecated_v1
+  def testShapeInference(self):
+    dtype_list = [dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64,
+                  dtypes.uint8, dtypes.uint16]
+
+    with self.session() as sess:
+      for dtype in dtype_list:
+        lhs = constant_op.constant([[0], [3], [5]], dtype=dtype)
+        rhs = constant_op.constant([[1, 2, 4]], dtype=dtype)
+
+        and_tensor = bitwise_ops.bitwise_and(lhs, rhs)
+        or_tensor = bitwise_ops.bitwise_or(lhs, rhs)
+        xor_tensor = bitwise_ops.bitwise_xor(lhs, rhs)
+        ls_tensor = bitwise_ops.left_shift(lhs, rhs)
+        rs_tensor = bitwise_ops.right_shift(lhs, rhs)
+
+        and_result, or_result, xor_result, ls_result, rs_result = sess.run(
+            [and_tensor, or_tensor, xor_tensor, ls_tensor, rs_tensor])
+
+        # Compare shape inference with result
+        self.assertAllEqual(and_tensor.get_shape().as_list(), and_result.shape)
+        self.assertAllEqual(and_tensor.get_shape().as_list(), [3, 3])
+        self.assertAllEqual(or_tensor.get_shape().as_list(), or_result.shape)
+        self.assertAllEqual(or_tensor.get_shape().as_list(), [3, 3])
+        self.assertAllEqual(xor_tensor.get_shape().as_list(), xor_result.shape)
+        self.assertAllEqual(xor_tensor.get_shape().as_list(), [3, 3])
+        self.assertAllEqual(ls_tensor.get_shape().as_list(), ls_result.shape)
+        self.assertAllEqual(ls_tensor.get_shape().as_list(), [3, 3])
+        self.assertAllEqual(rs_tensor.get_shape().as_list(), rs_result.shape)
+        self.assertAllEqual(rs_tensor.get_shape().as_list(), [3, 3])
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/python/ops/boosted_trees_ops.py

@@ -0,0 +1,311 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Ops for boosted_trees."""
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_boosted_trees_ops
+from tensorflow.python.ops import resources
+
+# Re-exporting ops used by other modules.
+# pylint: disable=unused-import
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_aggregate_stats
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_bucketize
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_calculate_best_feature_split as calculate_best_feature_split
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_calculate_best_feature_split_v2 as calculate_best_feature_split_v2
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_calculate_best_gains_per_feature as calculate_best_gains_per_feature
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_center_bias as center_bias
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_create_quantile_stream_resource as create_quantile_stream_resource
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_example_debug_outputs as example_debug_outputs
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_make_quantile_summaries as make_quantile_summaries
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_make_stats_summary as make_stats_summary
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_predict as predict
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_quantile_stream_resource_add_summaries as quantile_add_summaries
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_quantile_stream_resource_deserialize as quantile_resource_deserialize
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_quantile_stream_resource_flush as quantile_flush
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_quantile_stream_resource_get_bucket_boundaries as get_bucket_boundaries
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_quantile_stream_resource_handle_op as quantile_resource_handle_op
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_sparse_aggregate_stats
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_sparse_calculate_best_feature_split as sparse_calculate_best_feature_split
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_training_predict as training_predict
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_update_ensemble as update_ensemble
+from tensorflow.python.ops.gen_boosted_trees_ops import boosted_trees_update_ensemble_v2 as update_ensemble_v2
+from tensorflow.python.ops.gen_boosted_trees_ops import is_boosted_trees_quantile_stream_resource_initialized as is_quantile_resource_initialized
+# pylint: enable=unused-import
+
+from tensorflow.python.training import saver
+
+
+class PruningMode:
+  """Class for working with Pruning modes."""
+  NO_PRUNING, PRE_PRUNING, POST_PRUNING = range(0, 3)
+
+  _map = {'none': NO_PRUNING, 'pre': PRE_PRUNING, 'post': POST_PRUNING}
+
+  @classmethod
+  def from_str(cls, mode):
+    if mode in cls._map:
+      return cls._map[mode]
+    else:
+      raise ValueError(
+          'pruning_mode mode must be one of: {}. Found: {}'.format(', '.join(
+              sorted(cls._map)), mode))
+
+
+class QuantileAccumulatorSaveable(saver.BaseSaverBuilder.SaveableObject):
+  """SaveableObject implementation for QuantileAccumulator."""
+
+  def __init__(self, resource_handle, create_op, num_streams, name):
+    self.resource_handle = resource_handle
+    self._num_streams = num_streams
+    self._create_op = create_op
+    bucket_boundaries = get_bucket_boundaries(self.resource_handle,
+                                              self._num_streams)
+    slice_spec = ''
+    specs = []
+
+    def make_save_spec(tensor, suffix):
+      return saver.BaseSaverBuilder.SaveSpec(tensor, slice_spec, name + suffix)
+
+    for i in range(self._num_streams):
+      specs += [
+          make_save_spec(bucket_boundaries[i], '_bucket_boundaries_' + str(i))
+      ]
+    super(QuantileAccumulatorSaveable, self).__init__(self.resource_handle,
+                                                      specs, name)
+
+  def restore(self, restored_tensors, unused_tensor_shapes):
+    bucket_boundaries = restored_tensors
+    with ops.control_dependencies([self._create_op]):
+      return quantile_resource_deserialize(
+          self.resource_handle, bucket_boundaries=bucket_boundaries)
+
+
+class QuantileAccumulator():
+  """SaveableObject implementation for QuantileAccumulator.
+
+     The bucket boundaries are serialized and deserialized from checkpointing.
+  """
+
+  def __init__(self,
+               epsilon,
+               num_streams,
+               num_quantiles,
+               name=None,
+               max_elements=None):
+    del max_elements  # Unused.
+
+    self._eps = epsilon
+    self._num_streams = num_streams
+    self._num_quantiles = num_quantiles
+
+    with ops.name_scope(name, 'QuantileAccumulator') as name:
+      self._name = name
+      self.resource_handle = self._create_resource()
+      self._init_op = self._initialize()
+      is_initialized_op = self.is_initialized()
+    resources.register_resource(self.resource_handle, self._init_op,
+                                is_initialized_op)
+    ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS,
+                          QuantileAccumulatorSaveable(
+                              self.resource_handle, self._init_op,
+                              self._num_streams, self.resource_handle.name))
+
+  def _create_resource(self):
+    return quantile_resource_handle_op(
+        container='', shared_name=self._name, name=self._name)
+
+  def _initialize(self):
+    return create_quantile_stream_resource(self.resource_handle, self._eps,
+                                           self._num_streams)
+
+  @property
+  def initializer(self):
+    if self._init_op is None:
+      self._init_op = self._initialize()
+    return self._init_op
+
+  def is_initialized(self):
+    return is_quantile_resource_initialized(self.resource_handle)
+
+  def _serialize_to_tensors(self):
+    raise NotImplementedError('When the need arises, TF2 compatibility can be '
+                              'added by implementing this method, along with '
+                              '_restore_from_tensors below.')
+
+  def _restore_from_tensors(self, restored_tensors):
+    raise NotImplementedError('When the need arises, TF2 compatibility can be '
+                              'added by implementing this method, along with '
+                              '_serialize_to_tensors above.')
+
+  def add_summaries(self, float_columns, example_weights):
+    summaries = make_quantile_summaries(float_columns, example_weights,
+                                        self._eps)
+    summary_op = quantile_add_summaries(self.resource_handle, summaries)
+    return summary_op
+
+  def flush(self):
+    return quantile_flush(self.resource_handle, self._num_quantiles)
+
+  def get_bucket_boundaries(self):
+    return get_bucket_boundaries(self.resource_handle, self._num_streams)
+
+
+class _TreeEnsembleSavable(saver.BaseSaverBuilder.SaveableObject):
+  """SaveableObject implementation for TreeEnsemble."""
+
+  def __init__(self, resource_handle, create_op, name):
+    """Creates a _TreeEnsembleSavable object.
+
+    Args:
+      resource_handle: handle to the decision tree ensemble variable.
+      create_op: the op to initialize the variable.
+      name: the name to save the tree ensemble variable under.
+    """
+    stamp_token, serialized = (
+        gen_boosted_trees_ops.boosted_trees_serialize_ensemble(resource_handle))
+    # slice_spec is useful for saving a slice from a variable.
+    # It's not meaningful the tree ensemble variable. So we just pass an empty
+    # value.
+    slice_spec = ''
+    specs = [
+        saver.BaseSaverBuilder.SaveSpec(stamp_token, slice_spec,
+                                        name + '_stamp'),
+        saver.BaseSaverBuilder.SaveSpec(serialized, slice_spec,
+                                        name + '_serialized'),
+    ]
+    super(_TreeEnsembleSavable, self).__init__(resource_handle, specs, name)
+    self.resource_handle = resource_handle
+    self._create_op = create_op
+
+  def restore(self, restored_tensors, unused_restored_shapes):
+    """Restores the associated tree ensemble from 'restored_tensors'.
+
+    Args:
+      restored_tensors: the tensors that were loaded from a checkpoint.
+      unused_restored_shapes: the shapes this object should conform to after
+        restore. Not meaningful for trees.
+
+    Returns:
+      The operation that restores the state of the tree ensemble variable.
+    """
+    with ops.control_dependencies([self._create_op]):
+      return gen_boosted_trees_ops.boosted_trees_deserialize_ensemble(
+          self.resource_handle,
+          stamp_token=restored_tensors[0],
+          tree_ensemble_serialized=restored_tensors[1])
+
+
+class TreeEnsemble():
+  """Creates TreeEnsemble resource."""
+
+  def __init__(self, name, stamp_token=0, is_local=False, serialized_proto=''):
+    self._stamp_token = stamp_token
+    self._serialized_proto = serialized_proto
+    self._is_local = is_local
+    with ops.name_scope(name, 'TreeEnsemble') as name:
+      self._name = name
+      self.resource_handle = self._create_resource()
+      self._init_op = self._initialize()
+      is_initialized_op = self.is_initialized()
+      # Adds the variable to the savable list.
+      if not is_local:
+        ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS,
+                              _TreeEnsembleSavable(
+                                  self.resource_handle, self.initializer,
+                                  self.resource_handle.name))
+      resources.register_resource(
+          self.resource_handle,
+          self.initializer,
+          is_initialized_op,
+          is_shared=not is_local)
+
+  def _create_resource(self):
+    return gen_boosted_trees_ops.boosted_trees_ensemble_resource_handle_op(
+        container='', shared_name=self._name, name=self._name)
+
+  def _initialize(self):
+    return gen_boosted_trees_ops.boosted_trees_create_ensemble(
+        self.resource_handle,
+        self._stamp_token,
+        tree_ensemble_serialized=self._serialized_proto)
+
+  @property
+  def initializer(self):
+    if self._init_op is None:
+      self._init_op = self._initialize()
+    return self._init_op
+
+  def is_initialized(self):
+    return gen_boosted_trees_ops.is_boosted_trees_ensemble_initialized(
+        self.resource_handle)
+
+  def _serialize_to_tensors(self):
+    raise NotImplementedError('When the need arises, TF2 compatibility can be '
+                              'added by implementing this method, along with '
+                              '_restore_from_tensors below.')
+
+  def _restore_from_tensors(self, restored_tensors):
+    raise NotImplementedError('When the need arises, TF2 compatibility can be '
+                              'added by implementing this method, along with '
+                              '_serialize_to_tensors above.')
+
+  def get_stamp_token(self):
+    """Returns the current stamp token of the resource."""
+    stamp_token, _, _, _, _ = (
+        gen_boosted_trees_ops.boosted_trees_get_ensemble_states(
+            self.resource_handle))
+    return stamp_token
+
+  def get_states(self):
+    """Returns states of the tree ensemble.
+
+    Returns:
+      stamp_token, num_trees, num_finalized_trees, num_attempted_layers and
+      range of the nodes in the latest layer.
+    """
+    (stamp_token, num_trees, num_finalized_trees, num_attempted_layers,
+     nodes_range) = (
+         gen_boosted_trees_ops.boosted_trees_get_ensemble_states(
+             self.resource_handle))
+    # Use identity to give names.
+    return (array_ops.identity(stamp_token, name='stamp_token'),
+            array_ops.identity(num_trees, name='num_trees'),
+            array_ops.identity(num_finalized_trees, name='num_finalized_trees'),
+            array_ops.identity(
+                num_attempted_layers, name='num_attempted_layers'),
+            array_ops.identity(nodes_range, name='last_layer_nodes_range'))
+
+  def serialize(self):
+    """Serializes the ensemble into proto and returns the serialized proto.
+
+    Returns:
+      stamp_token: int64 scalar Tensor to denote the stamp of the resource.
+      serialized_proto: string scalar Tensor of the serialized proto.
+    """
+    return gen_boosted_trees_ops.boosted_trees_serialize_ensemble(
+        self.resource_handle)
+
+  def deserialize(self, stamp_token, serialized_proto):
+    """Deserialize the input proto and resets the ensemble from it.
+
+    Args:
+      stamp_token: int64 scalar Tensor to denote the stamp of the resource.
+      serialized_proto: string scalar Tensor of the serialized proto.
+
+    Returns:
+      Operation (for dependencies).
+    """
+    return gen_boosted_trees_ops.boosted_trees_deserialize_ensemble(
+        self.resource_handle, stamp_token, serialized_proto)

tensorflow/python/ops/candidate_sampling_ops.py

@@ -0,0 +1,513 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Wrappers for candidate sampling operations."""
+
+from tensorflow.python.framework import random_seed
+from tensorflow.python.ops import array_ops  # pylint: disable=unused-import
+from tensorflow.python.ops import gen_candidate_sampling_ops
+from tensorflow.python.ops import math_ops  # pylint: disable=unused-import
+from tensorflow.python.util import deprecation
+from tensorflow.python.util import dispatch
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(
+    'random.uniform_candidate_sampler',
+    v1=['random.uniform_candidate_sampler', 'nn.uniform_candidate_sampler'])
+@dispatch.add_dispatch_support
+@deprecation.deprecated_endpoints('nn.uniform_candidate_sampler')
+def uniform_candidate_sampler(true_classes, num_true, num_sampled, unique,
+                              range_max, seed=None, name=None):
+  """Samples a set of classes using a uniform base distribution.
+
+  This operation randomly samples a tensor of sampled classes
+  (`sampled_candidates`) from the range of integers `[0, range_max)`.
+
+  See the [Candidate Sampling Algorithms
+  Reference](http://www.tensorflow.org/extras/candidate_sampling.pdf)
+  for a quick course on Candidate Sampling.
+
+  The elements of `sampled_candidates` are drawn without replacement
+  (if `unique=True`) or with replacement (if `unique=False`) from
+  the base distribution.
+
+  The base distribution for this operation is the uniform distribution
+  over the range of integers `[0, range_max)`.
+
+  In addition, this operation returns tensors `true_expected_count`
+  and `sampled_expected_count` representing the number of times each
+  of the target classes (`true_classes`) and the sampled
+  classes (`sampled_candidates`) is expected to occur in an average
+  tensor of sampled classes. These values correspond to `Q(y|x)`
+  defined in the [Candidate Sampling Algorithms
+  Reference](http://www.tensorflow.org/extras/candidate_sampling.pdf).
+  If `unique=True`, then these are post-rejection probabilities and we
+  compute them approximately.
+
+  Note that this function (and also other `*_candidate_sampler`
+  functions) only gives you the ingredients to implement the various
+  Candidate Sampling algorithms listed in the big table in the
+  [Candidate Sampling Algorithms
+  Reference](http://www.tensorflow.org/extras/candidate_sampling.pdf). You
+  still need to implement the algorithms yourself.
+
+  For example, according to that table, the phrase "negative samples"
+  may mean different things in different algorithms. For instance, in
+  NCE, "negative samples" means `S_i` (which is just the sampled
+  classes) which may overlap with true classes, while in Sampled
+  Logistic, "negative samples" means `S_i - T_i` which excludes the
+  true classes. The return value `sampled_candidates` corresponds to
+  `S_i`, not to any specific definition of "negative samples" in any
+  specific algorithm. It's your responsibility to pick an algorithm
+  and calculate the "negative samples" defined by that algorithm
+  (e.g. `S_i - T_i`).
+
+  As another example, the `true_classes` argument is for calculating
+  the `true_expected_count` output (as a by-product of this function's
+  main calculation), which may be needed by some algorithms (according
+  to that table). It's not for excluding true classes in the return
+  value `sampled_candidates`. Again that step is algorithm-specific
+  and should be carried out by you.
+
+  Args:
+    true_classes: A `Tensor` of type `int64` and shape `[batch_size,
+      num_true]`. The target classes.
+    num_true: An `int`.  The number of target classes per training example.
+    num_sampled: An `int`.  The number of classes to randomly sample. The
+      `sampled_candidates` return value will have shape `[num_sampled]`. If
+      `unique=True`, `num_sampled` must be less than or equal to `range_max`.
+    unique: A `bool`. Determines whether all sampled classes in a batch are
+      unique.
+    range_max: An `int`. The number of possible classes.
+    seed: An `int`. An operation-specific seed. Default is 0.
+    name: A name for the operation (optional).
+
+  Returns:
+    sampled_candidates: A tensor of type `int64` and shape
+      `[num_sampled]`. The sampled classes, either with possible
+      duplicates (`unique=False`) or all unique (`unique=True`). As
+      noted above, `sampled_candidates` may overlap with true classes.
+    true_expected_count: A tensor of type `float`.  Same shape as
+      `true_classes`. The expected counts under the sampling distribution
+      of each of `true_classes`.
+    sampled_expected_count: A tensor of type `float`. Same shape as
+      `sampled_candidates`. The expected counts under the sampling distribution
+      of each of `sampled_candidates`.
+  """
+  seed1, seed2 = random_seed.get_seed(seed)
+  return gen_candidate_sampling_ops.uniform_candidate_sampler(
+      true_classes, num_true, num_sampled, unique, range_max, seed=seed1,
+      seed2=seed2, name=name)
+
+
+@tf_export(
+    'random.log_uniform_candidate_sampler',
+    v1=[
+        'random.log_uniform_candidate_sampler',
+        'nn.log_uniform_candidate_sampler'
+    ])
+@dispatch.add_dispatch_support
+@deprecation.deprecated_endpoints('nn.log_uniform_candidate_sampler')
+def log_uniform_candidate_sampler(true_classes, num_true, num_sampled, unique,
+                                  range_max, seed=None, name=None):
+  """Samples a set of classes using a log-uniform (Zipfian) base distribution.
+
+  This operation randomly samples a tensor of sampled classes
+  (`sampled_candidates`) from the range of integers `[0, range_max)`.
+
+  See the [Candidate Sampling Algorithms
+  Reference](http://www.tensorflow.org/extras/candidate_sampling.pdf)
+  for a quick course on Candidate Sampling.
+
+  The elements of `sampled_candidates` are drawn without replacement
+  (if `unique=True`) or with replacement (if `unique=False`) from
+  the base distribution.
+
+  The base distribution for this operation is an approximately log-uniform
+  or Zipfian distribution:
+
+  `P(class) = (log(class + 2) - log(class + 1)) / log(range_max + 1)`
+
+  This sampler is useful when the target classes approximately follow such
+  a distribution - for example, if the classes represent words in a lexicon
+  sorted in decreasing order of frequency. If your classes are not ordered by
+  decreasing frequency, do not use this op.
+
+  In addition, this operation returns tensors `true_expected_count`
+  and `sampled_expected_count` representing the number of times each
+  of the target classes (`true_classes`) and the sampled
+  classes (`sampled_candidates`) is expected to occur in an average
+  tensor of sampled classes.  These values correspond to `Q(y|x)`
+  defined in the [Candidate Sampling Algorithms
+  Reference](http://www.tensorflow.org/extras/candidate_sampling.pdf).
+  If `unique=True`, then these are post-rejection probabilities and we
+  compute them approximately.
+
+  Note that this function (and also other `*_candidate_sampler`
+  functions) only gives you the ingredients to implement the various
+  Candidate Sampling algorithms listed in the big table in the
+  [Candidate Sampling Algorithms
+  Reference](http://www.tensorflow.org/extras/candidate_sampling.pdf). You
+  still need to implement the algorithms yourself.
+
+  For example, according to that table, the phrase "negative samples"
+  may mean different things in different algorithms. For instance, in
+  NCE, "negative samples" means `S_i` (which is just the sampled
+  classes) which may overlap with true classes, while in Sampled
+  Logistic, "negative samples" means `S_i - T_i` which excludes the
+  true classes. The return value `sampled_candidates` corresponds to
+  `S_i`, not to any specific definition of "negative samples" in any
+  specific algorithm. It's your responsibility to pick an algorithm
+  and calculate the "negative samples" defined by that algorithm
+  (e.g. `S_i - T_i`).
+
+  As another example, the `true_classes` argument is for calculating
+  the `true_expected_count` output (as a by-product of this function's
+  main calculation), which may be needed by some algorithms (according
+  to that table). It's not for excluding true classes in the return
+  value `sampled_candidates`. Again that step is algorithm-specific
+  and should be carried out by you.
+
+  Args:
+    true_classes: A `Tensor` of type `int64` and shape `[batch_size,
+      num_true]`. The target classes.
+    num_true: An `int`.  The number of target classes per training example.
+    num_sampled: An `int`.  The number of classes to randomly sample.
+    unique: A `bool`. Determines whether all sampled classes in a batch are
+      unique.
+    range_max: An `int`. The number of possible classes.
+    seed: An `int`. An operation-specific seed. Default is 0.
+    name: A name for the operation (optional).
+
+  Returns:
+    sampled_candidates: A tensor of type `int64` and shape
+      `[num_sampled]`. The sampled classes. As noted above,
+      `sampled_candidates` may overlap with true classes.
+    true_expected_count: A tensor of type `float`.  Same shape as
+      `true_classes`. The expected counts under the sampling distribution
+      of each of `true_classes`.
+    sampled_expected_count: A tensor of type `float`. Same shape as
+      `sampled_candidates`. The expected counts under the sampling distribution
+      of each of `sampled_candidates`.
+  """
+  seed1, seed2 = random_seed.get_seed(seed)
+  return gen_candidate_sampling_ops.log_uniform_candidate_sampler(
+      true_classes, num_true, num_sampled, unique, range_max, seed=seed1,
+      seed2=seed2, name=name)
+
+
+@tf_export(
+    'random.learned_unigram_candidate_sampler',
+    'nn.learned_unigram_candidate_sampler')
+@dispatch.add_dispatch_support
+@deprecation.deprecated_endpoints(['nn.learned_unigram_candidate_sampler'])
+def learned_unigram_candidate_sampler(true_classes, num_true, num_sampled,
+                                      unique, range_max, seed=None, name=None):
+  """Samples a set of classes from a distribution learned during training.
+
+  This operation randomly samples a tensor of sampled classes
+  (`sampled_candidates`) from the range of integers `[0, range_max)`.
+
+  See the [Candidate Sampling Algorithms
+  Reference](http://www.tensorflow.org/extras/candidate_sampling.pdf)
+  for a quick course on Candidate Sampling.
+
+  The elements of `sampled_candidates` are drawn without replacement
+  (if `unique=True`) or with replacement (if `unique=False`) from
+  the base distribution.
+
+  The base distribution for this operation is constructed on the fly
+  during training.  It is a unigram distribution over the target
+  classes seen so far during training.  Every integer in `[0, range_max)`
+  begins with a weight of 1, and is incremented by 1 each time it is
+  seen as a target class.  The base distribution is not saved to checkpoints,
+  so it is reset when the model is reloaded.
+
+  In addition, this operation returns tensors `true_expected_count`
+  and `sampled_expected_count` representing the number of times each
+  of the target classes (`true_classes`) and the sampled
+  classes (`sampled_candidates`) is expected to occur in an average
+  tensor of sampled classes.  These values correspond to `Q(y|x)`
+  defined in the [Candidate Sampling Algorithms
+  Reference](http://www.tensorflow.org/extras/candidate_sampling.pdf).
+  If `unique=True`, then these are post-rejection probabilities and we
+  compute them approximately.
+
+  Note that this function (and also other `*_candidate_sampler`
+  functions) only gives you the ingredients to implement the various
+  Candidate Sampling algorithms listed in the big table in the
+  [Candidate Sampling Algorithms
+  Reference](http://www.tensorflow.org/extras/candidate_sampling.pdf). You
+  still need to implement the algorithms yourself.
+
+  For example, according to that table, the phrase "negative samples"
+  may mean different things in different algorithms. For instance, in
+  NCE, "negative samples" means `S_i` (which is just the sampled
+  classes) which may overlap with true classes, while in Sampled
+  Logistic, "negative samples" means `S_i - T_i` which excludes the
+  true classes. The return value `sampled_candidates` corresponds to
+  `S_i`, not to any specific definition of "negative samples" in any
+  specific algorithm. It's your responsibility to pick an algorithm
+  and calculate the "negative samples" defined by that algorithm
+  (e.g. `S_i - T_i`).
+
+  As another example, the `true_classes` argument is for calculating
+  the `true_expected_count` output (as a by-product of this function's
+  main calculation), which may be needed by some algorithms (according
+  to that table). It's not for excluding true classes in the return
+  value `sampled_candidates`. Again that step is algorithm-specific
+  and should be carried out by you.
+
+  Args:
+    true_classes: A `Tensor` of type `int64` and shape `[batch_size,
+      num_true]`. The target classes.
+    num_true: An `int`.  The number of target classes per training example.
+    num_sampled: An `int`.  The number of classes to randomly sample.
+    unique: A `bool`. Determines whether all sampled classes in a batch are
+      unique.
+    range_max: An `int`. The number of possible classes.
+    seed: An `int`. An operation-specific seed. Default is 0.
+    name: A name for the operation (optional).
+
+  Returns:
+    sampled_candidates: A tensor of type `int64` and shape
+      `[num_sampled]`. The sampled classes. As noted above,
+      `sampled_candidates` may overlap with true classes.
+    true_expected_count: A tensor of type `float`.  Same shape as
+      `true_classes`. The expected counts under the sampling distribution
+      of each of `true_classes`.
+    sampled_expected_count: A tensor of type `float`. Same shape as
+      `sampled_candidates`. The expected counts under the sampling distribution
+      of each of `sampled_candidates`.
+
+  """
+  seed1, seed2 = random_seed.get_seed(seed)
+  # Limiting to Max int32 value
+  if range_max > 2147483647:
+    raise ValueError(f'Value of range_max:{range_max} is too large to handle')
+  return gen_candidate_sampling_ops.learned_unigram_candidate_sampler(
+      true_classes, num_true, num_sampled, unique, range_max, seed=seed1,
+      seed2=seed2, name=name)
+
+
+@tf_export('random.fixed_unigram_candidate_sampler',
+           'nn.fixed_unigram_candidate_sampler')
+@dispatch.add_dispatch_support
+def fixed_unigram_candidate_sampler(true_classes,
+                                    num_true,
+                                    num_sampled,
+                                    unique,
+                                    range_max,
+                                    vocab_file='',
+                                    distortion=1.0,
+                                    num_reserved_ids=0,
+                                    num_shards=1,
+                                    shard=0,
+                                    unigrams=(),
+                                    seed=None,
+                                    name=None):
+  """Samples a set of classes using the provided (fixed) base distribution.
+
+  This operation randomly samples a tensor of sampled classes
+  (`sampled_candidates`) from the range of integers `[0, range_max)`.
+
+  See the [Candidate Sampling Algorithms
+  Reference](http://www.tensorflow.org/extras/candidate_sampling.pdf)
+  for a quick course on Candidate Sampling.
+
+  The elements of `sampled_candidates` are drawn without replacement
+  (if `unique=True`) or with replacement (if `unique=False`) from
+  the base distribution.
+
+  The base distribution is read from a file or passed in as an
+  in-memory array. There is also an option to skew the distribution by
+  applying a distortion power to the weights.
+
+  In addition, this operation returns tensors `true_expected_count`
+  and `sampled_expected_count` representing the number of times each
+  of the target classes (`true_classes`) and the sampled
+  classes (`sampled_candidates`) is expected to occur in an average
+  tensor of sampled classes.  These values correspond to `Q(y|x)`
+  defined in the [Candidate Sampling Algorithms
+  Reference](http://www.tensorflow.org/extras/candidate_sampling.pdf).
+  If `unique=True`, then these are post-rejection probabilities and we
+  compute them approximately.
+
+  Note that this function (and also other `*_candidate_sampler`
+  functions) only gives you the ingredients to implement the various
+  Candidate Sampling algorithms listed in the big table in the
+  [Candidate Sampling Algorithms
+  Reference](http://www.tensorflow.org/extras/candidate_sampling.pdf). You
+  still need to implement the algorithms yourself.
+
+  For example, according to that table, the phrase "negative samples"
+  may mean different things in different algorithms. For instance, in
+  NCE, "negative samples" means `S_i` (which is just the sampled
+  classes) which may overlap with true classes, while in Sampled
+  Logistic, "negative samples" means `S_i - T_i` which excludes the
+  true classes. The return value `sampled_candidates` corresponds to
+  `S_i`, not to any specific definition of "negative samples" in any
+  specific algorithm. It's your responsibility to pick an algorithm
+  and calculate the "negative samples" defined by that algorithm
+  (e.g. `S_i - T_i`).
+
+  As another example, the `true_classes` argument is for calculating
+  the `true_expected_count` output (as a by-product of this function's
+  main calculation), which may be needed by some algorithms (according
+  to that table). It's not for excluding true classes in the return
+  value `sampled_candidates`. Again that step is algorithm-specific
+  and should be carried out by you.
+
+  Args:
+    true_classes: A `Tensor` of type `int64` and shape `[batch_size,
+      num_true]`. The target classes.
+    num_true: An `int`.  The number of target classes per training example.
+    num_sampled: An `int`.  The number of classes to randomly sample.
+    unique: A `bool`. Determines whether all sampled classes in a batch are
+      unique.
+    range_max: An `int`. The number of possible classes.
+    vocab_file: Each valid line in this file (which should have a CSV-like
+      format) corresponds to a valid word ID. IDs are in sequential order,
+      starting from num_reserved_ids. The last entry in each line is expected
+      to be a value corresponding to the count or relative probability. Exactly
+      one of `vocab_file` and `unigrams` needs to be passed to this operation.
+    distortion: The distortion is used to skew the unigram probability
+      distribution.  Each weight is first raised to the distortion's power
+      before adding to the internal unigram distribution. As a result,
+      `distortion = 1.0` gives regular unigram sampling (as defined by the vocab
+      file), and `distortion = 0.0` gives a uniform distribution.
+    num_reserved_ids: Optionally some reserved IDs can be added in the range
+      `[0, num_reserved_ids)` by the users. One use case is that a special
+      unknown word token is used as ID 0. These IDs will have a sampling
+      probability of 0.
+    num_shards: A sampler can be used to sample from a subset of the original
+      range in order to speed up the whole computation through parallelism. This
+      parameter (together with `shard`) indicates the number of partitions that
+      are being used in the overall computation.
+    shard: A sampler can be used to sample from a subset of the original range
+      in order to speed up the whole computation through parallelism. This
+      parameter (together with `num_shards`) indicates the particular partition
+      number of the operation, when partitioning is being used.
+    unigrams: A list of unigram counts or probabilities, one per ID in
+      sequential order. Exactly one of `vocab_file` and `unigrams` should be
+      passed to this operation.
+    seed: An `int`. An operation-specific seed. Default is 0.
+    name: A name for the operation (optional).
+
+  Returns:
+    sampled_candidates: A tensor of type `int64` and shape
+      `[num_sampled]`. The sampled classes. As noted above,
+      `sampled_candidates` may overlap with true classes.
+    true_expected_count: A tensor of type `float`.  Same shape as
+      `true_classes`. The expected counts under the sampling distribution
+      of each of `true_classes`.
+    sampled_expected_count: A tensor of type `float`. Same shape as
+      `sampled_candidates`. The expected counts under the sampling distribution
+      of each of `sampled_candidates`.
+
+  """
+  seed1, seed2 = random_seed.get_seed(seed)
+  return gen_candidate_sampling_ops.fixed_unigram_candidate_sampler(
+      true_classes, num_true, num_sampled, unique, range_max,
+      vocab_file=vocab_file, distortion=distortion,
+      num_reserved_ids=num_reserved_ids, num_shards=num_shards, shard=shard,
+      unigrams=unigrams, seed=seed1, seed2=seed2, name=name)
+
+
+@tf_export('random.all_candidate_sampler', 'nn.all_candidate_sampler')
+def all_candidate_sampler(true_classes, num_true, num_sampled, unique,
+                          seed=None, name=None):
+  """Generate the set of all classes.
+
+  Deterministically generates and returns the set of all possible classes.
+  For testing purposes.  There is no need to use this, since you might as
+  well use full softmax or full logistic regression.
+
+  Args:
+    true_classes: A `Tensor` of type `int64` and shape `[batch_size,
+      num_true]`. The target classes.
+    num_true: An `int`.  The number of target classes per training example.
+    num_sampled: An `int`.  The number of possible classes.
+    unique: A `bool`. Ignored.
+      unique.
+    seed: An `int`. An operation-specific seed. Default is 0.
+    name: A name for the operation (optional).
+
+  Returns:
+    sampled_candidates: A tensor of type `int64` and shape `[num_sampled]`.
+      This operation deterministically returns the entire range
+      `[0, num_sampled]`.
+    true_expected_count: A tensor of type `float`.  Same shape as
+      `true_classes`. The expected counts under the sampling distribution
+      of each of `true_classes`. All returned values are 1.0.
+    sampled_expected_count: A tensor of type `float`. Same shape as
+      `sampled_candidates`. The expected counts under the sampling distribution
+      of each of `sampled_candidates`. All returned values are 1.0.
+  """
+  seed1, seed2 = random_seed.get_seed(seed)
+  return gen_candidate_sampling_ops.all_candidate_sampler(
+      true_classes, num_true, num_sampled, unique, seed=seed1, seed2=seed2,
+      name=name)
+
+
+@tf_export('nn.compute_accidental_hits')
+@dispatch.add_dispatch_support
+def compute_accidental_hits(true_classes, sampled_candidates, num_true,
+                            seed=None, name=None):
+  """Compute the position ids in `sampled_candidates` matching `true_classes`.
+
+  In Candidate Sampling, this operation facilitates virtually removing
+  sampled classes which happen to match target classes.  This is done
+  in Sampled Softmax and Sampled Logistic.
+
+  See our [Candidate Sampling Algorithms
+  Reference](http://www.tensorflow.org/extras/candidate_sampling.pdf).
+
+  We presuppose that the `sampled_candidates` are unique.
+
+  We call it an 'accidental hit' when one of the target classes
+  matches one of the sampled classes.  This operation reports
+  accidental hits as triples `(index, id, weight)`, where `index`
+  represents the row number in `true_classes`, `id` represents the
+  position in `sampled_candidates`, and weight is `-FLOAT_MAX`.
+
+  The result of this op should be passed through a `sparse_to_dense`
+  operation, then added to the logits of the sampled classes. This
+  removes the contradictory effect of accidentally sampling the true
+  target classes as noise classes for the same example.
+
+  Args:
+    true_classes: A `Tensor` of type `int64` and shape `[batch_size,
+      num_true]`. The target classes.
+    sampled_candidates: A tensor of type `int64` and shape `[num_sampled]`.
+      The sampled_candidates output of CandidateSampler.
+    num_true: An `int`.  The number of target classes per training example.
+    seed: An `int`. An operation-specific seed. Default is 0.
+    name: A name for the operation (optional).
+
+  Returns:
+    indices: A `Tensor` of type `int32` and shape `[num_accidental_hits]`.
+      Values indicate rows in `true_classes`.
+    ids: A `Tensor` of type `int64` and shape `[num_accidental_hits]`.
+      Values indicate positions in `sampled_candidates`.
+    weights: A `Tensor` of type `float` and shape `[num_accidental_hits]`.
+      Each value is `-FLOAT_MAX`.
+
+  """
+  seed1, seed2 = random_seed.get_seed(seed)
+  return gen_candidate_sampling_ops.compute_accidental_hits(
+      true_classes, sampled_candidates, num_true, seed=seed1, seed2=seed2,
+      name=name)

tensorflow/python/ops/clip_ops.py

@@ -0,0 +1,440 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Operations for clipping (gradient, weight) tensors to min/max values."""
+
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.ops import gen_nn_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.util import deprecation
+from tensorflow.python.util import dispatch
+from tensorflow.python.util.compat import collections_abc
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("clip_by_value")
+@dispatch.register_unary_elementwise_api
+@dispatch.add_dispatch_support
+def clip_by_value(t, clip_value_min, clip_value_max,
+                  name=None):
+  """Clips tensor values to a specified min and max.
+
+  Given a tensor `t`, this operation returns a tensor of the same type and
+  shape as `t` with its values clipped to `clip_value_min` and `clip_value_max`.
+  Any values less than `clip_value_min` are set to `clip_value_min`. Any values
+  greater than `clip_value_max` are set to `clip_value_max`.
+
+  Note: `clip_value_min` needs to be smaller or equal to `clip_value_max` for
+  correct results.
+
+  For example:
+
+  Basic usage passes a scalar as the min and max value.
+
+  >>> t = tf.constant([[-10., -1., 0.], [0., 2., 10.]])
+  >>> t2 = tf.clip_by_value(t, clip_value_min=-1, clip_value_max=1)
+  >>> t2.numpy()
+  array([[-1., -1.,  0.],
+         [ 0.,  1.,  1.]], dtype=float32)
+
+  The min and max can be the same size as `t`, or broadcastable to that size.
+
+  >>> t = tf.constant([[-1, 0., 10.], [-1, 0, 10]])
+  >>> clip_min = [[2],[1]]
+  >>> t3 = tf.clip_by_value(t, clip_value_min=clip_min, clip_value_max=100)
+  >>> t3.numpy()
+  array([[ 2.,  2., 10.],
+         [ 1.,  1., 10.]], dtype=float32)
+
+  Broadcasting fails, intentionally, if you would expand the dimensions of `t`
+
+  >>> t = tf.constant([[-1, 0., 10.], [-1, 0, 10]])
+  >>> clip_min = [[[2, 1]]] # Has a third axis
+  >>> t4 = tf.clip_by_value(t, clip_value_min=clip_min, clip_value_max=100)
+  Traceback (most recent call last):
+  ...
+  InvalidArgumentError: Incompatible shapes: [2,3] vs. [1,1,2]
+
+  It throws a `TypeError` if you try to clip an `int` to a `float` value
+  (`tf.cast` the input to `float` first).
+
+  >>> t = tf.constant([[1, 2], [3, 4]], dtype=tf.int32)
+  >>> t5 = tf.clip_by_value(t, clip_value_min=-3.1, clip_value_max=3.1)
+  Traceback (most recent call last):
+  ...
+  TypeError: Cannot convert ...
+
+
+  Args:
+    t: A `Tensor` or `IndexedSlices`.
+    clip_value_min: The minimum value to clip to. A scalar `Tensor` or one that
+      is broadcastable to the shape of `t`.
+    clip_value_max: The maximum value to clip to. A scalar `Tensor` or one that
+      is broadcastable to the shape of `t`.
+    name: A name for the operation (optional).
+
+  Returns:
+    A clipped `Tensor` or `IndexedSlices`.
+
+  Raises:
+    `tf.errors.InvalidArgumentError`: If the clip tensors would trigger array
+      broadcasting that would make the returned tensor larger than the input.
+    TypeError: If dtype of the input is `int32` and dtype of
+      the `clip_value_min` or `clip_value_max` is `float32`
+  """
+  with ops.name_scope(name, "clip_by_value",
+                      [t, clip_value_min, clip_value_max]) as name:
+    values = ops.convert_to_tensor(
+        t.values if isinstance(t, indexed_slices.IndexedSlices) else t,
+        name="t")
+
+    # Go through list of tensors, for each value in each tensor clip
+    t_min = math_ops.minimum(values, clip_value_max)
+    # Assert that the shape is compatible with the initial shape,
+    # to prevent unintentional broadcasting.
+    values.shape.assert_is_compatible_with(t_min.shape)
+
+    t_max = math_ops.maximum(t_min, clip_value_min, name=name)
+    values.shape.assert_is_compatible_with(t_max.shape)
+
+    if isinstance(t, indexed_slices.IndexedSlices):
+      t_max = indexed_slices.IndexedSlices(t_max, t.indices, t.dense_shape)
+
+  return t_max
+  # TODO(scottzhu): switch to use new implementation in 2 weeks.
+  # return gen_math_ops.clip_by_value(
+  #     t, clip_value_min, clip_value_max, name=name)
+
+
+@ops.RegisterGradient("ClipByValue")
+def _clip_by_value_grad(op, grad):
+  """Returns grad of clip_by_value."""
+  x = op.inputs[0]
+  y = op.inputs[1]
+  z = op.inputs[2]
+  gdtype = grad.dtype
+  sx = array_ops.shape(x)
+  sy = array_ops.shape(y)
+  sz = array_ops.shape(z)
+  gradshape = array_ops.shape(grad)
+  zeros = array_ops.zeros(gradshape, gdtype)
+  xymask = math_ops.less(x, y)
+  xzmask = math_ops.greater(x, z)
+  _, ry = gen_array_ops.broadcast_gradient_args(sx, sy)
+  _, rz = gen_array_ops.broadcast_gradient_args(sx, sz)
+  xgrad = array_ops.where(math_ops.logical_or(xymask, xzmask), zeros, grad)
+  ygrad = array_ops.where(xymask, grad, zeros)
+  zgrad = array_ops.where(xzmask, grad, zeros)
+  gy = array_ops.reshape(math_ops.reduce_sum(ygrad, ry), sy)
+  gz = array_ops.reshape(math_ops.reduce_sum(zgrad, rz), sz)
+  return xgrad, gy, gz
+
+
+@tf_export("clip_by_norm")
+@dispatch.add_dispatch_support
+def clip_by_norm(t, clip_norm, axes=None, name=None):
+  """Clips tensor values to a maximum L2-norm.
+
+  Given a tensor `t`, and a maximum clip value `clip_norm`, this operation
+  normalizes `t` so that its L2-norm is less than or equal to `clip_norm`,
+  along the dimensions given in `axes`. Specifically, in the default case
+  where all dimensions are used for calculation, if the L2-norm of `t` is
+  already less than or equal to `clip_norm`, then `t` is not modified. If
+  the L2-norm is greater than `clip_norm`, then this operation returns a
+  tensor of the same type and shape as `t` with its values set to:
+
+  `t * clip_norm / l2norm(t)`
+
+  In this case, the L2-norm of the output tensor is `clip_norm`.
+
+  As another example, if `t` is a matrix and `axes == [1]`, then each row
+  of the output will have L2-norm less than or equal to `clip_norm`. If
+  `axes == [0]` instead, each column of the output will be clipped.
+
+  Code example:
+
+  >>> some_nums = tf.constant([[1, 2, 3, 4, 5]], dtype=tf.float32)
+  >>> tf.clip_by_norm(some_nums, 2.0).numpy()
+  array([[0.26967996, 0.5393599 , 0.80903983, 1.0787199 , 1.3483998 ]],
+        dtype=float32)
+
+  This operation is typically used to clip gradients before applying them with
+  an optimizer.  Most gradient data is a collection of different shaped tensors
+  for different parts of the model.  Thus, this is a common usage:
+
+  ```
+  # Get your gradients after training
+  loss_value, grads = grad(model, features, labels)
+
+  # Apply some clipping
+  grads = [tf.clip_by_norm(g, norm)
+               for g in grads]
+
+  # Continue on with training
+  optimizer.apply_gradients(grads)
+  ```
+
+  Args:
+    t: A `Tensor` or `IndexedSlices`.  This must be a floating point type.
+    clip_norm: A 0-D (scalar) `Tensor` > 0. A maximum clipping value, also
+      floating point.
+      Note: If a negative clip_norm is provided, it will be treated as zero.
+    axes: A 1-D (vector) `Tensor` of type int32 containing the dimensions to use
+      for computing the L2-norm. If `None` (the default), uses all dimensions.
+    name: A name for the operation (optional).
+
+  Returns:
+    A clipped `Tensor` or `IndexedSlices`.
+
+  Raises:
+    ValueError: If the clip_norm tensor is not a 0-D scalar tensor.
+    TypeError: If dtype of the input is not a floating point or
+      complex type.
+  """
+  with ops.name_scope(name, "clip_by_norm", [t, clip_norm]) as name:
+    values = ops.convert_to_tensor(
+        t.values if isinstance(t, indexed_slices.IndexedSlices) else t,
+        name="t")
+
+    if np.isscalar(clip_norm):
+      if clip_norm < 0:
+        clip_norm = 0
+    else:
+      clip_norm = math_ops.cast(
+          math_ops.maximum(clip_norm, 0), dtype=values.dtype
+      )
+
+    # Calculate L2-norm, clip elements by ratio of clip_norm to L2-norm
+    l2sum = math_ops.reduce_sum(values * values, axes, keepdims=True)
+    pred = l2sum > 0
+    # Two-tap tf.where trick to bypass NaN gradients
+    l2sum_safe = array_ops.where(pred, l2sum, array_ops.ones_like(l2sum))
+    l2norm = array_ops.where(pred, math_ops.sqrt(l2sum_safe), l2sum)
+    intermediate = values * clip_norm
+    # Assert that the shape is compatible with the initial shape,
+    # to prevent unintentional broadcasting.
+    values.shape.assert_is_compatible_with(intermediate.shape)
+    values_clip = array_ops.identity(
+        intermediate / math_ops.maximum(l2norm, clip_norm), name=name)
+
+    if isinstance(t, indexed_slices.IndexedSlices):
+      return indexed_slices.IndexedSlices(values_clip, t.indices, t.dense_shape)
+
+    return values_clip
+
+
+@tf_export("linalg.global_norm", v1=["linalg.global_norm", "global_norm"])
+@dispatch.add_dispatch_support
+@deprecation.deprecated_endpoints("global_norm")
+def global_norm(t_list, name=None):
+  """Computes the global norm of multiple tensors.
+
+  Given a tuple or list of tensors `t_list`, this operation returns the
+  global norm of the elements in all tensors in `t_list`. The global norm is
+  computed as:
+
+  `global_norm = sqrt(sum([l2norm(t)**2 for t in t_list]))`
+
+  Any entries in `t_list` that are of type None are ignored.
+
+  Args:
+    t_list: A tuple or list of mixed `Tensors`, `IndexedSlices`, or None.
+    name: A name for the operation (optional).
+
+  Returns:
+    A 0-D (scalar) `Tensor` of type `float`.
+
+  Raises:
+    TypeError: If `t_list` is not a sequence.
+  """
+  if (not isinstance(t_list, collections_abc.Sequence) or
+      isinstance(t_list, str)):
+    raise TypeError("`t_list` should be a sequence of tensors. Received "
+                    f"{type(t_list)}.")
+  t_list = list(t_list)
+  with ops.name_scope(name, "global_norm", t_list) as name:
+    values = [
+        ops.convert_to_tensor(
+            t.values if isinstance(t, indexed_slices.IndexedSlices) else t,
+            name="t_%d" % i) if t is not None else t
+        for i, t in enumerate(t_list)
+    ]
+    half_squared_norms = []
+    for v in values:
+      if v is not None:
+        with ops.colocate_with(v):
+          half_squared_norms.append(gen_nn_ops.l2_loss(v))
+
+    half_squared_norm = math_ops.reduce_sum(
+        array_ops_stack.stack(half_squared_norms))
+
+    norm = math_ops.sqrt(
+        half_squared_norm *
+        constant_op.constant(2.0, dtype=half_squared_norm.dtype),
+        name="global_norm")
+
+  return norm
+
+
+@tf_export("clip_by_global_norm")
+@dispatch.add_dispatch_support
+def clip_by_global_norm(t_list, clip_norm, use_norm=None, name=None):
+  """Clips values of multiple tensors by the ratio of the sum of their norms.
+
+  Given a tuple or list of tensors `t_list`, and a clipping ratio `clip_norm`,
+  this operation returns a list of clipped tensors `list_clipped`
+  and the global norm (`global_norm`) of all tensors in `t_list`. Optionally,
+  if you've already computed the global norm for `t_list`, you can specify
+  the global norm with `use_norm`.
+
+  To perform the clipping, the values `t_list[i]` are set to:
+
+      t_list[i] * clip_norm / max(global_norm, clip_norm)
+
+  where:
+
+      global_norm = sqrt(sum([l2norm(t)**2 for t in t_list]))
+
+  If `clip_norm > global_norm` then the entries in `t_list` remain as they are,
+  otherwise they're all shrunk by the global ratio.
+
+  If `global_norm == infinity` then the entries in `t_list` are all set to `NaN`
+  to signal that an error occurred.
+
+  Any of the entries of `t_list` that are of type `None` are ignored.
+
+  This is the correct way to perform gradient clipping (Pascanu et al., 2012).
+
+  However, it is slower than `clip_by_norm()` because all the parameters must be
+  ready before the clipping operation can be performed.
+
+  Args:
+    t_list: A tuple or list of mixed `Tensors`, `IndexedSlices`, or None.
+    clip_norm: A 0-D (scalar) `Tensor` > 0. The clipping ratio.
+    use_norm: A 0-D (scalar) `Tensor` of type `float` (optional). The global
+      norm to use. If not provided, `global_norm()` is used to compute the norm.
+    name: A name for the operation (optional).
+
+  Returns:
+    list_clipped: A list of `Tensors` of the same type as `list_t`.
+    global_norm: A 0-D (scalar) `Tensor` representing the global norm.
+
+  Raises:
+    TypeError: If `t_list` is not a sequence.
+
+  References:
+    On the difficulty of training Recurrent Neural Networks:
+      [Pascanu et al., 2012](http://proceedings.mlr.press/v28/pascanu13.html)
+      ([pdf](http://proceedings.mlr.press/v28/pascanu13.pdf))
+  """
+  if (not isinstance(t_list, collections_abc.Sequence) or
+      isinstance(t_list, str)):
+    raise TypeError("`t_list` should be a sequence of tensors. Received "
+                    f"{type(t_list)}.")
+  t_list = list(t_list)
+  if use_norm is None:
+    use_norm = global_norm(t_list, name)
+
+  with ops.name_scope(name, "clip_by_global_norm",
+                      t_list + [clip_norm]) as name:
+    # Calculate L2-norm, clip elements by ratio of clip_norm to L2-norm
+    scale_for_finite = clip_norm * math_ops.minimum(
+        1.0 / use_norm,
+        constant_op.constant(1.0, dtype=use_norm.dtype) / clip_norm)
+    # If use_norm is any finite number, this is a no-op. For inf/-inf/NaN,
+    # this will make scale NaN.
+    scale = scale_for_finite + (use_norm - use_norm)
+
+    values = [
+        ops.convert_to_tensor(
+            t.values if isinstance(t, indexed_slices.IndexedSlices) else t,
+            name="t_%d" % i) if t is not None else t
+        for i, t in enumerate(t_list)
+    ]
+
+    values_clipped = []
+    for i, v in enumerate(values):
+      if v is None:
+        values_clipped.append(None)
+      else:
+        with ops.colocate_with(v):
+          values_clipped.append(
+              array_ops.identity(
+                  v * math_ops.cast(scale, v.dtype), name="%s_%d" % (name, i)
+              )
+          )
+
+    list_clipped = [
+        indexed_slices.IndexedSlices(c_v, t.indices, t.dense_shape)
+        if isinstance(t, indexed_slices.IndexedSlices) else c_v
+        for (c_v, t) in zip(values_clipped, t_list)
+    ]
+
+  return list_clipped, use_norm
+
+
+@deprecation.deprecated(
+    date=None,
+    instructions="clip_by_average_norm is deprecated in TensorFlow 2.0. Please "
+    "use clip_by_norm(t, clip_norm * tf.cast(tf.size(t), tf.float32), name) "
+    "instead.")
+@tf_export(v1=["clip_by_average_norm"])
+@dispatch.add_dispatch_support
+def clip_by_average_norm(t, clip_norm, name=None):
+  """Clips tensor values to a maximum average L2-norm.
+
+  Given a tensor `t`, and a maximum clip value `clip_norm`, this operation
+  normalizes `t` so that its average L2-norm is less than or equal to
+  `clip_norm`. Specifically, if the average L2-norm is already less than or
+  equal to `clip_norm`, then `t` is not modified. If the average L2-norm is
+  greater than `clip_norm`, then this operation returns a tensor of the same
+  type and shape as `t` with its values set to:
+
+  `t * clip_norm / l2norm_avg(t)`
+
+  In this case, the average L2-norm of the output tensor is `clip_norm`.
+
+  This operation is typically used to clip gradients before applying them with
+  an optimizer.
+
+  Args:
+    t: A `Tensor`.
+    clip_norm: A 0-D (scalar) `Tensor` > 0. A maximum clipping value.
+    name: A name for the operation (optional).
+
+  Returns:
+    A clipped `Tensor`.
+  """
+  with ops.name_scope(name, "clip_by_average_norm", [t, clip_norm]) as name:
+    t = ops.convert_to_tensor(t, name="t")
+
+    # Calculate L2-norm per element, clip elements by ratio of clip_norm to
+    # L2-norm per element
+    n_element = math_ops.cast(array_ops.size(t), dtypes.float32)
+    l2norm_inv = math_ops.rsqrt(
+        math_ops.reduce_sum(t * t, math_ops.range(array_ops.rank(t))))
+    tclip = array_ops.identity(
+        t * clip_norm * math_ops.minimum(
+            l2norm_inv * n_element, constant_op.constant(1.0) / clip_norm),
+        name=name)
+
+  return tclip

tensorflow/python/ops/clip_ops_test.py

@@ -0,0 +1,150 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for Clip Operations."""
+
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices as indexed_slices_lib
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import clip_ops
+from tensorflow.python.ops import numerics
+from tensorflow.python.platform import test
+
+
+class ClipOpsTest(test.TestCase):
+
+  def __init__(self, method_name="runTest"):
+    super(ClipOpsTest, self).__init__(method_name)
+
+  def _testClipTensorByNorm(self, inputs, max_norm, expected):
+    input_op = constant_op.constant(inputs)
+    clipped = clip_ops.clip_by_norm(input_op, max_norm)
+    check_op = numerics.add_check_numerics_ops()
+    result, _ = self.evaluate([clipped, check_op])
+    self.assertAllClose(result, expected)
+
+  def _testClipTensorByGlobalNorm(self, inputs, max_norm, expected):
+    clipped = clip_ops.clip_by_global_norm(inputs, max_norm)
+    result, _ = self.evaluate(clipped)
+    self.assertAllClose(result, expected)
+
+  def _testNonFiniteClippingByGlobalNorm(self, inputs, max_norm):
+    clipped = clip_ops.clip_by_global_norm(inputs, max_norm)
+    result, _ = self.evaluate(clipped)
+    self.assertTrue(np.all(np.isnan(result)))
+
+  def _testClipIndexedSlicesByNorm(self, values, indices, shape, max_norm,
+                                   axes):
+    values = constant_op.constant(values)
+    indices = constant_op.constant(indices)
+    shape = constant_op.constant(shape)
+    # IndexedSlices mode
+    indexed_slices = indexed_slices_lib.IndexedSlices(values, indices, shape)
+    clipped = clip_ops.clip_by_norm(indexed_slices, max_norm, axes)
+    # clipped should be IndexedSlices
+    self.assertIsInstance(clipped, indexed_slices_lib.IndexedSlices)
+    clipped = ops.convert_to_tensor(clipped)
+
+    # Tensor mode
+    dense_tensor = ops.convert_to_tensor(indexed_slices)
+    dense_clipped = clip_ops.clip_by_norm(dense_tensor, max_norm, axes)
+    result, expected = self.evaluate([clipped, dense_clipped])
+    self.assertAllClose(result, expected)
+
+  @test_util.run_deprecated_v1
+  def testClipTensorByNorm(self):
+    # Simple example
+    self._testClipTensorByNorm([[-3.0, 0.0, 0.0], [4.0, 0.0, 0.0]], 4.0,
+                               [[-2.4, 0.0, 0.0], [3.2, 0.0, 0.0]])
+    # No clipping.
+    self._testClipTensorByNorm([[1.0, 0.0, 0.0], [1.0, 0.0, 0.0]], 4.0,
+                               [[1.0, 0.0, 0.0], [1.0, 0.0, 0.0]])
+    # Zero norm
+    self._testClipTensorByNorm([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], 4.0,
+                               [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]])
+
+  @test_util.run_deprecated_v1
+  def testClipTensorByGlobalNorm(self):
+    # Simple example
+    self._testClipTensorByGlobalNorm([[-3.0, 0.0, 0.0], [4.0, 0.0, 0.0]], 4.0,
+                                     [[-2.4, 0.0, 0.0], [3.2, 0.0, 0.0]])
+    # No clipping.
+    self._testClipTensorByGlobalNorm([[1.0, 0.0, 0.0], [1.0, 0.0, 0.0]], 4.0,
+                                     [[1.0, 0.0, 0.0], [1.0, 0.0, 0.0]])
+    # Zero norm.
+    self._testClipTensorByGlobalNorm([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], 4.0,
+                                     [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]])
+
+  @test_util.run_deprecated_v1
+  def testClipTensorByGlobalNormMultipleDtypes(self):
+    (a, b), _ = self.evaluate(
+        clip_ops.clip_by_global_norm(
+            [
+                ops.convert_to_tensor([-3.0, 0.0, 0.0], dtype=dtypes.float32),
+                ops.convert_to_tensor([4.0, 0.0, 0.0], dtype=dtypes.bfloat16),
+            ],
+            clip_norm=4.0,
+            use_norm=ops.convert_to_tensor(5.0, dtype=dtypes.float32),
+        )
+    )
+    self.assertAllClose(a, [-2.4, 0.0, 0.0])
+    self.assertEqual(a.dtype, dtypes.float32)
+    self.assertAllClose(b, [3.203125, 0.0, 0.0])
+    self.assertEqual(b.dtype, dtypes.bfloat16)
+
+  @test_util.run_deprecated_v1
+  def testGlobalClipWithNonfinite(self):
+    self._testNonFiniteClippingByGlobalNorm(
+        [[-3.0, 0.0, 0.0], [float("inf"), 0.0, 0.0]], 4.0)
+    self._testNonFiniteClippingByGlobalNorm(
+        [[-3.0, 0.0, 0.0], [float("-inf"), 0.0, 0.0]], 4.0)
+    self._testNonFiniteClippingByGlobalNorm(
+        [[-3.0, 0.0, 0.0], [float("nan"), 0.0, 0.0]], 4.0)
+
+  def testClipIndexedSlicesByNorm(self):
+    values = [[[-3.0, 0.0, 0.0], [4.0, 0.0, 0.0]],
+              [[0.0, 2.0, 0.0], [0.0, 0.0, -1.0]]]
+    indices = [2, 6]
+    shape = [10, 2, 3]
+    # Axes == None
+    self._testClipIndexedSlicesByNorm(values, indices, shape, 4.0, None)
+
+    # Axes == 0
+    self._testClipIndexedSlicesByNorm(values, indices, shape, 4.0, 0)
+
+    # Axes == 1
+    self._testClipIndexedSlicesByNorm(values, indices, shape, 4.0, 1)
+
+    # Axes == 2
+    self._testClipIndexedSlicesByNorm(values, indices, shape, 4.0, 1)
+
+    # Axes == [0, 1]
+    self._testClipIndexedSlicesByNorm(values, indices, shape, 4.0, [0, 1])
+
+    # Axes == [0, 1]
+    self._testClipIndexedSlicesByNorm(values, indices, shape, 4.0, [0, 2])
+
+    # Axes == [0, 1]
+    self._testClipIndexedSlicesByNorm(values, indices, shape, 4.0, [1, 2])
+
+    # Axes == [0, 1]
+    self._testClipIndexedSlicesByNorm(values, indices, shape, 4.0, [0, 1, 2])
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/ops/clustering_ops.py

@@ -0,0 +1,774 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Clustering Operations."""
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import random_seed as random_seed_ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import gen_clustering_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_impl
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import while_loop
+from tensorflow.python.ops.embedding_ops import embedding_lookup
+# go/tf-wildcard-import
+# pylint: disable=wildcard-import
+from tensorflow.python.ops.gen_clustering_ops import *
+# pylint: enable=wildcard-import
+
+# Euclidean distance between vectors U and V is defined as \\(||U - V||_F\\)
+# which is the square root of the sum of the absolute squares of the elements
+# difference.
+SQUARED_EUCLIDEAN_DISTANCE = 'squared_euclidean'
+# Cosine distance between vectors U and V is defined as
+# \\(1 - (U \dot V) / (||U||_F ||V||_F)\\)
+COSINE_DISTANCE = 'cosine'
+
+RANDOM_INIT = 'random'
+KMEANS_PLUS_PLUS_INIT = 'kmeans_plus_plus'
+KMC2_INIT = 'kmc2'
+
+CLUSTERS_VAR_NAME = 'clusters'
+
+
+class KMeans:
+  """Creates the graph for k-means clustering."""
+
+  def __init__(self,
+               inputs,
+               num_clusters,
+               initial_clusters=RANDOM_INIT,
+               distance_metric=SQUARED_EUCLIDEAN_DISTANCE,
+               use_mini_batch=False,
+               mini_batch_steps_per_iteration=1,
+               random_seed=0,
+               kmeans_plus_plus_num_retries=2,
+               kmc2_chain_length=200):
+    """Creates an object for generating KMeans clustering graph.
+
+    This class implements the following variants of K-means algorithm:
+
+    If use_mini_batch is False, it runs standard full batch K-means. Each step
+    runs a single iteration of K-Means. This step can be run sharded across
+    multiple workers by passing a list of sharded inputs to this class. Note
+    however that a single step needs to process the full input at once.
+
+    If use_mini_batch is True, it runs a generalization of the mini-batch
+    K-means algorithm. It runs multiple iterations, where each iteration is
+    composed of mini_batch_steps_per_iteration steps. Two copies of cluster
+    centers are maintained: one that is updated at the end of each iteration,
+    and one that is updated every step. The first copy is used to compute
+    cluster allocations for each step, and for inference, while the second copy
+    is the one updated each step using the mini-batch update rule. After each
+    iteration is complete, this second copy is copied back the first copy.
+
+    Note that for use_mini_batch=True, when mini_batch_steps_per_iteration=1,
+    the algorithm reduces to the standard mini-batch algorithm. Also by setting
+    mini_batch_steps_per_iteration = num_inputs / batch_size, the algorithm
+    becomes an asynchronous version of the full-batch algorithm. Note however
+    that there is no guarantee by this implementation that each input is seen
+    exactly once per iteration. Also, different updates are applied
+    asynchronously without locking. So this asynchronous version may not behave
+    exactly like a full-batch version.
+
+    Args:
+      inputs: An input tensor or list of input tensors. It is assumed that the
+        data points have been previously randomly permuted.
+      num_clusters: An integer tensor specifying the number of clusters. This
+        argument is ignored if initial_clusters is a tensor or numpy array.
+      initial_clusters: Specifies the clusters used during initialization. One
+        of the following: - a tensor or numpy array with the initial cluster
+          centers. - a function f(inputs, k) that returns up to k centers from
+          `inputs`.
+        - "random": Choose centers randomly from `inputs`.
+        - "kmeans_plus_plus": Use kmeans++ to choose centers from `inputs`.
+        - "kmc2": Use the fast k-MC2 algorithm to choose centers from `inputs`.
+          In the last three cases, one batch of `inputs` may not yield
+          `num_clusters` centers, in which case initialization will require
+          multiple batches until enough centers are chosen. In the case of
+          "random" or "kmeans_plus_plus", if the input size is <= `num_clusters`
+          then the entire batch is chosen to be cluster centers.
+      distance_metric: Distance metric used for clustering. Supported options:
+        "squared_euclidean", "cosine".
+      use_mini_batch: If true, use the mini-batch k-means algorithm. Else assume
+        full batch.
+      mini_batch_steps_per_iteration: Number of steps after which the updated
+        cluster centers are synced back to a master copy.
+      random_seed: Seed for PRNG used to initialize seeds.
+      kmeans_plus_plus_num_retries: For each point that is sampled during
+        kmeans++ initialization, this parameter specifies the number of
+        additional points to draw from the current distribution before selecting
+        the best. If a negative value is specified, a heuristic is used to
+        sample O(log(num_to_sample)) additional points.
+      kmc2_chain_length: Determines how many candidate points are used by the
+        k-MC2 algorithm to produce one new cluster centers. If a (mini-)batch
+        contains less points, one new cluster center is generated from the
+        (mini-)batch.
+
+    Raises:
+      ValueError: An invalid argument was passed to initial_clusters or
+        distance_metric.
+    """
+    initialization_algorithms = [RANDOM_INIT, KMEANS_PLUS_PLUS_INIT, KMC2_INIT]
+    if isinstance(initial_clusters,
+                  str) and initial_clusters not in initialization_algorithms:
+      raise ValueError(
+          f'Unsupported initialization algorithm `{initial_clusters}`,'
+          f'must be one of `{initialization_algorithms}`.')
+
+    distance_metrics = [SQUARED_EUCLIDEAN_DISTANCE, COSINE_DISTANCE]
+    if distance_metric not in distance_metrics:
+      raise ValueError(f'Unsupported distance metric `{distance_metric}`,'
+                       f'must be one of `{distance_metrics}`.')
+    self._inputs = inputs if isinstance(inputs, list) else [inputs]
+    self._num_clusters = num_clusters
+    self._initial_clusters = initial_clusters
+    self._distance_metric = distance_metric
+    self._use_mini_batch = use_mini_batch
+    self._mini_batch_steps_per_iteration = int(mini_batch_steps_per_iteration)
+    self._seed = random_seed_ops.get_seed(random_seed)[0]
+    self._kmeans_plus_plus_num_retries = kmeans_plus_plus_num_retries
+    self._kmc2_chain_length = kmc2_chain_length
+
+  @classmethod
+  def _distance_graph(cls, inputs, clusters, distance_metric):
+    """Computes distance between each input and each cluster center.
+
+    Args:
+      inputs: list of input Tensors.
+      clusters: cluster Tensor.
+      distance_metric: distance metric used for clustering
+
+    Returns:
+      list of Tensors, where each element corresponds to each element in inputs.
+      The value is the distance of each row to all the cluster centers.
+      Currently only Euclidean distance and cosine distance are supported.
+    """
+    assert isinstance(inputs, list)
+    if distance_metric == SQUARED_EUCLIDEAN_DISTANCE:
+      return cls._compute_euclidean_distance(inputs, clusters)
+    elif distance_metric == COSINE_DISTANCE:
+      return cls._compute_cosine_distance(
+          inputs, clusters, inputs_normalized=True)
+    else:
+      assert False, str(distance_metric)
+
+  @classmethod
+  def _compute_euclidean_distance(cls, inputs, clusters):
+    """Computes Euclidean distance between each input and each cluster center.
+
+    Args:
+      inputs: list of input Tensors.
+      clusters: cluster Tensor.
+
+    Returns:
+      list of Tensors, where each element corresponds to each element in inputs.
+      The value is the distance of each row to all the cluster centers.
+    """
+    output = []
+    for inp in inputs:
+      with ops.colocate_with(inp, ignore_existing=True):
+        # Computes Euclidean distance. Note the first and third terms are
+        # broadcast additions.
+        squared_distance = (
+            math_ops.reduce_sum(math_ops.square(inp), 1, keepdims=True) -
+            2 * math_ops.matmul(inp, clusters, transpose_b=True) +
+            array_ops.transpose(
+                math_ops.reduce_sum(
+                    math_ops.square(clusters), 1, keepdims=True)))
+        output.append(squared_distance)
+
+    return output
+
+  @classmethod
+  def _compute_cosine_distance(cls, inputs, clusters, inputs_normalized=True):
+    """Computes cosine distance between each input and each cluster center.
+
+    Args:
+      inputs: list of input Tensor.
+      clusters: cluster Tensor
+      inputs_normalized: if True, it assumes that inp and clusters are
+        normalized and computes the dot product which is equivalent to the
+        cosine distance. Else it L2 normalizes the inputs first.
+
+    Returns:
+      list of Tensors, where each element corresponds to each element in inp.
+      The value is the distance of each row to all the cluster centers.
+    """
+    output = []
+    if not inputs_normalized:
+      with ops.colocate_with(clusters, ignore_existing=True):
+        clusters = nn_impl.l2_normalize(clusters, axis=1)
+    for inp in inputs:
+      with ops.colocate_with(inp, ignore_existing=True):
+        if not inputs_normalized:
+          inp = nn_impl.l2_normalize(inp, axis=1)
+        output.append(1 - math_ops.matmul(inp, clusters, transpose_b=True))
+    return output
+
+  def _infer_graph(self, inputs, clusters):
+    """Maps input to closest cluster and the score.
+
+    Args:
+      inputs: list of input Tensors.
+      clusters: Tensor of cluster centers.
+
+    Returns:
+      List of tuple, where each value in tuple corresponds to a value in inp.
+      The tuple has following three elements:
+      all_scores: distance of each input to each cluster center.
+      score: distance of each input to closest cluster center.
+      cluster_idx: index of cluster center closest to the corresponding input.
+    """
+    assert isinstance(inputs, list)
+    # Pairwise distances are used only by transform(). In all other cases, this
+    # sub-graph is not evaluated.
+    scores = self._distance_graph(inputs, clusters, self._distance_metric)
+    output = []
+    if (self._distance_metric == COSINE_DISTANCE and
+        not self._clusters_l2_normalized()):
+      # The cosine distance between normalized vectors x and y is the same as
+      # 2 * squared_euclidean_distance. We are using this fact and reusing the
+      # nearest_neighbors op.
+      # TODO(ands): Support COSINE distance in nearest_neighbors and remove
+      # this.
+      with ops.colocate_with(clusters, ignore_existing=True):
+        clusters = nn_impl.l2_normalize(clusters, axis=1)
+    for inp, score in zip(inputs, scores):
+      with ops.colocate_with(inp, ignore_existing=True):
+        (indices,
+         distances) = gen_clustering_ops.nearest_neighbors(inp, clusters, 1)
+        if self._distance_metric == COSINE_DISTANCE:
+          distances *= 0.5
+        output.append(
+            (score, array_ops.squeeze(distances,
+                                      [-1]), array_ops.squeeze(indices, [-1])))
+    return zip(*output)
+
+  def _clusters_l2_normalized(self):
+    """Returns True if clusters centers are kept normalized."""
+    return (self._distance_metric == COSINE_DISTANCE and
+            (not self._use_mini_batch or
+             self._mini_batch_steps_per_iteration > 1))
+
+  def _create_variables(self, num_clusters):
+    """Creates variables.
+
+    Args:
+      num_clusters: an integer Tensor providing the number of clusters.
+
+    Returns:
+      Tuple with following elements:
+      - cluster_centers: a Tensor for storing cluster centers
+      - cluster_centers_initialized: bool Variable indicating whether clusters
+            are initialized.
+      - cluster_counts: a Tensor for storing counts of points assigned to this
+            cluster. This is used by mini-batch training.
+      - cluster_centers_updated: Tensor representing copy of cluster centers
+            that are updated every step.
+      - update_in_steps: numbers of steps left before we sync
+            cluster_centers_updated back to cluster_centers.
+    """
+    init_value = array_ops.placeholder_with_default([], shape=None)
+    cluster_centers = variable_v1.VariableV1(
+        init_value, name=CLUSTERS_VAR_NAME, validate_shape=False)
+    cluster_centers_initialized = variable_v1.VariableV1(
+        False, dtype=dtypes.bool, name='initialized')
+
+    if self._use_mini_batch and self._mini_batch_steps_per_iteration > 1:
+      # Copy of cluster centers actively updated each step according to
+      # mini-batch update rule.
+      cluster_centers_updated = variable_v1.VariableV1(
+          init_value, name='clusters_updated', validate_shape=False)
+      # How many steps till we copy the updated clusters to cluster_centers.
+      update_in_steps = variable_v1.VariableV1(
+          self._mini_batch_steps_per_iteration,
+          dtype=dtypes.int64,
+          name='update_in_steps')
+      # Count of points assigned to cluster_centers_updated.
+      cluster_counts = variable_v1.VariableV1(
+          array_ops.zeros([num_clusters], dtype=dtypes.int64))
+    else:
+      cluster_centers_updated = cluster_centers
+      update_in_steps = None
+      cluster_counts = (
+          variable_v1.VariableV1(
+              array_ops.ones([num_clusters], dtype=dtypes.int64))
+          if self._use_mini_batch else None)
+    return (cluster_centers, cluster_centers_initialized, cluster_counts,
+            cluster_centers_updated, update_in_steps)
+
+  @classmethod
+  def _l2_normalize_data(cls, inputs):
+    """Normalized the input data."""
+    output = []
+    for inp in inputs:
+      with ops.colocate_with(inp, ignore_existing=True):
+        output.append(nn_impl.l2_normalize(inp, dim=1))
+    return output
+
+  def training_graph(self):
+    """Generate a training graph for kmeans algorithm.
+
+    This returns, among other things, an op that chooses initial centers
+    (init_op), a boolean variable that is set to True when the initial centers
+    are chosen (cluster_centers_initialized), and an op to perform either an
+    entire Lloyd iteration or a mini-batch of a Lloyd iteration (training_op).
+    The caller should use these components as follows. A single worker should
+    execute init_op multiple times until cluster_centers_initialized becomes
+    True. Then multiple workers may execute training_op any number of times.
+
+    Returns:
+      A tuple consisting of:
+      all_scores: A matrix (or list of matrices) of dimensions (num_input,
+        num_clusters) where the value is the distance of an input vector and a
+        cluster center.
+      cluster_idx: A vector (or list of vectors). Each element in the vector
+        corresponds to an input row in 'inp' and specifies the cluster id
+        corresponding to the input.
+      scores: Similar to cluster_idx but specifies the distance to the
+        assigned cluster instead.
+      cluster_centers_initialized: scalar indicating whether clusters have been
+        initialized.
+      init_op: an op to initialize the clusters.
+      training_op: an op that runs an iteration of training.
+    """
+    # Implementation of kmeans.
+    if (isinstance(self._initial_clusters, str) or
+        callable(self._initial_clusters)):
+      initial_clusters = self._initial_clusters
+      num_clusters = ops.convert_to_tensor(self._num_clusters)
+    else:
+      initial_clusters = ops.convert_to_tensor(self._initial_clusters)
+      num_clusters = array_ops.shape(initial_clusters)[0]
+
+    inputs = self._inputs
+    (cluster_centers_var, cluster_centers_initialized, total_counts,
+     cluster_centers_updated,
+     update_in_steps) = self._create_variables(num_clusters)
+    init_op = _InitializeClustersOpFactory(
+        self._inputs, num_clusters, initial_clusters, self._distance_metric,
+        self._seed, self._kmeans_plus_plus_num_retries, self._kmc2_chain_length,
+        cluster_centers_var, cluster_centers_updated,
+        cluster_centers_initialized).op()
+    cluster_centers = cluster_centers_var
+
+    if self._distance_metric == COSINE_DISTANCE:
+      inputs = self._l2_normalize_data(inputs)
+      if not self._clusters_l2_normalized():
+        cluster_centers = nn_impl.l2_normalize(cluster_centers, dim=1)
+
+    all_scores, scores, cluster_idx = self._infer_graph(inputs, cluster_centers)
+    if self._use_mini_batch:
+      sync_updates_op = self._mini_batch_sync_updates_op(
+          update_in_steps, cluster_centers_var, cluster_centers_updated,
+          total_counts)
+      assert sync_updates_op is not None
+      with ops.control_dependencies([sync_updates_op]):
+        training_op = self._mini_batch_training_op(inputs, cluster_idx,
+                                                   cluster_centers_updated,
+                                                   total_counts)
+    else:
+      assert cluster_centers == cluster_centers_var
+      training_op = self._full_batch_training_op(inputs, num_clusters,
+                                                 cluster_idx,
+                                                 cluster_centers_var)
+
+    return (all_scores, cluster_idx, scores, cluster_centers_initialized,
+            init_op, training_op)
+
+  def _mini_batch_sync_updates_op(self, update_in_steps, cluster_centers_var,
+                                  cluster_centers_updated, total_counts):
+    if self._use_mini_batch and self._mini_batch_steps_per_iteration > 1:
+      assert update_in_steps is not None
+      with ops.colocate_with(update_in_steps, ignore_existing=True):
+
+        def _f():
+          # Note that there is a race condition here, so we do a best effort
+          # updates here. We reset update_in_steps first so that other workers
+          # don't duplicate the updates. Also we update cluster_center_vars
+          # before resetting total_counts to avoid large updates to
+          # cluster_centers_updated based on partially updated
+          # cluster_center_vars.
+          with ops.control_dependencies([
+              state_ops.assign(update_in_steps,
+                               self._mini_batch_steps_per_iteration - 1)
+          ]):
+            with ops.colocate_with(
+                cluster_centers_updated, ignore_existing=True):
+              if self._distance_metric == COSINE_DISTANCE:
+                cluster_centers = nn_impl.l2_normalize(
+                    cluster_centers_updated, dim=1)
+              else:
+                cluster_centers = cluster_centers_updated
+            with ops.colocate_with(cluster_centers_var, ignore_existing=True):
+              with ops.control_dependencies(
+                  [state_ops.assign(cluster_centers_var, cluster_centers)]):
+                with ops.colocate_with(None, ignore_existing=True):
+                  with ops.control_dependencies([
+                      state_ops.assign(total_counts,
+                                       array_ops.zeros_like(total_counts))
+                  ]):
+                    return array_ops.identity(update_in_steps)
+
+        return cond.cond(
+            update_in_steps <= 0, _f,
+            lambda: state_ops.assign_sub(update_in_steps, 1))
+    else:
+      return control_flow_ops.no_op()
+
+  def _mini_batch_training_op(self, inputs, cluster_idx_list, cluster_centers,
+                              total_counts):
+    """Creates an op for training for mini batch case.
+
+    Args:
+      inputs: list of input Tensors.
+      cluster_idx_list: A vector (or list of vectors). Each element in the
+        vector corresponds to an input row in 'inp' and specifies the cluster id
+        corresponding to the input.
+      cluster_centers: Tensor Ref of cluster centers.
+      total_counts: Tensor Ref of cluster counts.
+
+    Returns:
+      An op for doing an update of mini-batch k-means.
+    """
+    update_ops = []
+    for inp, cluster_idx in zip(inputs, cluster_idx_list):
+      with ops.colocate_with(inp, ignore_existing=True):
+        assert total_counts is not None
+        cluster_idx = array_ops.reshape(cluster_idx, [-1])
+        # Dedupe the unique ids of cluster_centers being updated so that updates
+        # can be locally aggregated.
+        unique_ids, unique_idx = array_ops.unique(cluster_idx)
+        num_unique_cluster_idx = array_ops.size(unique_ids)
+        # Fetch the old values of counts and cluster_centers.
+        with ops.colocate_with(total_counts, ignore_existing=True):
+          old_counts = array_ops.gather(total_counts, unique_ids)
+        # TODO(agarwal): This colocation seems to run into problems. Fix it.
+        with ops.colocate_with(cluster_centers, ignore_existing=True):
+          old_cluster_centers = array_ops.gather(cluster_centers, unique_ids)
+        # Locally aggregate the increment to counts.
+        count_updates = math_ops.unsorted_segment_sum(
+            array_ops.ones_like(unique_idx, dtype=total_counts.dtype),
+            unique_idx, num_unique_cluster_idx)
+        # Locally compute the sum of inputs mapped to each id.
+        # For a cluster with old cluster value x, old count n, and with data
+        # d_1,...d_k newly assigned to it, we recompute the new value as
+        # \\(x += (sum_i(d_i) - k * x) / (n + k)\\).
+        # Compute \\(sum_i(d_i)\\), see comment above.
+        cluster_center_updates = math_ops.unsorted_segment_sum(
+            inp, unique_idx, num_unique_cluster_idx)
+        # Shape to enable broadcasting count_updates and learning_rate to inp.
+        # It extends the shape with 1's to match the rank of inp.
+        broadcast_shape = array_ops.concat([
+            array_ops.reshape(num_unique_cluster_idx, [1]),
+            array_ops.ones(
+                array_ops.reshape(array_ops.rank(inp) - 1, [1]),
+                dtype=dtypes.int32)
+        ], 0)
+        # Subtract k * x, see comment above.
+        cluster_center_updates -= math_ops.cast(
+            array_ops.reshape(count_updates, broadcast_shape),
+            inp.dtype) * old_cluster_centers
+        learning_rate = math_ops.reciprocal(
+            math_ops.cast(old_counts + count_updates, inp.dtype))
+        learning_rate = array_ops.reshape(learning_rate, broadcast_shape)
+        # scale by 1 / (n + k), see comment above.
+        cluster_center_updates *= learning_rate
+        # Apply the updates.
+      update_counts = state_ops.scatter_add(total_counts, unique_ids,
+                                            count_updates)
+      update_cluster_centers = state_ops.scatter_add(cluster_centers,
+                                                     unique_ids,
+                                                     cluster_center_updates)
+      update_ops.extend([update_counts, update_cluster_centers])
+    return control_flow_ops.group(*update_ops)
+
+  def _full_batch_training_op(self, inputs, num_clusters, cluster_idx_list,
+                              cluster_centers):
+    """Creates an op for training for full batch case.
+
+    Args:
+      inputs: list of input Tensors.
+      num_clusters: an integer Tensor providing the number of clusters.
+      cluster_idx_list: A vector (or list of vectors). Each element in the
+        vector corresponds to an input row in 'inp' and specifies the cluster id
+        corresponding to the input.
+      cluster_centers: Tensor Ref of cluster centers.
+
+    Returns:
+      An op for doing an update of mini-batch k-means.
+    """
+    cluster_sums = []
+    cluster_counts = []
+    epsilon = constant_op.constant(1e-6, dtype=inputs[0].dtype)
+    for inp, cluster_idx in zip(inputs, cluster_idx_list):
+      with ops.colocate_with(inp, ignore_existing=True):
+        cluster_sums.append(
+            math_ops.unsorted_segment_sum(inp, cluster_idx, num_clusters))
+        cluster_counts.append(
+            math_ops.unsorted_segment_sum(
+                array_ops.reshape(
+                    array_ops.ones(
+                        array_ops.reshape(array_ops.shape(inp)[0], [-1])),
+                    [-1, 1]), cluster_idx, num_clusters))
+    with ops.colocate_with(cluster_centers, ignore_existing=True):
+      new_clusters_centers = math_ops.add_n(cluster_sums) / (
+          math_ops.cast(math_ops.add_n(cluster_counts), cluster_sums[0].dtype) +
+          epsilon)
+      if self._clusters_l2_normalized():
+        new_clusters_centers = nn_impl.l2_normalize(new_clusters_centers, dim=1)
+    return state_ops.assign(cluster_centers, new_clusters_centers)
+
+
+class _InitializeClustersOpFactory:
+  """Internal class to create the op to initialize the clusters.
+
+    The op performs this algorithm (see constructor args):
+
+    num_remaining = num_clusters - length(cluster_centers)
+    if num_remaining == 0:
+      assert that cluster_centers_initialized is true
+    else:
+      assert that num_remaining > 0
+      new_centers = choose up to num_remaining initial centers
+      l2-normalize new_centers if using cosine distance
+      all_centers = concat(cluster_centers, new_centers)
+      cluster_centers := all_centers
+      if there is a cluster_centers_updated variable:
+        cluster_centers_updated := cluster_centers
+      num_now_remaining = num_clusters - length(cluster_centers)
+      if num_now_remaining == 0:
+        cluster_centers_initialized := true
+  """
+
+  # TODO(ccolby): Refactor this class so that kmc2 isn't so much a special case.
+
+  def __init__(self, inputs, num_clusters, initial_clusters, distance_metric,
+               random_seed, kmeans_plus_plus_num_retries, kmc2_chain_length,
+               cluster_centers, cluster_centers_updated,
+               cluster_centers_initialized):
+    """Creates an op factory.
+
+    Args:
+      inputs: See KMeans constructor.
+      num_clusters: An integer Tensor providing the number of clusters.
+      initial_clusters: See KMeans constructor.
+      distance_metric: See KMeans constructor.
+      random_seed: See KMeans constructor.
+      kmeans_plus_plus_num_retries: See KMeans constructor.
+      kmc2_chain_length: See KMeans constructor.
+      cluster_centers: The TF variable holding the initial centers. It may
+        already contain some centers when the op is executed.
+      cluster_centers_updated: A second TF variable to hold a copy of the
+        initial centers, used for full-batch mode. In mini-batch mode,
+        cluster_centers_updated is the same variable as cluster_centers.
+      cluster_centers_initialized: A boolean TF variable that will be set to
+        true when all the initial centers have been chosen.
+    """
+    # All of these instance variables are constants.
+    self._inputs = inputs
+    self._num_clusters = num_clusters
+    self._initial_clusters = initial_clusters
+    self._distance_metric = distance_metric
+    self._seed = random_seed
+    self._kmeans_plus_plus_num_retries = kmeans_plus_plus_num_retries
+    self._kmc2_chain_length = kmc2_chain_length
+    self._cluster_centers = cluster_centers
+    self._cluster_centers_updated = cluster_centers_updated
+    self._cluster_centers_initialized = cluster_centers_initialized
+
+    self._num_selected = array_ops.shape(self._cluster_centers)[0]
+    self._num_remaining = self._num_clusters - self._num_selected
+    self._num_data = math_ops.add_n(
+        [array_ops.shape(i)[0] for i in self._inputs])
+
+  def _random(self):
+    indices = random_ops.random_uniform(
+        array_ops.reshape(self._num_remaining, [-1]),
+        minval=0,
+        maxval=math_ops.cast(self._num_data, dtypes.int64),
+        seed=self._seed,
+        dtype=dtypes.int64)
+    return embedding_lookup(self._inputs, indices, partition_strategy='div')
+
+  def _kmeans_plus_plus(self):
+    # Points from only the first shard are used for initializing centers.
+    # TODO(ands): Use all points.
+    inp = self._inputs[0]
+    if self._distance_metric == COSINE_DISTANCE:
+      inp = nn_impl.l2_normalize(inp, dim=1)
+    return gen_clustering_ops.kmeans_plus_plus_initialization(
+        inp, math_ops.cast(self._num_remaining, dtypes.int64), self._seed,
+        self._kmeans_plus_plus_num_retries)
+
+  def _kmc2_multiple_centers(self):
+    """Adds new initial cluster centers using the k-MC2 algorithm.
+
+    In each call to the op, the provided batch is split into subsets based on
+    the specified `kmc2_chain_length`. On each subset, a single Markov chain of
+    the k-MC2 algorithm is used to add *one* new center cluster center. If there
+    are less than `kmc2_chain_length` points in the subset, a single center is
+    added using one Markov chain on the full input. It is assumed that the
+    provided batch has previously been randomly permuted. Otherwise, k-MC2 may
+    return suboptimal centers.
+
+    Returns:
+      An op that adds new cluster centers.
+    """
+    # The op only operates on the first shard of data.
+    first_shard = self._inputs[0]
+    # Number of points in the input that can be used.
+    batch_size = array_ops.shape(first_shard)[0]
+    # Maximum number of subsets such that the size of each subset is at least
+    # `kmc2_chain_length`. Final subsets may be larger.
+    max_to_sample = math_ops.cast(
+        batch_size / self._kmc2_chain_length, dtype=dtypes.int32)
+    # We sample at least one new center and at most all remaining centers.
+    num_to_sample = math_ops.maximum(
+        math_ops.minimum(self._num_remaining, max_to_sample), 1)
+
+    def _cond(i, _):
+      """Stopping condition for the while loop."""
+      return math_ops.less(i, num_to_sample)
+
+    def _body(i, _):
+      """Body that adds a single new center based on a subset."""
+
+      def _sample_random():
+        """Returns a random point as a cluster center."""
+        # By assumption the batch is reshuffled and _sample_random is always
+        # called for i=0. Hence, we simply return the first point.
+        new_center = array_ops.reshape(first_shard[0], [1, -1])
+        if self._distance_metric == COSINE_DISTANCE:
+          new_center = nn_impl.l2_normalize(new_center, dim=1)
+        return new_center
+
+      def _sample_kmc2_chain():
+        """Returns previous centers as well as a new center sampled using k-MC2."""
+        # Extract the subset from the underlying batch.
+        start = i * self._kmc2_chain_length
+        end = start + self._kmc2_chain_length
+        subset = first_shard[start:end]
+        # Compute the distances from points in the subset to previous centers.
+        _, distances = gen_clustering_ops.nearest_neighbors(
+            subset, self._cluster_centers, 1)
+        # Sample index of new center using k-MC2 Markov chain.
+        new_center_index = gen_clustering_ops.kmc2_chain_initialization(
+            array_ops.squeeze(distances), self._seed)
+        # Extract actual new center.
+        newly_sampled_center = array_ops.reshape(subset[new_center_index],
+                                                 [1, -1])
+        # Return concatenation with previously sampled centers.
+        if self._distance_metric == COSINE_DISTANCE:
+          newly_sampled_center = nn_impl.l2_normalize(
+              newly_sampled_center, dim=1)
+        return array_ops.concat([self._cluster_centers, newly_sampled_center],
+                                0)
+
+      # Obtain a random point if there are no previously sampled centers.
+      # Otherwise, construct a k-MC2 Markov chain.
+      new_centers = cond.cond(
+          math_ops.equal(self._num_selected, 0), _sample_random,
+          _sample_kmc2_chain)
+      # Assign new cluster centers to underlying variable.
+      assigned_centers = state_ops.assign(
+          self._cluster_centers, new_centers, validate_shape=False)
+      if self._cluster_centers_updated is not self._cluster_centers:
+        assigned_centers = state_ops.assign(
+            self._cluster_centers_updated,
+            assigned_centers,
+            validate_shape=False)
+      return i + 1, self._num_clusters - array_ops.shape(assigned_centers)[0]
+
+    # Add num_to_sample new data points.
+    _, num_remaining = while_loop.while_loop(_cond, _body, [0, 0])
+    return num_remaining
+
+  def _greedy_batch_sampler(self, sampler):
+    # If the input dataset size is smaller than the number of centers
+    # remaining, choose the entire input dataset as centers. This can happen
+    # with mini-batch. Otherwise, sample the batch according to the provided
+    # sampler.
+    return cond.cond(self._num_data <= self._num_remaining,
+                     lambda: array_ops.concat(self._inputs, 0),
+                     sampler)
+
+  def _single_batch_sampler(self, sampler):
+    # Enforce that there are at least as many data points as centers
+    # remaining. This gives the provided sampler the chance to select all
+    # remaining centers from a single batch.
+    with ops.control_dependencies(
+        [check_ops.assert_greater_equal(self._num_data, self._num_remaining)]):
+      return sampler()
+
+  def _choose_initial_centers(self):
+    if isinstance(self._initial_clusters, str):
+      if self._initial_clusters == RANDOM_INIT:
+        return self._greedy_batch_sampler(self._random)
+      else:  # self._initial_clusters == KMEANS_PLUS_PLUS_INIT
+        return self._single_batch_sampler(self._kmeans_plus_plus)
+    elif callable(self._initial_clusters):
+      return self._initial_clusters(self._inputs, self._num_remaining)
+    else:
+      with ops.control_dependencies([
+          check_ops.assert_equal(self._num_remaining,
+                                 array_ops.shape(self._initial_clusters)[0])
+      ]):
+        return self._initial_clusters
+
+  def _add_new_centers(self):
+    """Adds some centers and returns the number of centers remaining."""
+    new_centers = self._choose_initial_centers()
+    if self._distance_metric == COSINE_DISTANCE:
+      new_centers = nn_impl.l2_normalize(new_centers, dim=1)
+    # If cluster_centers is empty, it doesn't have the right shape for concat.
+    all_centers = cond.cond(
+        math_ops.equal(self._num_selected, 0), lambda: new_centers,
+        lambda: array_ops.concat([self._cluster_centers, new_centers], 0))
+    # TODO(ccolby): De-dupe all_centers?
+    a = state_ops.assign(
+        self._cluster_centers, all_centers, validate_shape=False)
+    if self._cluster_centers_updated is not self._cluster_centers:
+      a = state_ops.assign(
+          self._cluster_centers_updated, a, validate_shape=False)
+    return self._num_clusters - array_ops.shape(a)[0]
+
+  def _initialize(self):
+    with ops.control_dependencies([
+        check_ops.assert_positive(self._num_remaining),
+    ]):
+      if self._initial_clusters == KMC2_INIT:
+        num_now_remaining = self._kmc2_multiple_centers()
+      else:
+        num_now_remaining = self._add_new_centers()
+      return cond.cond(
+          math_ops.equal(num_now_remaining, 0),
+          lambda: state_ops.assign(self._cluster_centers_initialized, True),
+          control_flow_ops.no_op)
+
+  def op(self):
+    """Returns the cluster initializer op."""
+    return cond.cond(
+        math_ops.equal(self._num_remaining, 0),
+        lambda: check_ops.assert_equal(self._cluster_centers_initialized, True),
+        self._initialize)

tensorflow/python/ops/clustering_ops_test.py

@@ -0,0 +1,208 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not
+# use this file except in compliance with the License.  You may obtain a copy of
+# the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+# License for the specific language governing permissions and limitations under
+# the License.
+# ==============================================================================
+"""Tests for clustering_ops."""
+
+import numpy as np
+
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import clustering_ops
+from tensorflow.python.platform import test
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class KmeansPlusPlusInitializationTest(test.TestCase):
+
+  # All but one input point are close to (101, 1). With uniform random sampling,
+  # it is highly improbable for (-1, -1) to be selected.
+  def setUp(self):
+    self._points = np.array([[100., 0.],
+                             [101., 2.],
+                             [102., 0.],
+                             [100., 1.],
+                             [100., 2.],
+                             [101., 0.],
+                             [101., 0.],
+                             [101., 1.],
+                             [102., 0.],
+                             [-1., -1.]]).astype(np.float32)
+
+  def runTestWithSeed(self, seed):
+    with self.cached_session():
+      sampled_points = clustering_ops.kmeans_plus_plus_initialization(
+          self._points, 3, seed, (seed % 5) - 1)
+      self.assertAllClose(
+          sorted(self.evaluate(sampled_points).tolist()),
+          [[-1., -1.], [101., 1.], [101., 1.]],
+          atol=1.0)
+
+  def testBasic(self):
+    for seed in range(100):
+      self.runTestWithSeed(seed)
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class KMC2InitializationTest(test.TestCase):
+
+  def runTestWithSeed(self, seed):
+    with self.cached_session():
+      distances = np.zeros(1000).astype(np.float32)
+      distances[6] = 10e7
+      distances[4] = 10e3
+
+      sampled_point = clustering_ops.kmc2_chain_initialization(distances, seed)
+      self.assertAllEqual(sampled_point, 6)
+      distances[6] = 0.0
+      sampled_point = clustering_ops.kmc2_chain_initialization(distances, seed)
+      self.assertAllEqual(sampled_point, 4)
+
+  def testBasic(self):
+    for seed in range(100):
+      self.runTestWithSeed(seed)
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class KMC2InitializationLargeTest(test.TestCase):
+
+  def setUp(self):
+    self._distances = np.zeros(1001)
+    self._distances[500] = 100.0
+    self._distances[1000] = 50.0
+
+  def testBasic(self):
+    with self.cached_session():
+      counts = {}
+      seed = 0
+      for i in range(50):
+        sample = self.evaluate(
+            clustering_ops.kmc2_chain_initialization(self._distances, seed + i))
+        counts[sample] = counts.get(sample, 0) + 1
+      self.assertEqual(len(counts), 2)
+      self.assertTrue(500 in counts)
+      self.assertTrue(1000 in counts)
+      self.assertGreaterEqual(counts[500], 5)
+      self.assertGreaterEqual(counts[1000], 5)
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class KMC2InitializationCornercaseTest(test.TestCase):
+
+  def setUp(self):
+    self._distances = np.zeros(10)
+
+  def runTestWithSeed(self, seed):
+    with self.cached_session():
+      sampled_point = clustering_ops.kmc2_chain_initialization(
+          self._distances, seed)
+      self.assertAllEqual(sampled_point, 0)
+
+  def testBasic(self):
+    for seed in range(100):
+      self.runTestWithSeed(seed)
+
+
+@test_util.run_all_in_graph_and_eager_modes
+# A simple test that can be verified by hand.
+class NearestCentersTest(test.TestCase):
+
+  def setUp(self):
+    self._points = np.array([[100., 0.],
+                             [101., 2.],
+                             [99., 2.],
+                             [1., 1.]]).astype(np.float32)
+
+    self._centers = np.array([[100., 0.],
+                              [99., 1.],
+                              [50., 50.],
+                              [0., 0.],
+                              [1., 1.]]).astype(np.float32)
+
+  def testNearest1(self):
+    with self.cached_session():
+      [indices, distances] = clustering_ops.nearest_neighbors(self._points,
+                                                              self._centers, 1)
+      self.assertAllClose(indices, [[0], [0], [1], [4]])
+      self.assertAllClose(distances, [[0.], [5.], [1.], [0.]])
+
+  def testNearest2(self):
+    with self.cached_session():
+      [indices, distances] = clustering_ops.nearest_neighbors(self._points,
+                                                              self._centers, 2)
+      self.assertAllClose(indices, [[0, 1], [0, 1], [1, 0], [4, 3]])
+      self.assertAllClose(distances, [[0., 2.], [5., 5.], [1., 5.], [0., 2.]])
+
+
+@test_util.run_all_in_graph_and_eager_modes
+# A test with large inputs.
+class NearestCentersLargeTest(test.TestCase):
+
+  def setUp(self):
+    num_points = 1000
+    num_centers = 2000
+    num_dim = 100
+    max_k = 5
+    # Construct a small number of random points and later tile them.
+    points_per_tile = 10
+    assert num_points % points_per_tile == 0
+    points = np.random.standard_normal(
+        [points_per_tile, num_dim]).astype(np.float32)
+    # Construct random centers.
+    self._centers = np.random.standard_normal(
+        [num_centers, num_dim]).astype(np.float32)
+
+    # Exhaustively compute expected nearest neighbors.
+    def squared_distance(x, y):
+      return np.linalg.norm(x - y, ord=2)**2
+
+    nearest_neighbors = [
+        sorted([(squared_distance(point, self._centers[j]), j)
+                for j in range(num_centers)])[:max_k] for point in points
+    ]
+    expected_nearest_neighbor_indices = np.array(
+        [[i for _, i in nn] for nn in nearest_neighbors])
+    expected_nearest_neighbor_squared_distances = np.array(
+        [[dist for dist, _ in nn] for nn in nearest_neighbors])
+    # Tile points and expected results to reach requested size (num_points)
+    (self._points, self._expected_nearest_neighbor_indices,
+     self._expected_nearest_neighbor_squared_distances) = (
+         np.tile(x, (int(num_points / points_per_tile), 1))
+         for x in (points, expected_nearest_neighbor_indices,
+                   expected_nearest_neighbor_squared_distances))
+
+  def testNearest1(self):
+    with self.cached_session():
+      [indices, distances] = clustering_ops.nearest_neighbors(self._points,
+                                                              self._centers, 1)
+      self.assertAllClose(
+          indices,
+          self._expected_nearest_neighbor_indices[:, [0]])
+      self.assertAllClose(
+          distances,
+          self._expected_nearest_neighbor_squared_distances[:, [0]])
+
+  def testNearest5(self):
+    with self.cached_session():
+      [indices, distances] = clustering_ops.nearest_neighbors(self._points,
+                                                              self._centers, 5)
+      self.assertAllClose(
+          indices,
+          self._expected_nearest_neighbor_indices[:, 0:5])
+      self.assertAllClose(
+          distances,
+          self._expected_nearest_neighbor_squared_distances[:, 0:5])
+
+
+if __name__ == "__main__":
+  np.random.seed(0)
+  test.main()

tensorflow/python/ops/collective_ops.py

@@ -0,0 +1,578 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""TensorFlow collective Ops."""
+from tensorflow.python.ops import gen_collective_ops
+
+
+def all_reduce(t,
+               group_size,
+               group_key,
+               instance_key,
+               merge_op='Add',
+               final_op='Id',
+               subdiv_offsets=(0,),
+               communication_hint='auto',
+               timeout=0):
+  """Reduces tensors collectively, across devices.
+
+  Args:
+    t: the tensor to be reduced.
+    group_size: the total number of tensors to be collectively reduced.
+      Each must reside on a different device.  Should be a positive integer.
+    group_key: an integer identifying the group of devices.
+    instance_key: an integer identifying the participating group of Ops.
+    merge_op: string naming the binary Op to be applied to compute each
+      partial reduction.
+    final_op: string naming the unary Op to be applied to each fully
+      reduced value.  Can be 'Id' for no operation.
+    subdiv_offsets: a list of integer offsets into the tensor at which each
+      independent subdivision should begin.  Use [0] if no subdivision should
+      be done.
+    communication_hint: preferred collective communication.  The implementation
+      may fall back to another mechanism.  Options include `auto`, `ring`, and
+      `nccl`.
+    timeout: a float. If set to a non zero, set a completion timeout to detect
+      staleness.  If the timer goes off, a DeadlineExceededError is raised.  The
+      timeout value in seconds. This feature is experimental.
+
+  Returns:
+    An Op implementing the distributed reduction.
+
+  Raises:
+    ValueError: if any of the input parameter constraints are not met.
+  """
+  if group_size < 1:
+    raise ValueError('Parameter `group_size` to all_reduce must be at least 1. '
+                     f'Received: {group_size}.')
+  return gen_collective_ops.collective_reduce(
+      t,
+      group_size=group_size,
+      group_key=group_key,
+      instance_key=instance_key,
+      merge_op=merge_op,
+      final_op=final_op,
+      subdiv_offsets=subdiv_offsets,
+      communication_hint=communication_hint.lower(),
+      timeout_seconds=timeout)
+
+
+def assign_group_v2(group_assignment, device_index, base_key):
+  """Assign group key based on group_assignment.
+
+  Args:
+    group_assignment: a 2 dimensional integer Tensor that encodes which devices
+      belong to the same group. The values are indices of the devices within 0
+      to number of devices.
+    device_index: integer for the index of the current device
+    base_key: integer to offset the resulted group_key. The base key shall be
+      unique for different values of group_assignment in the same tf.function.
+  Notes: The device_index argument must be consistent with the index of the
+    device of this Op in the device assignment list. The behavior of this Op is
+    undefined if they are inconsistent.
+
+  Returns:
+    group_size, group_key: The group size and group key for the current device.
+  """
+  group_size, group_key = gen_collective_ops.collective_assign_group_v2(
+      group_assignment=group_assignment,
+      device_index=device_index,
+      base_key=base_key)
+  return group_size, group_key
+
+
+def all_reduce_v2(t,
+                  group_size,
+                  group_key,
+                  instance_key,
+                  merge_op='Add',
+                  final_op='Id',
+                  communication_hint='auto',
+                  timeout=0,
+                  ordering_token=None,
+                  max_subdivs_per_device=-1,
+                  name=None):
+  """Reduces tensors collectively, across devices.
+
+  Args:
+    t: the tensor to be reduced.
+    group_size: an int32 tensor. The total number of tensors to be collectively
+      reduced.  Each must reside on a different device.  Should be a positive
+      integer.
+    group_key: an int32 tensor identifying the group of devices.
+    instance_key: an int32 tensor identifying the participating group of Ops.
+    merge_op: string naming the binary Op to be applied to compute each partial
+      reduction.
+    final_op: string naming the unary Op to be applied to each fully reduced
+      value.  Can be 'Id' for no operation.
+    communication_hint: preferred collective communication.  The implementation
+      may fall back to another mechanism.  Options include `auto`, `ring`, and
+      `nccl`.
+    timeout: a float. If set to a non zero, set a completion timeout to detect
+      staleness.  If the timer goes off, a DeadlineExceededError is raised.  The
+      timeout value in seconds. This feature is experimental.
+    ordering_token: a resource tensor on the same device as the op to order the
+      collectives in a per-device manner by auto control dependency. This
+      argument can be omitted when there is one collective Op per `tf.function`,
+      or when explicit control dependency is used instead of auto control
+      dependency.
+    max_subdivs_per_device: int specifying the maximum number of subdivisions a
+      tensor on a device can be divided into. The runtime uses this constraint
+      to parallelize processing of each per-device tensor. Setting to -1
+      disables subdivision and reverts to previous behavior of not sub-dividing
+      tensor. Setting to 0 uses system defaults.
+    name: name of the Op.
+
+  Returns:
+    An Op implementing the distributed reduction.
+  """
+  if ordering_token is not None:
+    ordering_token = [ordering_token]
+  else:
+    ordering_token = []
+
+  return gen_collective_ops.collective_reduce_v2(
+      t,
+      group_size=group_size,
+      group_key=group_key,
+      instance_key=instance_key,
+      merge_op=merge_op,
+      final_op=final_op,
+      communication_hint=communication_hint.lower(),
+      timeout_seconds=timeout,
+      is_stateless=False,
+      ordering_token=ordering_token,
+      max_subdivs_per_device=max_subdivs_per_device,
+      name=name)
+
+
+def all_gather(t,
+               group_size,
+               group_key,
+               instance_key,
+               communication_hint='auto',
+               timeout=0):
+  """Accumulates tensors collectively, across devices, along first dimension.
+
+  Args:
+    t: the tensor to participate in the accumulation.
+    group_size: the total number of tensors to be collectively accumulated.
+      Each must reside on a different device. Should be a positive integer.
+    group_key: an integer identifying the group of devices.
+    instance_key: an integer identifying the participating group of Ops.
+    communication_hint: preferred collective communication. The implementation
+      may fall back to another mechanism. Options include `auto`, `ring`, and
+      `nccl`.
+    timeout: a float. If set to a non zero, set a completion timeout to detect
+      staleness. If the timer goes off, a DeadlineExceededError is raised. The
+      timeout value in seconds. This feature is experimental.
+
+  Returns:
+    An Op implementing the distributed operation.
+
+  Raises:
+    ValueError: if any of the input parameter constraints are not met.
+  """
+  if group_size < 1:
+    raise ValueError('Parameter `group_size` to all_gather must be at least 1.'
+                     f' Received: {group_size}.')
+  return gen_collective_ops.collective_gather(
+      t,
+      shape=[0],
+      group_size=group_size,
+      group_key=group_key,
+      instance_key=instance_key,
+      communication_hint=communication_hint.lower(),
+      timeout_seconds=timeout)
+
+
+def all_gather_v2(t,
+                  group_size,
+                  group_key,
+                  instance_key,
+                  communication_hint='auto',
+                  timeout=0,
+                  ordering_token=None,
+                  name=None):
+  """Accumulates tensors collectively, across devices, along first dimension.
+
+  Args:
+    t: the tensor to participate in the accumulation.
+    group_size: an int32 tensor, the total number of tensors to be collectively
+      accumulated. Each must reside on a different device. Should be a positive
+      integer.
+    group_key: an int32 tensor identifying the group of devices.
+    instance_key: an int32 tensor identifying the participating group of Ops.
+    communication_hint: preferred collective communication. The implementation
+      may fall back to another mechanism. Options include `auto`, `ring`, and
+      `nccl`.
+    timeout: a float. If set to a non zero, set a completion timeout to detect
+      staleness. If the timer goes off, a DeadlineExceededError is raised. The
+      timeout value in seconds. This feature is experimental.
+    ordering_token: a resource tensor on the same device as the op to order the
+      collectives in a per-device manner by auto control dependency. This
+      argument can be omitted when there is one collective Op per `tf.function`,
+      or when explicit control dependency is used instead of auto control
+      dependency.
+    name: name of the Op.
+
+  Returns:
+    An Op implementing the distributed operation.
+  """
+  if ordering_token is not None:
+    ordering_token = [ordering_token]
+  else:
+    ordering_token = []
+
+  return gen_collective_ops.collective_gather_v2(
+      t,
+      group_size=group_size,
+      group_key=group_key,
+      instance_key=instance_key,
+      communication_hint=communication_hint.lower(),
+      timeout_seconds=timeout,
+      is_stateless=False,
+      ordering_token=ordering_token,
+      name=name)
+
+
+def broadcast_send(t,
+                   shape,
+                   dtype,
+                   group_size,
+                   group_key,
+                   instance_key,
+                   communication_hint='auto',
+                   timeout=0):
+  """Broadcasts one tensor to a group of others, across devices.
+
+  Args:
+    t: the tensor to be sent.
+    shape: the shape of the tensor being sent, which must agree with t.
+    dtype: the type of the tensor being sent, which must agree with t.
+    group_size: one plus the number of receiving tensors, i.e. the total
+      number of devices participating.  Each tensor must reside on a
+      different device.
+    group_key: an integer identifying the group of devices.
+    instance_key: an integer identifying the participating group of Ops.
+    communication_hint: preferred collective communication.  The implementation
+      may fall back to another mechanism.  Options include `auto`, `ring`, and
+      `nccl`.
+    timeout: If set to a non zero, set a completion timeout to detect staleness.
+      If the timer goes off, a DeadlineExceededError is raised.
+      The timeout value in seconds. This feature is experimental.
+
+  Returns:
+    An Op implementing the distributed broadcast send.
+
+  Raises:
+    ValueError: if any of the input parameter constraints are not met.
+
+  Note that the shape and dtype arguments appear redundant since they
+  should be obtainable from t.  The are two reasons for including
+  them.  First, the shape and type of tensors passed via broadcast must
+  be known ahead of time in their most specific form so that the receive
+  side can allocate memory for the operation and shape/type inference can
+  carry forward from there.  Including the same declarations on the
+  send side clarifies a commitment already made.  Secondly, having nearly
+  identical use syntax for send and receive sides may simplify tool-driven
+  generation of broadcast.
+  """
+  if group_size <= 1:
+    raise ValueError(
+        'Parameter `group_size` to broadcast_send must be at least 2. '
+        f'Received: {group_size}.')
+  if t.shape != shape:
+    raise ValueError(
+        'Shape of broadcast_send tensor `t` not equal to declared shape. '
+        f'Received {t.shape}, expected {shape}.')
+  if t.dtype != dtype:
+    raise ValueError(
+        'Type of broadcast_send tensor `t` not equal to declared type. '
+        f'Received {t.dtype}, expected {dtype}.')
+  return gen_collective_ops.collective_bcast_send(
+      t,
+      shape=shape,
+      group_size=group_size,
+      group_key=group_key,
+      instance_key=instance_key,
+      communication_hint=communication_hint.lower(),
+      timeout_seconds=timeout)
+
+
+def broadcast_send_v2(t,
+                      group_size,
+                      group_key,
+                      instance_key,
+                      communication_hint='auto',
+                      timeout=0):
+  """Broadcasts one tensor to a group of others, across devices.
+
+  Args:
+    t: the tensor to be sent.
+    group_size: an int32 tensor.  One plus the number of receiving tensors, i.e.
+        the total number of devices participating.  Each tensor must reside on a
+        different device.
+    group_key: an int32 tensor identifying the group of devices.
+    instance_key: an int32 tensor identifying the participating group of Ops.
+    communication_hint: preferred collective communication.  The implementation
+      may fall back to another mechanism.  Options include `auto`, `ring`, and
+      `nccl`.
+    timeout: If set to a non zero, set a completion timeout to detect staleness.
+      If the timer goes off, a DeadlineExceededError is raised.
+      The timeout value in seconds. This feature is experimental.
+
+  Returns:
+    An Op implementing the distributed broadcast send.
+  """
+  return gen_collective_ops.collective_bcast_send_v2(
+      t,
+      group_size=group_size,
+      group_key=group_key,
+      instance_key=instance_key,
+      communication_hint=communication_hint.lower(),
+      timeout_seconds=timeout)
+
+
+def broadcast_recv(shape,
+                   dtype,
+                   group_size,
+                   group_key,
+                   instance_key,
+                   communication_hint='auto',
+                   timeout=0):
+  """Receives a broadcasts tensor, across devices.
+
+  Args:
+    shape: Shape of the tensor to be received.
+    dtype: Type of the tensor to be received.
+    group_size: one plus the number of receiving tensors, i.e. the total
+      number of devices participating.  Each tensor must reside on a
+      different device.
+    group_key: an integer identifying the group of devices.
+    instance_key: an integer identifying the participating group of Ops.
+    communication_hint: preferred collective communication.  The implementation
+      may fall back to another mechanism.  Options include `auto`, `ring`, and
+      `nccl`.
+    timeout: If set to a non zero, set a completion timeout to detect staleness.
+      If the timer goes off, a DeadlineExceededError is raised.
+      The timeout value in seconds. This feature is experimental.
+
+  Returns:
+    An Op implementing the broadcast receive.
+
+  Raises:
+    ValueError: if any of the input parameter constraints are not met.
+  """
+  if group_size <= 1:
+    raise ValueError(
+        'Parameter `group_size` to broadcast_send must be at least 2. '
+        f'Received: {group_size}.')
+  return gen_collective_ops.collective_bcast_recv(
+      shape=shape,
+      T=dtype,
+      group_size=group_size,
+      group_key=group_key,
+      instance_key=instance_key,
+      communication_hint=communication_hint.lower(),
+      timeout_seconds=timeout)
+
+
+def broadcast_recv_v2(shape,
+                      dtype,
+                      group_size,
+                      group_key,
+                      instance_key,
+                      communication_hint='auto',
+                      timeout=0):
+  """Receives a broadcasts tensor, across devices.
+
+  Args:
+    shape: an int tensor.  Shape of the tensor to be received.
+    dtype: Type of the tensor to be received.
+    group_size: an int32 tensor.  One plus the number of receiving tensors, i.e.
+        the total number of devices participating.  Each tensor must reside on a
+        different device.
+    group_key: an int32 tensor identifying the group of devices.
+    instance_key: an int32 tensor identifying the participating group of Ops.
+    communication_hint: preferred collective communication.  The implementation
+      may fall back to another mechanism.  Options include `auto`, `ring`, and
+      `nccl`.
+    timeout: If set to a non zero, set a completion timeout to detect staleness.
+      If the timer goes off, a DeadlineExceededError is raised.
+      The timeout value in seconds. This feature is experimental.
+
+  Returns:
+    An Op implementing the broadcast receive.
+  """
+  return gen_collective_ops.collective_bcast_recv_v2(
+      T=dtype,
+      group_size=group_size,
+      group_key=group_key,
+      instance_key=instance_key,
+      shape=shape,
+      communication_hint=communication_hint.lower(),
+      timeout_seconds=timeout)
+
+
+def initialize_communicator(group_key,
+                            rank,
+                            group_size,
+                            communication_hint='auto',
+                            timeout_seconds=0):
+  """Initializes a collective communicator.
+
+  This creates a collective communicator, which represents membership to a
+  collective group identified by the group_key. It should be called once per
+  member of the group, and each member needs to be on a different device.
+  It blocks until all members of the group run this op.
+
+  Communicators of a group can only be initialized once. Trying to initialize
+  communicators for an existing group key will result in an error.
+
+  Args:
+    group_key: an int32 `tf.Tensor` identifying the group.
+    rank: an `tf.Tensor` specifying the rank of this device in the group. If
+      specified, the rank is required to be unique in the group.
+    group_size: an int32 `tf.Tensor`. The size of the group.
+    communication_hint: preferred collective communication.  The implementation
+      may fall back to another mechanism.  Options include `auto`, `ring`, and
+      `nccl`.
+    timeout_seconds: If set to a non zero, set a completion timeout to detect
+      staleness. If the timer goes off, a DeadlineExceededError is raised. The
+      timeout value in seconds. This feature is experimental.
+
+
+  Returns:
+    A resource `tf.Tensor`.
+  """
+  return gen_collective_ops.collective_initialize_communicator(
+      group_key=group_key,
+      rank=rank,
+      group_size=group_size,
+      communication_hint=communication_hint,
+      timeout_seconds=timeout_seconds)
+
+
+def all_reduce_v3(communicator,
+                  t,
+                  reduction='Add',
+                  group_assignment=None,
+                  timeout_seconds=None):
+  """Reduces tensors mutually.
+
+  Args:
+    communicator: the resource `tf.Tensor` returned from
+      `initialize_communicator`.
+    t: the `tf.Tensor` to be reduced.
+    reduction: a string. The name of the operation to reduce the values.
+      Accepted values are `"min"`, `"max"`, `"mul"`, `"add"`.
+    group_assignment: Optional int32 `tf.Tensor` with shape [num_groups,
+      num_ranks_per_group]. `group_assignment[i]` represents the ranks in the
+      `ith` subgroup.
+    timeout_seconds: If set to a non zero, set a completion timeout to detect
+      staleness. If the timer goes off, a DeadlineExceededError is raised. The
+      timeout value in seconds. This feature is experimental.
+
+  Returns:
+    The reduced `tf.Tensor`.
+  """
+  if group_assignment is None:
+    group_assignment = []
+  return gen_collective_ops.collective_reduce_v3(
+      communicator=communicator,
+      input=t,
+      group_assignment=group_assignment,
+      reduction=reduction,
+      timeout_seconds=timeout_seconds)
+
+
+def all_to_all_v2(
+    t,
+    group_size,
+    group_key,
+    instance_key,
+    communication_hint='auto',
+    timeout=0,
+    ordering_token=None,
+    name=None,
+):
+  """Exchanges tensors mutually.
+
+  Args:
+    t: a `tf.Tensor`. The first dimension should have the length as the size of
+      the group. `t[i]` is sent to `rank i` within the group.
+    group_size: an int32 tensor, the total number of tensors to be mutually
+      exchanged. Each must reside on a different device. Should be a positive
+      integer.
+    group_key: an int32 tensor identifying the group of devices.
+    instance_key: an int32 tensor identifying the participating group of Ops.
+    communication_hint: preferred collective communication. The implementation
+      may fall back to another mechanism. Options include `auto` and `nccl`.
+    timeout: a float. If set to a non zero, set a completion timeout to detect
+      staleness. If the timer goes off, a DeadlineExceededError is raised. The
+      timeout value in seconds. This feature is experimental.
+    ordering_token: a resource tensor on the same device as the op to order the
+      collectives in a per-device manner by auto control dependency. This
+      argument can be omitted when there is one collective Op per `tf.function`,
+      or when explicit control dependency is used instead of auto control
+      dependency.
+    name: name of the Op.
+
+  Returns:
+    An Op implementing the distributed operation.
+  """
+  if ordering_token is not None:
+    ordering_token = [ordering_token]
+  else:
+    ordering_token = []
+
+  return gen_collective_ops.collective_all_to_all_v2(
+      t,
+      group_size=group_size,
+      group_key=group_key,
+      instance_key=instance_key,
+      communication_hint=communication_hint.lower(),
+      timeout_seconds=timeout,
+      is_stateless=False,
+      ordering_token=ordering_token,
+      name=name,
+  )
+
+
+def all_to_all_v3(communicator, t, group_assignment=None, timeout_seconds=None):
+  """Exchanges tensors mutually.
+
+  Args:
+    communicator: the resource `tf.Tensor` returned from
+      `initialize_communicator`.
+    t: a `tf.Tensor`. The first dimension should have the length as the size of
+      the group. `t[i]` is sent to `rank i` within the group.
+    group_assignment: Optional int32 `tf.Tensor` with shape [num_groups,
+      num_ranks_per_group]. `group_assignment[i]` represents the ranks in the
+      `ith` subgroup.
+    timeout_seconds: If set to a non zero, set a completion timeout to detect
+      staleness. If the timer goes off, a DeadlineExceededError is raised. The
+      timeout value in seconds. This feature is experimental.
+
+  Returns:
+    a `tf.Tensor`. `t[i]` is sent from `rank i` within the group.
+  """
+  if group_assignment is None:
+    group_assignment = []
+  return gen_collective_ops.collective_all_to_all_v3(
+      communicator=communicator,
+      input=t,
+      group_assignment=group_assignment,
+      timeout_seconds=timeout_seconds)

tensorflow/python/ops/collective_ops_benchmark.py

@@ -0,0 +1,83 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Local CPU benchmarks for collective ops."""
+
+import time
+
+import numpy as np
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.client import session
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import collective_ops
+from tensorflow.python.platform import test
+
+
+class CollectiveOpBenchmark(test.Benchmark):
+  """Benchmarks for local CPU collective op execution."""
+
+  def benchmark_collective(self):
+    """Measures the performance of local CPU collective execution."""
+    shapes = [(10,), (1000,), (1000000,)]
+    devices = [2, 4, 8]
+    collective_key_counter = 0
+
+    for group_size in devices:
+      group_key = collective_key_counter
+      instance_key = collective_key_counter
+      collective_key_counter += 1
+
+      for shape in shapes:
+        config = config_pb2.ConfigProto(device_count={"CPU": group_size})
+        with session.Session(config=config) as sess:
+          # Use a C++ callable to minimize the Python overhead in the benchmark.
+          callable_opts = config_pb2.CallableOptions()
+          reduce_ops = []
+          for device in range(group_size):
+            with ops.device("CPU:{}".format(device)):
+              t = constant_op.constant(np.multiply(range(shape[0]), 1.0))
+              r = collective_ops.all_reduce(t, group_size, group_key,
+                                            instance_key, "Add", "Div")
+              reduce_ops.append(r)
+              callable_opts.target.append(r.name)
+          op_callable = sess._make_callable_from_options(callable_opts)  # pylint: disable=protected-access
+
+          # Run five steps to warm up the session caches and do collective param
+          # resolution before taking the first measurement.
+          for _ in range(5):
+            op_callable()
+          deltas = []
+          overall_start = time.time()
+          # Run at least five repetitions and for at least five seconds.
+          while len(deltas) < 5 or time.time() - overall_start < 5.0:
+            start = time.time()
+            for _ in range(100):
+              op_callable()
+            end = time.time()
+            deltas.append(end - start)
+          del op_callable
+
+        median_wall_time = np.median(deltas) / 100.0
+        iters = len(deltas) * 100
+
+        self.report_benchmark(
+            iters=iters, wall_time=median_wall_time,
+            name="num_elements_{}_num_devices_{}".format(np.prod(shape),
+                                                         group_size))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/ops/collective_ops_gpu_test.py

@@ -0,0 +1,428 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for Collective Operations that require GPU."""
+
+import os
+import threading
+import time
+
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import collective_ops
+from tensorflow.python.platform import test
+
+
+class CollectiveOpGPUTest(test.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    """Set group_size = num_gpus = 2 for all tests in this class."""
+    super(CollectiveOpGPUTest, cls).setUpClass()
+    # Group size is the number of devices in a group communicating collectively.
+    # This will be passed into the collective ops in the tests below.
+    cls._group_size = 2
+    cls._devices = ['/device:GPU:{}'.format(i) for i in range(2)]
+    os.environ['NCCL_DEBUG'] = 'INFO'
+    os.environ['NCCL_LAUNCH_MODE'] = 'PARALLEL'
+
+  def _setup_context(self, num_gpus=2):
+    context._reset_context()
+    gpus = config.list_physical_devices('GPU')
+    if len(gpus) < num_gpus:
+      self.skipTest(
+          'Expected at least {} GPUs but found {} GPUs'.format(
+              num_gpus, len(gpus)
+          )
+      )
+    context.ensure_initialized()
+
+  def testBasicNcclAllReduce(self):
+    self._setup_context()
+
+    inputs = [
+        [0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
+        [0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3],
+    ]
+    expected = [0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2]
+    group_key = 1
+    instance_key = 1
+
+    @def_function.function
+    def run_basic_all_reduce():
+      collectives = []
+      for i in range(self._group_size):
+        with ops.device(self._devices[i]):
+          t = constant_op.constant(inputs[i])
+          collectives.append(
+              collective_ops.all_reduce(
+                  t, self._group_size, group_key, instance_key, 'Add', 'Div'
+              )
+          )
+      return collectives
+
+    for result in run_basic_all_reduce():
+      self.assertAllClose(result, expected, rtol=1e-5, atol=1e-5)
+
+  def testInt32Error(self):
+    self._setup_context()
+
+    inputs = [[0, 1], [2, 3]]
+    group_key = 1
+    instance_key = 50
+
+    @def_function.function
+    def run_int32_error():
+      for i in range(self._group_size):
+        with ops.device(self._devices[i]):
+          t = constant_op.constant(inputs[i], dtype=dtypes.int32)
+          collective_ops.all_reduce(
+              t, self._group_size, group_key, instance_key, 'Add', 'Div'
+          )
+
+    with self.assertRaisesRegex(
+        errors.InternalError, 'does not support datatype DT_INT32 on DEVICE_GPU'
+    ):
+      run_int32_error()
+
+  def testFp16Reduce(self):
+    self._setup_context()
+
+    inputs = [
+        [0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
+        [0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3],
+    ]
+    expected = [0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2]
+    group_key = 1
+    instance_key = 100
+
+    @def_function.function
+    def run_fp16_reduce():
+      collectives = []
+      for i in range(self._group_size):
+        with ops.device(self._devices[i]):
+          t = constant_op.constant(inputs[i], dtype=dtypes.float16)
+          collectives.append(
+              collective_ops.all_reduce(
+                  t, self._group_size, group_key, instance_key, 'Add', 'Div'
+              )
+          )
+      return collectives
+
+    for result in run_fp16_reduce():
+      self.assertAllClose(result, expected, rtol=1e-3, atol=1e-3)
+
+  def testNcclHintAllReduce(self):
+    self._setup_context()
+
+    inputs = [
+        [0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
+        [0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3],
+    ]
+    expected = [0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2]
+    group_key = 1
+    instance_key = 1
+
+    @def_function.function
+    def run_nccl_hint_all_reduce():
+      collectives = []
+      for i in range(self._group_size):
+        with ops.device(self._devices[i]):
+          t = constant_op.constant(inputs[i])
+          collectives.append(
+              collective_ops.all_reduce(
+                  t,
+                  self._group_size,
+                  group_key,
+                  instance_key,
+                  'Add',
+                  'Div',
+                  communication_hint='nccl',
+              )
+          )
+      return collectives
+
+    for result in run_nccl_hint_all_reduce():
+      self.assertAllClose(result, expected, rtol=1e-5, atol=1e-5)
+
+  def testBasicNcclBroadcast(self):
+    self._setup_context()
+
+    tensor_value = [0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1]
+    group_key = 1
+    instance_key = 1
+
+    @def_function.function
+    def run_basic_nccl_broadcast():
+      collectives = []
+      with ops.device(self._devices[0]):
+        t = constant_op.constant(tensor_value)
+        collectives.append(
+            collective_ops.broadcast_send(
+                t, t.shape, t.dtype, self._group_size, group_key, instance_key
+            )
+        )
+      with ops.device(self._devices[1]):
+        t = constant_op.constant(tensor_value)
+        collectives.append(
+            collective_ops.broadcast_recv(
+                t.shape, t.dtype, self._group_size, group_key, instance_key
+            )
+        )
+      return collectives
+
+    for result in run_basic_nccl_broadcast():
+      self.assertAllClose(result, tensor_value, rtol=1e-5, atol=1e-5)
+
+  def testNcclBroadcastDoubleRecv(self):
+    self._setup_context()
+
+    tensor_value = [0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1]
+    group_key = 1
+    instance_key = 1
+
+    @def_function.function
+    def run_nccl_broadcast_double_recv():
+      for device in self._devices:
+        with ops.device(device):
+          t = constant_op.constant(tensor_value)
+          collective_ops.broadcast_recv(
+              t.shape, t.dtype, self._group_size, group_key, instance_key
+          )
+
+    with self.assertRaisesRegex(errors.InternalError, 'found no source'):
+      run_nccl_broadcast_double_recv()
+
+  def testNcclBroadcastDoubleSend(self):
+    self._setup_context()
+
+    tensor_value = [0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1]
+    group_key = 1
+    instance_key = 1
+
+    @def_function.function
+    def run_nccl_broadcast_double_send():
+      for device in self._devices:
+        with ops.device(device):
+          t = constant_op.constant(tensor_value)
+          collective_ops.broadcast_send(
+              t, t.shape, t.dtype, self._group_size, group_key, instance_key
+          )
+
+    with self.assertRaisesRegex(errors.InternalError, 'already has source'):
+      run_nccl_broadcast_double_send()
+
+  def testBasicNcclAllGather(self):
+    self._setup_context()
+
+    inputs = [
+        [0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
+        [0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3],
+    ]
+    expected = [
+        0.1,
+        1.1,
+        2.1,
+        3.1,
+        4.1,
+        5.1,
+        6.1,
+        7.1,
+        0.3,
+        1.3,
+        2.3,
+        3.3,
+        4.3,
+        5.3,
+        6.3,
+        7.3,
+    ]
+    group_key = 1
+    instance_key = 1
+
+    @def_function.function
+    def run_basic_nccl_all_gather():
+      collectives = []
+      for i in range(self._group_size):
+        with ops.device(self._devices[i]):
+          t = constant_op.constant(inputs[i])
+          collectives.append(
+              collective_ops.all_gather(
+                  t, self._group_size, group_key, instance_key
+              )
+          )
+      return collectives
+
+    for result in run_basic_nccl_all_gather():
+      self.assertAllClose(result, expected, rtol=1e-5, atol=1e-5)
+
+  def testCollectiveDeviceMismatch(self):
+    self._setup_context()
+
+    group_key = 10
+    instance_key = 20
+    t0 = [1, 2, 3, 4]
+    t1 = [5, 6, 7, 8]
+
+    @def_function.function
+    def run_collective_device_mismatch():
+      with ops.device('/CPU:0'):
+        in0 = constant_op.constant(t0)
+        collective_ops.all_reduce(
+            in0, self._group_size, group_key, instance_key, 'Add', 'Id'
+        )
+      with ops.device('/GPU:0'):
+        in1 = constant_op.constant(t1)
+        collective_ops.all_reduce(
+            in1, self._group_size, group_key, instance_key, 'Add', 'Id'
+        )
+
+    with self.assertRaisesRegex(
+        errors.InternalError, 'but that group has type'
+    ):
+      run_collective_device_mismatch()
+
+  def testCollectiveReduceMinMax(self):
+    self._setup_context()
+
+    @def_function.function
+    def run_all_reduce(group_key, instance_key, merge_op):
+      t0 = [1.0, 20.0, 3.0, 40.0, 5.0]
+      t1 = [10.0, 2.0, 30.0, 4.0, 50.0]
+      with ops.device('/GPU:0'):
+        in0 = constant_op.constant(t0)
+        c0 = collective_ops.all_reduce(
+            in0,
+            self._group_size,
+            group_key,
+            instance_key,
+            merge_op,
+            final_op='Id',
+            communication_hint='nccl',
+        )
+      with ops.device('/GPU:1'):
+        in1 = constant_op.constant(t1)
+        c1 = collective_ops.all_reduce(
+            in1,
+            self._group_size,
+            group_key,
+            instance_key,
+            merge_op,
+            final_op='Id',
+            communication_hint='nccl',
+        )
+      return c0, c1
+
+    for combination in [
+        ('Max', [10.0, 20.0, 30.0, 40.0, 50.0]),
+        ('Min', [1.0, 2.0, 3.0, 4.0, 5.0]),
+    ]:
+      merge_op = combination[0]
+      results = run_all_reduce(group_key=10, instance_key=20, merge_op=merge_op)
+      expected = combination[1]
+      for result in results:
+        self.assertAllClose(result, expected, rtol=1e-5, atol=1e-5)
+
+  def testNcclStress(self):
+    self.skipTest(
+        'b/435404154: As we moved from NVIDIA CUDA base image to Ubuntu 22.04'
+        ' with NVIDIA Driver 580 installed for RBE, this test is failing and'
+        ' needs to be addressed as part of the bug.'
+    )
+    self._setup_context(num_gpus=1)
+
+    num_iters = 1000
+    for _ in range(num_iters):
+      with ops.device('/device:GPU:0'):
+        collective_ops.all_reduce(
+            [1.0],
+            group_size=1,
+            group_key=0,
+            instance_key=0,
+            merge_op='Add',
+            final_op='Id',
+            communication_hint='NCCL',
+        )
+
+  @test_util.run_v2_only
+  def testAbortNccl(self):
+    self._setup_context(num_gpus=2)
+
+    group_size = 2
+    group_key = 100
+    instance_key = 100
+    in_tensor = constant_op.constant(1.0)
+
+    # First perform a normal collective to finish resolution.
+    def collective_fn():
+      for device in ['GPU:0', 'GPU:1']:
+        with ops.device(device):
+          collective_ops.all_reduce(
+              in_tensor,
+              group_size,
+              group_key,
+              instance_key,
+              'Add',
+              'Id',
+              communication_hint='nccl',
+          )
+
+    def_function.function(collective_fn)()
+
+    # Launch a collective that hangs, and abort the collective executor after
+    # the launch.
+    def abort_fn():
+      time.sleep(2)
+      context.context().abort_collective_ops(errors.UNAVAILABLE, 'peer down')
+
+    t = threading.Thread(target=abort_fn)
+    t.start()
+
+    with self.assertRaisesRegex(errors.UnavailableError, 'peer down'):
+      collective_ops.all_reduce(
+          in_tensor,
+          group_size,
+          group_key,
+          instance_key,
+          'Add',
+          'Id',
+          communication_hint='nccl',
+      )
+
+    # After abortion, subsequent collectives should fail immediately.
+    with self.assertRaisesRegex(errors.UnavailableError, 'peer down'):
+      collective_ops.all_reduce(
+          in_tensor,
+          group_size,
+          group_key,
+          instance_key,
+          'Add',
+          'Id',
+          communication_hint='nccl',
+      )
+
+    t.join()
+    # Reset the context in order to reset the collective executor.
+    context._reset_context()  # pylint: disable=protected-access
+    def_function.function(collective_fn)()
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/ops/collective_ops_test.py

@@ -0,0 +1,568 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for Collective Operations."""
+
+import time
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.core.protobuf import rewriter_config_pb2
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import kernels
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import collective_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.ops import while_loop
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging as logging
+
+
+class CollectiveOpTest(test.TestCase):
+
+  def setUp(self):
+    context._reset_context()  # pylint: disable=protected-access
+    super(CollectiveOpTest, self).setUp()
+
+  def _testCollectiveReduce(self,
+                            inputs,
+                            expected,
+                            set_graph_key,
+                            communication_hint='auto',
+                            fp16=False,
+                            instance_key=1,
+                            merge_op='Add',
+                            final_op='Div',
+                            timeout=0,
+                            reported_group_size=None):
+    group_key = 1
+    group_size = len(inputs)
+    if reported_group_size is None:
+      reported_group_size = group_size
+    device_type = 'CPU'
+    config = config_pb2.ConfigProto(device_count={device_type: group_size})
+    devices = ['/{}:{}'.format(device_type, i) for i in range(group_size)]
+
+    with self.session(config=config) as sess:
+      colred = []
+      for i in range(group_size):
+        with ops.device(devices[i]):
+          tensor = constant_op.constant(inputs[i], dtype=(
+              dtypes.float16 if fp16 else dtypes.float32))
+          colred.append(
+              collective_ops.all_reduce(
+                  tensor,
+                  reported_group_size,
+                  group_key,
+                  instance_key,
+                  merge_op,
+                  final_op,
+                  communication_hint=communication_hint,
+                  timeout=timeout))
+      run_options = config_pb2.RunOptions()
+      if set_graph_key:
+        run_options.experimental.collective_graph_key = 1
+      results = sess.run(colred, options=run_options)
+    tolerance = 1e-3 if fp16 else 1e-5
+    for i in range(group_size):
+      logging.info('i {} result {} expected {}'.format(i, results[i], expected))
+      self.assertAllClose(results[i], expected, rtol=tolerance, atol=tolerance)
+
+  def _testMultipleConcurrentCollectiveReduce(self, t0, t1, expected):
+    group_key = 1
+    group_size = 2
+    num_instances = 2
+    all_reduces = []
+    config = config_pb2.ConfigProto(device_count={'CPU': group_size})
+    config.experimental.collective_deterministic_sequential_execution = True
+    with self.session(config=config) as sess:
+      for cpu in range(group_size):
+        with ops.device('/CPU:%d' % cpu):
+          in_tensor = constant_op.constant(t0 if cpu == 0 else t1)
+          for instance in range(num_instances):
+            all_reduces.append(collective_ops.all_reduce(
+                in_tensor, group_size, group_key, instance, 'Add', 'Div'))
+      results = sess.run(all_reduces)
+    for i in range(group_size * num_instances):
+      self.assertAllClose(results[i], expected, rtol=1e-5, atol=1e-5)
+
+  def testCollectiveReduce(self):
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveReduce(
+          inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
+                  [0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
+          expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
+          set_graph_key=True)
+
+  def testCollectiveAutoGraphKey(self):
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveReduce(
+          inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
+                  [0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
+          expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
+          set_graph_key=False)
+
+  def testFp16Reduce(self):
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveReduce(
+          inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
+                  [0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
+          expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
+          set_graph_key=True,
+          fp16=True)
+
+  def testCollectiveMultipleConcurrentReduce(self):
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testMultipleConcurrentCollectiveReduce(
+          [0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
+          [0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3],
+          [0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2])
+
+  def testCollectiveTimeoutV1(self):
+    timeout = 4.5
+    kwargs = dict(
+        inputs=[[i + j + 0.1 for i in range(8)] for j in range(3)],
+        expected=[1 + i + 0.1 for i in range(8)],
+        set_graph_key=True,
+        timeout=timeout)
+
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveReduce(**kwargs)
+
+    start_time = time.time()
+    with ops.Graph().as_default():
+      with self.assertRaisesRegex(
+          errors.DeadlineExceededError,
+          'Collective has timed out waiting for other workers'):
+        self._testCollectiveReduce(
+            reported_group_size=len(kwargs['inputs']) + 1, **kwargs)
+    elapsed = time.time() - start_time
+    self.assertAllGreaterEqual(elapsed, timeout)
+
+  def testNcclHintFallbackToRingReduce(self):
+    """Tests that setting `communication_hint=nccl` works on non-GPU builds."""
+    if kernels.get_registered_kernels_for_op('NcclAllReduce'):
+      self.skipTest('Run only on non-GPU environments')
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveReduce(
+          inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
+                  [0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
+          expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
+          set_graph_key=False,
+          communication_hint='nccl')
+
+  def _testWhile(self, num_vars, num_iterations, key_base):
+    group_size = 2
+    group_key = 1
+    instances = [(key_base + i) for i in range(num_vars)]
+    devices = ['CPU:{}'.format(i) for i in range(group_size)]
+
+    config = config_pb2.ConfigProto(device_count={'CPU': group_size})
+    rewrite_options = config.graph_options.rewrite_options
+    rewrite_options.scoped_allocator_optimization = (
+        rewriter_config_pb2.RewriterConfig.ON)
+    del rewrite_options.scoped_allocator_opts.enable_op[:]
+    rewrite_options.scoped_allocator_opts.enable_op.append('CollectiveReduce')
+
+    with self.session(config=config) as sess:
+      loop_vars = []
+      for device in devices:
+        with ops.device(device):
+          loop_vars.append(
+              [variable_v1.VariableV1((1 << i) * 1.) for i in range(num_vars)])
+      # This variable controls number of iterations.
+      loop_vars.append(variable_v1.VariableV1(0.))
+      def loop_body(dev0_tensors, dev1_tensors, loop_tensor):
+        return_ops = []
+        for i in range(len(devices)):
+          device = devices[i]
+          device_tensors = dev0_tensors if i == 0 else dev1_tensors
+          with ops.device(device):
+            device_collectives = []
+            for j in range(num_vars):
+              # NOTE(ayushd): we need the `cast` here to ensure that the input
+              # to `all_reduce` has an explicit device string.  We don't use
+              # `identity` because `cast` is more resilient to getting optimized
+              # away by various optimization passes.
+              input_tensor = math_ops.cast(device_tensors[j], dtypes.float16)
+              collective_op = collective_ops.all_reduce(
+                  input_tensor, group_size, group_key, instances[j],
+                  'Add', 'Id')
+              output_tensor = math_ops.cast(collective_op, dtypes.float32)
+              device_collectives.append(output_tensor)
+            return_ops.append(device_collectives)
+        return_ops.append(math_ops.add(loop_tensor, 1.))
+        return return_ops
+      # Run until last variable exceeds number of iterations.
+      loop_cond = lambda d0, d1, i: math_ops.less(i, num_iterations)
+      sess.run(variables.global_variables_initializer())
+      results = sess.run(while_loop.while_loop(loop_cond, loop_body, loop_vars))
+      self.assertEqual(results[:-1], [
+          [((1 << (num_iterations + v)) * 1.) for v in range(num_vars)]
+          for _ in range(group_size)])
+
+  def testSimpleWhile(self):
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testWhile(num_vars=1, num_iterations=4, key_base=20)
+
+  def testWhileMultipleAllReduce(self):
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testWhile(num_vars=2, num_iterations=4, key_base=20)
+
+  def testWhileWithScopedAllocator(self):
+    group_size = 2
+    group_key = 1
+    instance_key0 = 1
+    instance_key1 = 2
+
+    config = config_pb2.ConfigProto(device_count={'CPU': group_size})
+    rewrite_options = config.graph_options.rewrite_options
+    rewrite_options.scoped_allocator_optimization = (
+        rewriter_config_pb2.RewriterConfig.ON)
+    del rewrite_options.scoped_allocator_opts.enable_op[:]
+    rewrite_options.scoped_allocator_opts.enable_op.append('CollectiveReduce')
+
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      with self.session(config=config) as sess:
+        run_ops = []
+        for i in range(group_size):
+          with ops.device('CPU:%d' % i):
+            constant = constant_op.constant(0.)
+            cond = lambda i: math_ops.less(i, 10.)
+            body = lambda i: math_ops.add(i, 1.)
+            input0 = while_loop.while_loop(cond, body, [constant])
+            input1 = math_ops.add(constant, 5)
+            colred0 = collective_ops.all_reduce(input0, group_size, group_key,
+                                                instance_key0, 'Add', 'Id')
+            colred1 = collective_ops.all_reduce(input1, group_size, group_key,
+                                                instance_key1, 'Add', 'Id')
+            run_ops.append(math_ops.add_n([colred0, colred1]))
+        results = sess.run(run_ops)
+      self.assertEqual(results, [30., 30.])
+
+  def testCollectiveReduceScalar(self):
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveReduce(inputs=[0.1, 0.3], expected=0.2,
+                                 set_graph_key=True)
+
+  def testCollectiveReduceMaximum(self):
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveReduce(
+          inputs=[[1., 20., 3., 40., 5.], [10., 2., 30., 4., 50.]],
+          expected=[10., 20., 30., 40., 50.],
+          set_graph_key=True,
+          instance_key=30,
+          merge_op='Max',
+          final_op='Id')
+
+  def testCollectiveReduceMinimum(self):
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveReduce(
+          inputs=[[1., 20., 3., 40., 5.], [10., 2., 30., 4., 50.]],
+          expected=[1., 2., 3., 4., 5.],
+          set_graph_key=True,
+          instance_key=40,
+          merge_op='Min',
+          final_op='Id')
+
+  def _testCollectiveBroadcast(self, in_val):
+    group_key = 1
+    instance_key = 1
+    with self.session(
+        config=config_pb2.ConfigProto(device_count={'CPU': 2})) as sess:
+      with ops.device('/CPU:0'):
+        in0 = constant_op.constant(in_val)
+        out0 = collective_ops.broadcast_send(in0, in0.shape, in0.dtype,
+                                             2, group_key, instance_key)
+      with ops.device('/CPU:1'):
+        c1 = constant_op.constant(in_val)
+        out1 = collective_ops.broadcast_recv(c1.shape, c1.dtype,
+                                             2, group_key, instance_key)
+      run_options = config_pb2.RunOptions()
+      run_options.experimental.collective_graph_key = 1
+      results = sess.run([out0, out1], options=run_options)
+    self.assertAllClose(results[0], in_val, rtol=1e-5, atol=1e-5)
+    self.assertAllClose(results[1], in_val, rtol=1e-5, atol=1e-5)
+
+  def testCollectiveBroadcast(self):
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveBroadcast([0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1])
+
+  def testCollectiveBroadcastBool(self):
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveBroadcast([True, False])
+
+  def _testCollectiveGather(self, t0, t1, expected, set_graph_key):
+    group_key = 1
+    instance_key = 1
+    with self.session(
+        config=config_pb2.ConfigProto(device_count={'CPU': 2})) as sess:
+      with ops.device('/CPU:0'):
+        in0 = constant_op.constant(t0)
+        c0 = collective_ops.all_gather(in0, 2, group_key, instance_key)
+      with ops.device('/CPU:1'):
+        in1 = constant_op.constant(t1)
+        c1 = collective_ops.all_gather(in1, 2, group_key, instance_key)
+      run_options = config_pb2.RunOptions()
+      if set_graph_key:
+        run_options.experimental.collective_graph_key = 1
+      results = sess.run([c0, c1], options=run_options)
+    self.assertAllClose(results[0], expected, rtol=1e-5, atol=1e-5)
+    self.assertAllClose(results[1], expected, rtol=1e-5, atol=1e-5)
+
+  def testCollectiveGather(self):
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveGather([0, 1, 2, 3, 4, 5, 6, 7],
+                                 [10, 11, 12, 13, 14, 15, 16, 17],
+                                 [0, 1, 2, 3, 4, 5, 6, 7,
+                                  10, 11, 12, 13, 14, 15, 16, 17],
+                                 True)
+      self._testCollectiveGather([[0, 1, 2, 3], [4, 5, 6, 7]],
+                                 [[10, 11, 12, 13], [14, 15, 16, 17]],
+                                 [[0, 1, 2, 3], [4, 5, 6, 7],
+                                  [10, 11, 12, 13], [14, 15, 16, 17]],
+                                 True)
+      self._testCollectiveGather([[[0, 1], [2, 3]], [[4, 5], [6, 7]]],
+                                 [[[10, 11], [12, 13]], [[14, 15], [16, 17]]],
+                                 [[[0, 1], [2, 3]], [[4, 5], [6, 7]],
+                                  [[10, 11], [12, 13]], [[14, 15], [16, 17]]],
+                                 True)
+
+  def testCollectiveGatherShapeMismatch(self):
+    group_key = 1
+    instance_key = 1
+    t0 = [1, 2, 3, 4]
+    t1 = [5, 6, 7, 8]
+    t2 = [9, 10]
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      with self.session(
+          config=config_pb2.ConfigProto(device_count={'CPU': 2})) as sess:
+        with ops.device('/CPU:0'):
+          in0 = constant_op.constant(t0)
+          c0 = collective_ops.all_gather(in0, 2, group_key, instance_key)
+        with ops.device('/CPU:1'):
+          in1 = constant_op.constant(t1)
+          in2 = constant_op.constant(t2)
+          c1 = collective_ops.all_gather(in1, 2, group_key, instance_key)
+          c2 = collective_ops.all_gather(in2, 2, group_key, instance_key)
+        run_options = config_pb2.RunOptions()
+        run_options.experimental.collective_graph_key = 1
+        sess.run([c0, c1], options=run_options)
+        with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                    'Shape mismatch'):
+          sess.run([c0, c2], options=run_options)
+
+  def testCollectiveGatherShapeMismatchAcrossDevices(self):
+    group_key = 1
+    instance_key = 1
+    t0 = [1, 2, 3, 4]
+    t1 = [5, 6]
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      with self.session(
+          config=config_pb2.ConfigProto(device_count={'CPU': 2})) as sess:
+        with ops.device('/CPU:0'):
+          in0 = constant_op.constant(t0)
+          c0 = collective_ops.all_gather(in0, 2, group_key, instance_key)
+        with ops.device('/CPU:1'):
+          in1 = constant_op.constant(t1)
+          c1 = collective_ops.all_gather(in1, 2, group_key, instance_key)
+        run_options = config_pb2.RunOptions()
+        run_options.experimental.collective_graph_key = 1
+        with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                    'Shape mismatch'):
+          sess.run([c0, c1], options=run_options)
+
+  def testCollectiveGatherPolymorphicShape(self):
+    t0 = [0, 1, 2, 3, 4, 5, 6, 7]
+    t1 = [10, 11, 12, 13, 14, 15, 16, 17]
+    group_size = 2
+    group_key = 1
+    instance_key = 123
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      with self.session(
+          config=config_pb2.ConfigProto(
+              device_count={'CPU': group_size})) as sess:
+        with ops.device('/CPU:0'):
+          in0 = array_ops.placeholder(dtype=dtypes.int32, shape=[None])
+          c0 = collective_ops.all_gather(in0, group_size, group_key,
+                                         instance_key)
+        with ops.device('/CPU:1'):
+          in1 = array_ops.placeholder(dtype=dtypes.int32, shape=[None])
+          c1 = collective_ops.all_gather(in1, group_size, group_key,
+                                         instance_key)
+
+        results = sess.run([c0, c1], feed_dict={in0: t0, in1: t1})
+        results_ = sess.run([c0, c1], feed_dict={in0: t0[1:], in1: t1[1:]})
+
+    expected_output = [0, 1, 2, 3, 4, 5, 6, 7, 10, 11, 12, 13, 14, 15, 16, 17]
+    self.assertAllClose(results[0], expected_output, rtol=1e-5, atol=1e-5)
+    self.assertAllClose(results[1], expected_output, rtol=1e-5, atol=1e-5)
+
+    expected_output_ = [1, 2, 3, 4, 5, 6, 7, 11, 12, 13, 14, 15, 16, 17]
+    self.assertAllClose(results_[0], expected_output_, rtol=1e-5, atol=1e-5)
+    self.assertAllClose(results_[1], expected_output_, rtol=1e-5, atol=1e-5)
+
+  @test_util.run_v2_only
+  @test_util.disable_tfrt(
+      'b/177270918: TFRT has dead lock when executing collective ops.')
+  def testCollectiveGroupSizeMismatch(self):
+    cpus = config.list_physical_devices('CPU')
+    self.assertEqual(len(cpus), 1)
+    config.set_logical_device_configuration(cpus[0], [
+        context.LogicalDeviceConfiguration(),
+        context.LogicalDeviceConfiguration()
+    ])
+    context.ensure_initialized()
+
+  @test_util.run_v2_only
+  def testCollectiveGatherShapeCheckFailure(self):
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                'input should have rank > 0'):
+      collective_ops.gen_collective_ops.CollectiveGather(
+          input=1,
+          group_size=1,
+          group_key=1,
+          instance_key=1,
+          shape=(3, 3, 3),
+          communication_hint='auto',
+          timeout_seconds=0,
+          name='')
+
+    @def_function.function
+    def run_all_reduce():
+      group_key = 10
+      instance_key = 20
+      t0 = [1, 2, 3, 4]
+      t1 = [5, 6, 7, 8]
+      with ops.device('/CPU:0'):
+        in0 = constant_op.constant(t0)
+        c0 = collective_ops.all_reduce(
+            in0, group_size=2, group_key=group_key, instance_key=instance_key,
+            merge_op='Add', final_op='Id')
+      with ops.device('/CPU:1'):
+        in1 = constant_op.constant(t1)
+        c1 = collective_ops.all_reduce(
+            in1, group_size=3, group_key=group_key, instance_key=instance_key,
+            merge_op='Add', final_op='Id')
+      return c0, c1
+
+    with self.assertRaisesRegex(errors.InternalError,
+                                'but that group has size'):
+      run_all_reduce()
+
+  @test_util.run_v2_only
+  def testCollectiveTensorsHaveNoDeviceSpecified(self):
+    cpus = config.list_physical_devices('CPU')
+    self.assertEqual(len(cpus), 1)
+    config.set_logical_device_configuration(cpus[0], [
+        context.LogicalDeviceConfiguration(),
+        context.LogicalDeviceConfiguration()
+    ])
+    context.ensure_initialized()
+
+    group_size = 2
+    group_key = 1
+    instance_key = 1
+
+    @def_function.function
+    def fn(all_args):
+      results = []
+      # The inputs have no devices set. This is expected to be a trace-time
+      # check only.
+      self.assertEqual(all_args[0].device, '')
+      self.assertEqual(all_args[1].device, '')
+
+      with ops.device('/CPU:0'):
+        results.append(
+            collective_ops.all_reduce(all_args[0], group_size, group_key,
+                                      instance_key, 'Add', 'Div'))
+      with ops.device('/CPU:1'):
+        results.append(
+            collective_ops.all_reduce(all_args[1], group_size, group_key,
+                                      instance_key, 'Add', 'Div'))
+
+      return results
+
+    with ops.device('/CPU:0'):
+      in0 = constant_op.constant(1)
+    with ops.device('/CPU:1'):
+      in1 = constant_op.constant(3)
+    result = fn([in0, in1])
+    self.assertAllClose(result, [2, 2])
+
+  def testConstantWithScopedAllocator(self):
+    group_size = 2
+    group_key = 1
+    instance_key1 = 1
+    instance_key2 = 2
+
+    graph_options = config_pb2.GraphOptions(
+        optimizer_options=config_pb2.OptimizerOptions(do_constant_folding=True))
+    cfg = config_pb2.ConfigProto(device_count={'CPU': group_size},
+                                 graph_options=graph_options)
+    rewrite_options = cfg.graph_options.rewrite_options
+    rewrite_options.scoped_allocator_optimization = (
+        rewriter_config_pb2.RewriterConfig.ON)
+    del rewrite_options.scoped_allocator_opts.enable_op[:]
+    rewrite_options.scoped_allocator_opts.enable_op.append('CollectiveReduce')
+
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      with self.session(config=cfg) as sess:
+        run_ops = []
+        for i in range(group_size):
+          with ops.device('CPU:%d' % i):
+            constant = constant_op.constant(i + 1.)
+            input_tensor1 = array_ops.identity(constant)
+            input_tensor2 = array_ops.identity(constant)
+            reduced_tensor1 = collective_ops.all_reduce(
+                input_tensor1, group_size, group_key, instance_key1, 'Add',
+                'Id')
+            reduced_tensor2 = collective_ops.all_reduce(
+                input_tensor2, group_size, group_key, instance_key2, 'Add',
+                'Id')
+            run_ops.append(array_ops.identity(reduced_tensor1))
+            run_ops.append(array_ops.identity(reduced_tensor2))
+        results = sess.run(run_ops)
+    self.assertEqual(results, [3., 3., 3., 3.])
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/ops/collective_ops_xla_test.py

@@ -0,0 +1,74 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for Collective Operations with XLA."""
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.core.protobuf import rewriter_config_pb2
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import collective_ops
+from tensorflow.python.platform import test
+
+
+class CollectiveOpXlaTest(test.TestCase):
+
+  def testScopedAllocatorWithXla(self):
+    group_size = 2
+    group_key = 1
+    instance_key1 = 1
+    instance_key2 = 2
+    tensor_size = 10
+
+    graph_options = config_pb2.GraphOptions(
+        optimizer_options=config_pb2.OptimizerOptions(
+            do_constant_folding=False))
+    cfg = config_pb2.ConfigProto(device_count={'CPU': group_size},
+                                 graph_options=graph_options)
+    rewrite_options = cfg.graph_options.rewrite_options
+    rewrite_options.scoped_allocator_optimization = (
+        rewriter_config_pb2.RewriterConfig.ON)
+    del rewrite_options.scoped_allocator_opts.enable_op[:]
+    rewrite_options.scoped_allocator_opts.enable_op.append('CollectiveReduce')
+
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default(), self.session(config=cfg) as sess:
+      run_ops = []
+      for i in range(group_size):
+        with ops.device('CPU:%d' % i):
+          tensor_val = [i + 1.] * tensor_size
+          constant = constant_op.constant(tensor_val)
+
+          @def_function.function(jit_compile=True)
+          def f(x):
+            return 2 * x + 1
+
+          input_tensor1 = array_ops.identity(f(constant))
+          input_tensor2 = array_ops.identity(f(constant))
+          reduced_tensor1 = collective_ops.all_reduce(
+              input_tensor1, group_size, group_key, instance_key1, 'Add', 'Id')
+          reduced_tensor2 = collective_ops.all_reduce(
+              input_tensor2, group_size, group_key, instance_key2, 'Add', 'Id')
+          run_ops.append(array_ops.identity(reduced_tensor1))
+          run_ops.append(array_ops.identity(reduced_tensor2))
+      results = sess.run(run_ops)
+      for result in results:
+        for result_val in result:
+          self.assertEqual(result_val, 8.)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/ops/compiled_collective_ops_gpu_test.py

@@ -0,0 +1,82 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for Collective Operations that require GPU."""
+
+import os
+
+from tensorflow.python.distribute import mirrored_strategy
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import collective_ops
+from tensorflow.python.ops import while_loop
+from tensorflow.python.platform import test
+
+
+class CompiledCollectiveOpGPUTest(test.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    """Set group_size = num_gpus = 2 for all tests in this class."""
+    super(CompiledCollectiveOpGPUTest, cls).setUpClass()
+    # Group size is the number of devices in a group communicating collectively.
+    # This will be passed into the collective ops in the tests below.
+    cls._group_size = 2
+    cls._devices = ['/device:GPU:{}'.format(i) for i in range(2)]
+    os.environ['NCCL_DEBUG'] = 'INFO'
+    os.environ['NCCL_LAUNCH_MODE'] = 'PARALLEL'
+
+  def _setup_context(self, num_gpus=2):
+    context._reset_context()
+    gpus = config.list_physical_devices('GPU')
+    if len(gpus) < num_gpus:
+      self.skipTest('Expected at least {} GPUs but found {} GPUs'.format(
+          num_gpus, len(gpus)))
+    context.ensure_initialized()
+
+  def testCompiledAllReduce(self):
+    self._setup_context()
+
+    def all_reduce_sum(v):
+      return collective_ops.all_reduce_v2(
+          t=v,
+          group_size=2,
+          group_key=1,
+          instance_key=1,
+          merge_op='Add',
+          final_op='Id')
+
+    strategy = mirrored_strategy.MirroredStrategy(['GPU:0', 'GPU:1'])
+
+    @def_function.function(jit_compile=True)
+    def f():
+      return while_loop.while_loop(
+          lambda i, _: i < 5, lambda i, t: (i + 1, all_reduce_sum(t)),
+          (array_ops.zeros([]), constant_op.constant(1.0)))
+
+    @def_function.function
+    def run():
+      return strategy.run(f)
+
+    _, reduce = strategy.experimental_local_results(run())[0]
+    self.assertEqual(reduce.numpy(), 32.0)
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/ops/composite_tensor_ops.py

@@ -0,0 +1,118 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Operations for ExtensionTypes (aka Composite Tensors)."""
+
+from tensorflow.core.protobuf import composite_tensor_variant_pb2
+from tensorflow.python.framework import composite_tensor
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import gen_composite_tensor_ops
+from tensorflow.python.saved_model import nested_structure_coder
+from tensorflow.python.util import nest
+
+
+def composite_tensor_to_variants(value, type_spec=None, name=None):
+  """Encodes `value` as a scalar variant tensor.
+
+  Args:
+    value: The `ExtensionType` value to encode.
+    type_spec: Information about the value's type that should be included in the
+      encoding.
+    name: Optional name for the operation.
+
+  Returns:
+    A Tensor with shape=`()` and dtype=`tf.variant`.
+
+  Raises:
+    ValueError: If `type_spec` is not compatible with `value`.
+  """
+  if not isinstance(value, composite_tensor.CompositeTensor):
+    raise TypeError("Expected `value` to be a CompositeTensor. "
+                    f"Received {type(value)}.")
+
+  if type_spec is None:
+    type_spec = value._type_spec  # pylint: disable=protected-access
+  if not type_spec.is_compatible_with(value):
+    raise ValueError(f"`type_spec` {type_spec} is not compatible with `value` "
+                     f"{value!r}.")
+  metadata = composite_tensor_variant_pb2.CompositeTensorVariantMetadata()
+  metadata.type_spec_proto.CopyFrom(
+      nested_structure_coder.encode_structure(type_spec).type_spec_value)
+
+  return gen_composite_tensor_ops.CompositeTensorVariantFromComponents(
+      components=nest.flatten(value, expand_composites=True),
+      metadata=metadata.SerializeToString(),
+      name=name)
+
+
+def composite_tensor_from_variant(encoded, type_spec, name=None):
+  """Returns the `ExtensionType` value encoded by a variant scalar tensor.
+
+  Args:
+    encoded: A Tensor returned by `composite_tensor_to_variants`.
+    type_spec: The `TypeSpec` of the original value.  This is used to determine
+      the number and types of the component tensors that comprise the decoded
+      value.  Must be compatible with the `TypeSpec` serilized in `encoded`.
+    name: Optional name for the operation.
+
+  Returns:
+    An `ExtensionType` value that is compatible with `TypeSpec`.
+
+  Raises:
+    TypeError: If `encoded` is not a Tensor with dtype=variant.
+    InvalidArgumentError: If `encoded` is not compatible with `type_spec`.
+  """
+  if not isinstance(encoded, tensor.Tensor):
+    raise TypeError(f"Expected `encoded` to be a Tensor, got {encoded!r}.")
+  if encoded.dtype != dtypes.variant:
+    raise TypeError("Expected `encoded` to have dtype=variant, got "
+                    f"{encoded!r}.")
+  encoded.shape.assert_is_compatible_with(())
+
+  metadata = composite_tensor_variant_pb2.CompositeTensorVariantMetadata()
+  metadata.type_spec_proto.CopyFrom(
+      nested_structure_coder.encode_structure(type_spec).type_spec_value)
+
+  component_dtypes = [
+      t.dtype for t in nest.flatten(type_spec, expand_composites=True)
+  ]
+
+  components = gen_composite_tensor_ops.CompositeTensorVariantToComponents(
+      encoded=encoded,
+      metadata=metadata.SerializeToString(),
+      Tcomponents=component_dtypes,
+      name=name)
+  return nest.pack_sequence_as(type_spec, components, expand_composites=True)
+
+
+@ops.RegisterGradient("CompositeTensorVariantFromComponents")
+def _composite_tensor_to_variants_grad(op, grad):
+  return gen_composite_tensor_ops.CompositeTensorVariantToComponents(
+      encoded=grad,
+      metadata=op.get_attr("metadata"),
+      Tcomponents=op.get_attr("Tcomponents"))
+
+
+@ops.RegisterGradient("CompositeTensorVariantToComponents")
+def _composite_tensor_from_variant_grad(op, *grad):
+  assert len(grad) == len(op.outputs)
+  # `components` is `op.outputs`, but with any tensors for which we're
+  # taking the gradient replaced by the corresponding value from `grad`.
+  components = [
+      op.outputs[i] if grad[i] is None else grad[i] for i in range(len(grad))
+  ]
+  return gen_composite_tensor_ops.CompositeTensorVariantFromComponents(
+      components=components, metadata=op.get_attr("metadata"))

tensorflow/python/ops/concat_benchmark.py

@@ -0,0 +1,143 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Benchmark for split and grad of split."""
+
+import itertools
+import random
+import time
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.client import session as session_lib
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import gradients_impl
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+def build_graph(device, input_shape, variable, num_inputs, axis, grad):
+  """Build a graph containing a sequence of concat operations.
+
+  Args:
+    device: string, the device to run on.
+    input_shape: shape of the input tensors.
+    variable: whether or not to randomize the input shape
+    num_inputs: the number of inputs to concat
+    axis: axis to be concat'ed
+    grad: if True compute the gradient
+
+  Returns:
+    An array of tensors to run()
+  """
+  with ops.device("/%s:0" % device):
+    if not variable:
+      inputs = [array_ops.zeros(input_shape) for _ in range(num_inputs)]
+    else:
+      if axis == 1:
+        inputs = [
+            array_ops.zeros([
+                input_shape[0],
+                random.randint(max(1, input_shape[1] - 5), input_shape[1] + 5)
+            ]) for _ in range(num_inputs)
+        ]
+      else:
+        inputs = [
+            array_ops.zeros([
+                random.randint(max(1, input_shape[0] - 5), input_shape[0] + 5),
+                input_shape[1]
+            ]) for _ in range(num_inputs)
+        ]
+
+    outputs = [array_ops.concat(inputs, axis) for _ in range(100)]
+    if grad:
+      return control_flow_ops.group(*list(
+          itertools.chain.from_iterable([
+              gradients_impl.gradients(output, inputs) for output in outputs
+          ])))
+    else:
+      return control_flow_ops.group(*outputs)
+
+
+class ConcatBenchmark(test.Benchmark):
+  """Benchmark concat."""
+
+  def _run_graph(self, device, input_shape, variable, num_inputs, axis, grad,
+                 num_iters):
+    """Run the graph and print its execution time.
+
+    Args:
+      device: string, the device to run on.
+      input_shape: shape of the input tensors.
+      variable: whether or not the input shape should be fixed
+      num_inputs: the number of inputs to concat
+      axis: axis to be concat'ed
+      grad: if True compute the gradient
+      num_iters: number of steps to run.
+
+    Returns:
+      The duration of the run in seconds.
+    """
+    graph = ops.Graph()
+    with graph.as_default():
+      outputs = build_graph(device, input_shape, variable, num_inputs, axis,
+                            grad)
+    config = config_pb2.ConfigProto(graph_options=config_pb2.GraphOptions(
+        optimizer_options=config_pb2.OptimizerOptions(
+            opt_level=config_pb2.OptimizerOptions.L0)))
+    with session_lib.Session(graph=graph, config=config) as session:
+      variables.global_variables_initializer().run()
+      _ = session.run(outputs)  # warm up.
+      start_time = time.time()
+      for _ in range(num_iters):
+        _ = session.run(outputs)
+      duration = time.time() - start_time
+      print("%s shape:%d/%d var: %r #inputs:%d axis:%d grad:%r - %f secs - %f "
+            "GB/sec" % (device, input_shape[0], input_shape[1], variable,
+                        num_inputs, axis, grad, duration / num_iters,
+                        num_inputs * input_shape[0] * input_shape[1] * 4 * 2 *
+                        100 / (duration / num_iters) / 1e9))
+
+    name_template = (
+        "concat_bench_{device}_input_shape_{shape}_variable_{variable}"
+        "_num_inputs_{num_inputs}_axis_{axis}_grad_{grad}")
+
+    self.report_benchmark(name=name_template.format(
+        device=device,
+        num_inputs=num_inputs,
+        variable=variable,
+        grad=grad,
+        shape=str(input_shape).replace(" ", ""),
+        axis=str(axis),
+        iters=num_iters))
+
+    return duration
+
+  def benchmark_concat(self):
+    print("Forward vs backward concat")
+    shapes = [[2000, 8], [8, 2000], [100, 18], [1000, 18], [100, 97],
+              [1000, 97], [10000, 1], [1, 10000]]
+    axis_ = [0, 1]
+    num_inputs = 20
+    num_iters = [10] * len(shapes)
+    variable = [False, True]  # fixed input size or not
+    for shape, iters in zip(shapes, num_iters):
+      for axis in axis_:
+        for v in variable:
+          self._run_graph("cpu", shape, v, num_inputs, axis, True, iters)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/ops/cond.py

@@ -0,0 +1,409 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Cond function for Control Flow Operations."""
+
+from tensorflow.python.eager import context
+from tensorflow.python.eager.polymorphic_function import eager_function_run
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor as tensor_lib
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond_v2
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import control_flow_util as util
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.types import core
+from tensorflow.python.util import deprecation
+from tensorflow.python.util import dispatch
+from tensorflow.python.util import nest
+from tensorflow.python.util.tf_export import tf_export
+
+
+# pylint: disable=redefined-outer-name
+# pylint: disable=g-doc-args
+@tf_export(v1=["cond"])
+@dispatch.add_dispatch_support
+@deprecation.deprecated_args(
+    None, "fn1/fn2 are deprecated in favor of the true_fn/false_fn arguments.",
+    "fn1", "fn2")
+def cond(pred,
+         true_fn=None,
+         false_fn=None,
+         strict=False,
+         name=None,
+         fn1=None,
+         fn2=None):
+  """Return `true_fn()` if the predicate `pred` is true else `false_fn()`.
+
+  `true_fn` and `false_fn` both return lists of output tensors. `true_fn` and
+  `false_fn` must have the same non-zero number and type of outputs.
+
+  **WARNING**: Any Tensors or Operations created outside of `true_fn` and
+  `false_fn` will be executed regardless of which branch is selected at runtime.
+
+  Although this behavior is consistent with the dataflow model of TensorFlow,
+  it has frequently surprised users who expected a lazier semantics.
+  Consider the following simple program:
+
+  ```python
+  z = tf.multiply(a, b)
+  result = tf.cond(x < y, lambda: tf.add(x, z), lambda: tf.square(y))
+  ```
+
+  If `x < y`, the `tf.add` operation will be executed and `tf.square`
+  operation will not be executed. Since `z` is needed for at least one
+  branch of the `cond`, the `tf.multiply` operation is always executed,
+  unconditionally.
+
+  Note that `cond` calls `true_fn` and `false_fn` *exactly once* (inside the
+  call to `cond`, and not at all during `Session.run()`). `cond`
+  stitches together the graph fragments created during the `true_fn` and
+  `false_fn` calls with some additional graph nodes to ensure that the right
+  branch gets executed depending on the value of `pred`.
+
+  `tf.cond` supports nested structures as implemented in
+  `tensorflow.python.util.nest`. Both `true_fn` and `false_fn` must return the
+  same (possibly nested) value structure of lists, tuples, and/or named tuples.
+  Singleton lists and tuples form the only exceptions to this: when returned by
+  `true_fn` and/or `false_fn`, they are implicitly unpacked to single values.
+  This behavior is disabled by passing `strict=True`.
+
+  Args:
+    pred: A scalar determining whether to return the result of `true_fn` or
+      `false_fn`.
+    true_fn: The callable to be performed if pred is true.
+    false_fn: The callable to be performed if pred is false.
+    strict: A boolean that enables/disables 'strict' mode; see above.
+    name: Optional name prefix for the returned tensors.
+
+  Returns:
+    Tensors returned by the call to either `true_fn` or `false_fn`. If the
+    callables return a singleton list, the element is extracted from the list.
+
+  Raises:
+    TypeError: if `true_fn` or `false_fn` is not callable.
+    ValueError: if `true_fn` and `false_fn` do not return the same number of
+      tensors, or return tensors of different types.
+
+  Example:
+
+  ```python
+  x = tf.constant(2)
+  y = tf.constant(5)
+  def f1(): return tf.multiply(x, 17)
+  def f2(): return tf.add(y, 23)
+  r = tf.cond(tf.less(x, y), f1, f2)
+  # r is set to f1().
+  # Operations in f2 (e.g., tf.add) are not executed.
+  ```
+
+  """
+  # We needed to make true_fn/false_fn keyword arguments for
+  # backwards-compatibility. This check exists so that we can convert back to
+  # having them be positional arguments.
+  # TODO(josh11b): Make `true_fn` and `false_fn` positional arguments after
+  # `fn1` and `fn2` are deleted.
+  if fn1 is not None:
+    if true_fn is not None:
+      raise TypeError(
+          "cond(): 'true_fn' and 'fn1' may not be set simultaneously.")
+    true_fn = fn1
+  elif true_fn is None:
+    raise TypeError("cond(): 'true_fn' argument required")
+  if fn2 is not None:
+    if false_fn is not None:
+      raise TypeError(
+          "cond(): 'false_fn' and 'fn2' may not be set simultaneously.")
+    false_fn = fn2
+  elif false_fn is None:
+    raise TypeError("cond(): 'false_fn' argument required")
+
+  if not callable(true_fn):
+    raise TypeError("'true_fn' must be callable.")
+  if not callable(false_fn):
+    raise TypeError("'false_fn' must be callable.")
+
+  if context.executing_eagerly():
+    return _eager_cond_implementation(pred, true_fn, false_fn, strict, name)
+
+  # Always enable control flow v2 if building a function, regardless of toggle.
+  if util.EnableControlFlowV2(ops.get_default_graph()):
+    return cond_v2.cond_v2(pred, true_fn, false_fn, name)
+
+  with ops.name_scope(name, "cond", [pred]):
+    # Add the Switch to the graph.
+    if isinstance(pred, bool):
+      raise TypeError("'pred' must not be a Python bool.")
+    p_2, p_1 = control_flow_ops.switch(pred, pred)
+    pivot_1 = array_ops.identity(p_1, name="switch_t")
+    pivot_2 = array_ops.identity(p_2, name="switch_f")
+    pred = array_ops.identity(pred, name="pred_id")
+    # Disable the fetching of tensors that are only on one branch of cond.
+    for tensor in [p_1, p_2, pivot_1, pivot_2, pred]:
+      tensor.op.graph.prevent_fetching(tensor.op)
+
+    # Build the graph for the true branch in a new context.
+    context_t = control_flow_ops.CondContext(pred, pivot_1, branch=1)
+    try:
+      context_t.Enter()
+      orig_res_t, res_t = context_t.BuildCondBranch(true_fn)
+      if orig_res_t is None:
+        raise ValueError("'true_fn' must have a return value.")
+      context_t.ExitResult(res_t)
+    finally:
+      context_t.Exit()
+
+    # Build the graph for the false branch in a new context.
+    context_f = control_flow_ops.CondContext(pred, pivot_2, branch=0)
+    try:
+      context_f.Enter()
+      orig_res_f, res_f = context_f.BuildCondBranch(false_fn)
+      if orig_res_f is None:
+        raise ValueError("'false_fn' must have a return value.")
+      context_f.ExitResult(res_f)
+    finally:
+      context_f.Exit()
+
+    if not strict:
+      orig_res_t = _UnpackIfSingleton(orig_res_t)
+      orig_res_f = _UnpackIfSingleton(orig_res_f)
+
+    # Check that the return values of the two branches have the same structure.
+    try:
+      nest.assert_same_structure(orig_res_t, orig_res_f, expand_composites=True)
+    except (TypeError, ValueError):
+      nest.map_structure(_cast_indexed_slice_indices, orig_res_t, orig_res_f)
+      nest.map_structure(_cast_indexed_slice_indices, res_t, res_f)
+      try:
+        nest.assert_same_structure(orig_res_t, orig_res_f,
+                                   expand_composites=True)
+      except TypeError as e:
+        raise TypeError(
+            f"Incompatible return types of 'true_fn' and 'false_fn': {e}")
+      except ValueError as e:
+        raise ValueError(
+            f"Incompatible return values of 'true_fn' and 'false_fn': {e}")
+
+    # Add the final merge to the graph.
+    if not res_t:
+      raise ValueError(
+          "'true_fn' and 'false_fn' must return at least one result.")
+
+    res_t_flat = nest.flatten(res_t, expand_composites=True)
+    res_f_flat = nest.flatten(res_f, expand_composites=True)
+
+    for (x, y) in zip(res_t_flat, res_f_flat):
+      assert (
+          isinstance(x, tensor_lib.Tensor)
+          and isinstance(y, tensor_lib.Tensor))
+      if x.dtype.base_dtype != y.dtype.base_dtype:
+        raise ValueError(
+            "Outputs of 'true_fn' and 'false_fn' must have the same type(s). "
+            f"Received {x.dtype.name} from 'true_fn' "
+            f"and {y.dtype.name} from 'false_fn'.")
+
+    merges = [
+        control_flow_ops.merge(pair)[0] for pair in zip(res_f_flat, res_t_flat)]
+    merges = nest.map_structure(
+        control_flow_ops._convert_flow_to_tensorarray,  # pylint: disable=protected-access
+        nest.flatten(orig_res_t, expand_composites=True),
+        merges)
+
+    # Only add non-nested conds to the collection. Any nested control flow will
+    # be encapsulated in the root context.
+    assert context_t.outer_context == context_f.outer_context
+    if context_t.outer_context is None:
+      ops.add_to_collection(ops.GraphKeys.COND_CONTEXT, context_t)
+      ops.add_to_collection(ops.GraphKeys.COND_CONTEXT, context_f)
+
+    merges = nest.pack_sequence_as(
+        structure=orig_res_t, flat_sequence=merges, expand_composites=True)
+
+    # Singleton lists and tuples are automatically unpacked if strict == False.
+    if not strict:
+      merges = _UnpackIfSingleton(merges)
+    return merges
+
+
+@tf_export("cond", v1=[])
+@dispatch.add_dispatch_support
+def cond_for_tf_v2(pred, true_fn=None, false_fn=None, name=None):
+  """Return `true_fn()` if the predicate `pred` is true else `false_fn()`.
+
+  Note: This op is automatically used in a `tf.function` to convert Python
+  if-statements when the predicate is a `tf.Tensor`, unless `autograph=False` is
+  explicitly specified in `tf.function` args. For example, the following are
+  equivalent:
+
+  >>> @tf.function
+  ... def fun1(x,y):
+  ...   if x > 0:  # AutoGraph converts if-statement to tf.cond().
+  ...     z = y+1
+  ...   else:
+  ...     z = y-1
+  ...   return z
+  >>> print(fun1(tf.constant(7), tf.constant(3)).numpy())
+  4
+
+  >>> @tf.function
+  ... def fun2(x,y):
+  ...   pred = x > 0
+  ...   true_fn =  lambda: y+1
+  ...   false_fn = lambda: y-1
+  ...   return tf.cond(pred, true_fn, false_fn)  # Use tf.cond() explicitly.
+  >>> print(fun1(tf.constant(7), tf.constant(3)).numpy())
+  4
+
+  For more information, see [tf.function and AutoGraph guide](
+  https://www.tensorflow.org/guide/function#autograph_transformations).
+
+  `true_fn` and `false_fn` both return lists of output tensors. `true_fn` and
+  `false_fn` must have the same non-zero number and type of outputs.
+
+  **WARNING**: Any Tensors or Operations created outside of `true_fn` and
+  `false_fn` will be executed regardless of which branch is selected at runtime.
+
+  Although this behavior is consistent with the dataflow model of TensorFlow,
+  it has frequently surprised users who expected a lazier semantics.
+  Consider the following simple program:
+
+  >>> x, y = tf.constant(2, dtype=tf.int32), tf.constant(4, dtype=tf.int32)
+  >>> z = tf.multiply(x, y)
+  >>> r = tf.cond(x < y, lambda: tf.add(x, z), lambda: tf.square(y))
+  >>> print(r.numpy())
+  10
+
+  If `x < y`, the `tf.add` operation will be executed and `tf.square`
+  operation will not be executed. Since `z` is needed for at least one
+  branch of the `cond`, the `tf.multiply` operation is always executed,
+  unconditionally.
+
+  Note that `cond` calls `true_fn` and `false_fn` *exactly once* (inside the
+  call to `cond`, and not at all during `Session.run()`). `cond`
+  stitches together the graph fragments created during the `true_fn` and
+  `false_fn` calls with some additional graph nodes to ensure that the right
+  branch gets executed depending on the value of `pred`.
+
+  `tf.cond` supports nested structures as implemented in
+  `tensorflow.python.util.nest`. Both `true_fn` and `false_fn` must return the
+  same (possibly nested) value structure of lists, tuples, and/or named tuples.
+  Singleton lists and tuples form the only exceptions to this: when returned by
+  `true_fn` and/or `false_fn`, they are implicitly unpacked to single values.
+
+  Note: It is illegal to "directly" use tensors created inside a cond branch
+  outside it, e.g. by storing a reference to a branch tensor in the python
+  state. If you need to use a tensor created in a branch function you should
+  return it as an output of the branch function and use the output from
+  `tf.cond` instead.
+
+  Args:
+    pred: A scalar determining whether to return the result of `true_fn` or
+      `false_fn`.
+    true_fn: The callable to be performed if pred is true.
+    false_fn: The callable to be performed if pred is false.
+    name: Optional name prefix for the returned tensors.
+
+  Returns:
+    Tensors returned by the call to either `true_fn` or `false_fn`. If the
+    callables return a singleton list, the element is extracted from the list.
+
+  Raises:
+    TypeError: if `true_fn` or `false_fn` is not callable.
+    ValueError: if `true_fn` and `false_fn` do not return the same number of
+      tensors, or return tensors of different types.
+
+  Example:
+
+  >>> x = tf.constant(2)
+  >>> y = tf.constant(5)
+  >>> def f1(): return tf.multiply(x, 7)
+  >>> def f2(): return tf.add(y, 3)
+  >>> r = tf.cond(tf.less(x, y), f1, f2)
+  >>> # r is set to f1().
+  >>> # Operations in f2 (e.g., tf.add) are not executed.
+  >>> print(r.numpy())
+  14
+
+  """
+  return cond(pred, true_fn=true_fn, false_fn=false_fn, strict=True, name=name)
+
+
+def _UnpackIfSingleton(res):
+  if isinstance(res, (list, tuple)) and len(res) == 1:
+    return res[0]
+  else:
+    return res
+
+
+def _eager_cond_implementation(pred, true_fn, false_fn, strict, name):
+  """Special cases for `cond` when executing eagerly."""
+  pred = ops.convert_to_tensor(pred)
+  pred_constant_value = tensor_util.constant_value(pred)
+  if pred_constant_value is None:
+    # Eager tensors from a parallel device may not have a constant
+    # value. Running the cond op itself would work, but we don't have logic to
+    # build cond ops without wrapping in a function first.
+    if (not isinstance(true_fn, core.PolymorphicFunction)
+        or not isinstance(false_fn, core.PolymorphicFunction)):
+      raise TypeError("When running tf.cond on a parallel device, 'true_fn' "
+                      "and 'false_fn' must be decorated with `tf.function`.")
+    functions_run_eagerly = eager_function_run.functions_run_eagerly()
+    if functions_run_eagerly:
+      # We need to use tf.function to deal with variable creation inside the
+      # cond, and skipping it because of run_functions_eagerly would just
+      # crash immediately.
+      logging.warning(
+          "It looks like tf.function behavior was disabled, perhaps using "
+          "tf.config.run_functions_eagerly. Parallelized tf.cond requires "
+          "tf.function to work. This primitive will override the disable.")
+    eager_function_run.run_functions_eagerly(False)
+    try:
+      return cond_v2.cond_v2(pred, true_fn, false_fn, name)
+    finally:
+      if functions_run_eagerly is not None:
+        eager_function_run.run_functions_eagerly(functions_run_eagerly)
+  else:
+    # For conditions which are eager tensors with a constant value (most of
+    # them), we only call the relevant branch function and execute it eagerly.
+    with ops.name_scope(name, "cond", [pred]):
+      if pred_constant_value:
+        result = true_fn()
+      else:
+        result = false_fn()
+      if not strict:
+        result = _UnpackIfSingleton(result)
+      return result
+
+
+def _cast_indexed_slice_indices(a, b):
+  """Cast IndexedSlice.indices from int32 to int64 where necessary.
+
+  If `a` and `b` are both IndexedSlices, and their indices have different
+  dtypes, then cast both their dtypes to `int64` (modifies `a` and `b`
+  in-place).  Otherwise, does nothing.
+
+  Args:
+    a: A value, which may be an IndexedSlices.
+    b: A value, which may be an IndexedSlices.
+  """
+  if (isinstance(a, indexed_slices.IndexedSlices) and
+      isinstance(b, indexed_slices.IndexedSlices) and
+      a.indices.dtype != b.indices.dtype):
+    # pylint: disable=protected-access
+    a._indices = math_ops.cast(a.indices, dtypes.int64)
+    b._indices = math_ops.cast(b.indices, dtypes.int64)

tensorflow/python/ops/confusion_matrix.py

@@ -0,0 +1,259 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Confusion matrix related utilities."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.util import deprecation
+from tensorflow.python.util import dispatch
+from tensorflow.python.util.tf_export import tf_export
+
+
+def remove_squeezable_dimensions(
+    labels, predictions, expected_rank_diff=0, name=None):
+  """Squeeze last dim if ranks differ from expected by exactly 1.
+
+  In the common case where we expect shapes to match, `expected_rank_diff`
+  defaults to 0, and we squeeze the last dimension of the larger rank if they
+  differ by 1.
+
+  But, for example, if `labels` contains class IDs and `predictions` contains 1
+  probability per class, we expect `predictions` to have 1 more dimension than
+  `labels`, so `expected_rank_diff` would be 1. In this case, we'd squeeze
+  `labels` if `rank(predictions) - rank(labels) == 0`, and
+  `predictions` if `rank(predictions) - rank(labels) == 2`.
+
+  This will use static shape if available. Otherwise, it will add graph
+  operations, which could result in a performance hit.
+
+  Args:
+    labels: Label values, a `Tensor` whose dimensions match `predictions`.
+    predictions: Predicted values, a `Tensor` of arbitrary dimensions.
+    expected_rank_diff: Expected result of `rank(predictions) - rank(labels)`.
+    name: Name of the op.
+
+  Returns:
+    Tuple of `labels` and `predictions`, possibly with last dim squeezed.
+  """
+  with ops.name_scope(name, 'remove_squeezable_dimensions',
+                      [labels, predictions]):
+    predictions = ops.convert_to_tensor(predictions)
+    labels = ops.convert_to_tensor(labels)
+    predictions_shape = predictions.get_shape()
+    predictions_rank = predictions_shape.ndims
+    labels_shape = labels.get_shape()
+    labels_rank = labels_shape.ndims
+    if (labels_rank is not None) and (predictions_rank is not None):
+      # Use static rank.
+      rank_diff = predictions_rank - labels_rank
+      if (rank_diff == expected_rank_diff + 1 and
+          predictions_shape.dims[-1].is_compatible_with(1)):
+        predictions = array_ops.squeeze(predictions, [-1])
+      elif (rank_diff == expected_rank_diff - 1 and
+            labels_shape.dims[-1].is_compatible_with(1)):
+        labels = array_ops.squeeze(labels, [-1])
+      return labels, predictions
+
+    # Use dynamic rank.
+    rank_diff = array_ops.rank(predictions) - array_ops.rank(labels)
+    if (predictions_rank is None) or (
+        predictions_shape.dims[-1].is_compatible_with(1)):
+      predictions = cond.cond(
+          math_ops.equal(expected_rank_diff + 1, rank_diff),
+          lambda: array_ops.squeeze(predictions, [-1]),
+          lambda: predictions)
+    if (labels_rank is None) or (
+        labels_shape.dims[-1].is_compatible_with(1)):
+      labels = cond.cond(
+          math_ops.equal(expected_rank_diff - 1, rank_diff),
+          lambda: array_ops.squeeze(labels, [-1]),
+          lambda: labels)
+    return labels, predictions
+
+
+@tf_export('math.confusion_matrix', v1=[])
+@dispatch.add_dispatch_support
+def confusion_matrix(labels,
+                     predictions,
+                     num_classes=None,
+                     weights=None,
+                     dtype=dtypes.int32,
+                     name=None):
+  """Computes the confusion matrix from predictions and labels.
+
+  The matrix columns represent the prediction labels and the rows represent the
+  real labels. The confusion matrix is always a 2-D array of shape `[n, n]`,
+  where `n` is the number of valid labels for a given classification task. Both
+  prediction and labels must be 1-D arrays of the same shape in order for this
+  function to work.
+
+  If `num_classes` is `None`, then `num_classes` will be set to one plus the
+  maximum value in either predictions or labels. Class labels are expected to
+  start at 0. For example, if `num_classes` is 3, then the possible labels
+  would be `[0, 1, 2]`.
+
+  If `weights` is not `None`, then each prediction contributes its
+  corresponding weight to the total value of the confusion matrix cell.
+
+  For example:
+
+  ```python
+    tf.math.confusion_matrix([1, 2, 4], [2, 2, 4]) ==>
+        [[0 0 0 0 0]
+         [0 0 1 0 0]
+         [0 0 1 0 0]
+         [0 0 0 0 0]
+         [0 0 0 0 1]]
+  ```
+
+  Note that the possible labels are assumed to be `[0, 1, 2, 3, 4]`,
+  resulting in a 5x5 confusion matrix.
+
+  Args:
+    labels: 1-D `Tensor` of real labels for the classification task.
+    predictions: 1-D `Tensor` of predictions for a given classification.
+    num_classes: The possible number of labels the classification task can
+                 have. If this value is not provided, it will be calculated
+                 using both predictions and labels array.
+    weights: An optional `Tensor` whose shape matches `predictions`.
+    dtype: Data type of the confusion matrix.
+    name: Scope name.
+
+  Returns:
+    A `Tensor` of type `dtype` with shape `[n, n]` representing the confusion
+    matrix, where `n` is the number of possible labels in the classification
+    task.
+
+  Raises:
+    ValueError: If both predictions and labels are not 1-D vectors and have
+      mismatched shapes, or if `weights` is not `None` and its shape doesn't
+      match `predictions`.
+  """
+  with ops.name_scope(name, 'confusion_matrix',
+                      (predictions, labels, num_classes, weights)) as name:
+    labels, predictions = remove_squeezable_dimensions(
+        ops.convert_to_tensor(labels, name='labels'),
+        ops.convert_to_tensor(
+            predictions, name='predictions'))
+    predictions = math_ops.cast(predictions, dtypes.int64)
+    labels = math_ops.cast(labels, dtypes.int64)
+
+    # Sanity checks - underflow or overflow can cause memory corruption.
+    labels = control_flow_ops.with_dependencies(
+        [check_ops.assert_non_negative(
+            labels, message='`labels` contains negative values')],
+        labels)
+    predictions = control_flow_ops.with_dependencies(
+        [check_ops.assert_non_negative(
+            predictions, message='`predictions` contains negative values')],
+        predictions)
+
+    if num_classes is None:
+      num_classes = math_ops.maximum(math_ops.reduce_max(predictions),
+                                     math_ops.reduce_max(labels)) + 1
+    else:
+      num_classes_int64 = math_ops.cast(num_classes, dtypes.int64)
+      labels = control_flow_ops.with_dependencies(
+          [check_ops.assert_less(
+              labels, num_classes_int64, message='`labels` out of bound')],
+          labels)
+      predictions = control_flow_ops.with_dependencies(
+          [check_ops.assert_less(
+              predictions, num_classes_int64,
+              message='`predictions` out of bound')],
+          predictions)
+
+    if weights is not None:
+      weights = ops.convert_to_tensor(weights, name='weights')
+      predictions.get_shape().assert_is_compatible_with(weights.get_shape())
+      weights = math_ops.cast(weights, dtype)
+
+    shape = array_ops_stack.stack([num_classes, num_classes])
+    indices = array_ops_stack.stack([labels, predictions], axis=1)
+    values = (array_ops.ones_like(predictions, dtype)
+              if weights is None else weights)
+    return array_ops.scatter_nd(
+        indices=indices,
+        updates=values,
+        shape=math_ops.cast(shape, dtypes.int64))
+
+
+@tf_export(v1=['math.confusion_matrix', 'confusion_matrix'])
+@dispatch.add_dispatch_support
+@deprecation.deprecated_endpoints('confusion_matrix', 'train.confusion_matrix')
+def confusion_matrix_v1(labels,
+                        predictions,
+                        num_classes=None,
+                        dtype=dtypes.int32,
+                        name=None,
+                        weights=None):
+  """Computes the confusion matrix from predictions and labels.
+
+  The matrix columns represent the prediction labels and the rows represent the
+  real labels. The confusion matrix is always a 2-D array of shape `[n, n]`,
+  where `n` is the number of valid labels for a given classification task. Both
+  prediction and labels must be 1-D arrays of the same shape in order for this
+  function to work.
+
+  If `num_classes` is `None`, then `num_classes` will be set to one plus the
+  maximum value in either predictions or labels. Class labels are expected to
+  start at 0. For example, if `num_classes` is 3, then the possible labels
+  would be `[0, 1, 2]`.
+
+  If `weights` is not `None`, then each prediction contributes its
+  corresponding weight to the total value of the confusion matrix cell.
+
+  For example:
+
+  ```python
+    tf.math.confusion_matrix([1, 2, 4], [2, 2, 4]) ==>
+        [[0 0 0 0 0]
+         [0 0 1 0 0]
+         [0 0 1 0 0]
+         [0 0 0 0 0]
+         [0 0 0 0 1]]
+  ```
+
+  Note that the possible labels are assumed to be `[0, 1, 2, 3, 4]`,
+  resulting in a 5x5 confusion matrix.
+
+  Args:
+    labels: 1-D `Tensor` of real labels for the classification task.
+    predictions: 1-D `Tensor` of predictions for a given classification.
+    num_classes: The possible number of labels the classification task can have.
+      If this value is not provided, it will be calculated using both
+      predictions and labels array.
+    dtype: Data type of the confusion matrix.
+    name: Scope name.
+    weights: An optional `Tensor` whose shape matches `predictions`.
+
+  Returns:
+    A `Tensor` of type `dtype` with shape `[n, n]` representing the confusion
+    matrix, where `n` is the number of possible labels in the classification
+    task.
+
+  Raises:
+    ValueError: If both predictions and labels are not 1-D vectors and have
+      mismatched shapes, or if `weights` is not `None` and its shape doesn't
+      match `predictions`.
+  """
+  return confusion_matrix(labels, predictions, num_classes, weights, dtype,
+                          name)

tensorflow/python/ops/control_flow_assert.py

@@ -0,0 +1,130 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Assert functions for Control Flow Operations."""
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import gen_control_flow_ops
+from tensorflow.python.ops import gen_logging_ops
+from tensorflow.python.ops import gen_math_ops
+from tensorflow.python.util import dispatch
+from tensorflow.python.util import tf_should_use
+from tensorflow.python.util.tf_export import tf_export
+
+
+def _summarize_eager(tensor, summarize=None):
+  """Returns a summarized string representation of eager `tensor`.
+
+  Args:
+    tensor: EagerTensor to summarize
+    summarize: Include these many first elements of `array`
+  """
+  # Emulate the behavior of Tensor::SummarizeValue()
+  if summarize is None:
+    summarize = 3
+  elif summarize < 0:
+    summarize = array_ops.size(tensor)
+
+  # reshape((-1,)) is the fastest way to get a flat array view
+  if tensor._rank():  # pylint: disable=protected-access
+    flat = tensor.numpy().reshape((-1,))
+    lst = [str(x) for x in flat[:summarize]]
+    if len(lst) < flat.size:
+      lst.append("...")
+  else:
+    # tensor.numpy() returns a scalar for zero dimensional arrays
+    if gen_math_ops.not_equal(summarize, 0):
+      lst = [str(tensor.numpy())]
+    else:
+      lst = []
+
+  return ", ".join(lst)
+
+
+# Assert and Print are special symbols in python, so we must
+# use an upper-case version of them.
+@tf_export("debugging.Assert", "Assert")
+@dispatch.add_dispatch_support
+@tf_should_use.should_use_result
+def Assert(condition, data, summarize=None, name=None):
+  """Asserts that the given condition is true.
+
+  If `condition` evaluates to false, print the list of tensors in `data`.
+  `summarize` determines how many entries of the tensors to print.
+
+  Args:
+    condition: The condition to evaluate.
+    data: The tensors to print out when condition is false.
+    summarize: Print this many entries of each tensor.
+    name: A name for this operation (optional).
+
+  Returns:
+    assert_op: An `Operation` that, when executed, raises a
+    `tf.errors.InvalidArgumentError` if `condition` is not true.
+    @compatibility(eager)
+    returns None
+    @end_compatibility
+
+  Raises:
+    @compatibility(TF1)
+    When in TF V1 mode (that is, outside `tf.function`) Assert needs a control
+    dependency on the output to ensure the assertion executes:
+
+  ```python
+  # Ensure maximum element of x is smaller or equal to 1
+  assert_op = tf.Assert(tf.less_equal(tf.reduce_max(x), 1.), [x])
+  with tf.control_dependencies([assert_op]):
+    ... code using x ...
+  ```
+
+    @end_compatibility
+  """
+  if context.executing_eagerly():
+    if not condition:
+      xs = ops.convert_n_to_tensor(data)
+      data_str = [_summarize_eager(x, summarize) for x in xs]
+      raise errors.InvalidArgumentError(
+          node_def=None,
+          op=None,
+          message="Expected '%s' to be true. Summarized data: %s" %
+          (condition, "\n".join(data_str)))
+    return
+
+  with ops.name_scope(name, "Assert", [condition, data]) as name:
+    xs = ops.convert_n_to_tensor(data)
+    if all(x.dtype in {dtypes.string, dtypes.int32} for x in xs):
+      # As a simple heuristic, we assume that string and int32 are
+      # on host to avoid the need to use cond. If it is not case,
+      # we will pay the price copying the tensor to host memory.
+      return gen_logging_ops._assert(condition, data, summarize, name="Assert")  # pylint: disable=protected-access
+    else:
+      condition = ops.convert_to_tensor(condition, name="Condition")
+
+      def true_assert():
+        return gen_logging_ops._assert(  # pylint: disable=protected-access
+            condition, data, summarize, name="Assert")
+
+      guarded_assert = cond.cond(
+          condition,
+          gen_control_flow_ops.no_op,
+          true_assert,
+          name="AssertGuard")
+      if context.executing_eagerly():
+        return
+      return guarded_assert.op

tensorflow/python/ops/control_flow_case.py

@@ -0,0 +1,417 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Case functions for Control Flow Operations."""
+
+import collections
+import functools
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import control_flow_assert
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util import dispatch
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("case", v1=[])
+@dispatch.add_dispatch_support
+def case_v2(pred_fn_pairs,
+            default=None,
+            exclusive=False,
+            strict=False,
+            name="case"):
+  """Create a case operation.
+
+  See also `tf.switch_case`.
+
+  The `pred_fn_pairs` parameter is a list of pairs of size N.
+  Each pair contains a boolean scalar tensor and a python callable that
+  creates the tensors to be returned if the boolean evaluates to True.
+  `default` is a callable generating a list of tensors. All the callables
+  in `pred_fn_pairs` as well as `default` (if provided) should return the same
+  number and types of tensors.
+
+  If `exclusive==True`, all predicates are evaluated, and an exception is
+  thrown if more than one of the predicates evaluates to `True`.
+  If `exclusive==False`, execution stops at the first predicate which
+  evaluates to True, and the tensors generated by the corresponding function
+  are returned immediately. If none of the predicates evaluate to True, this
+  operation returns the tensors generated by `default`.
+
+  `tf.case` supports nested structures as implemented in
+  `tf.nest`. All of the callables must return the same (possibly nested) value
+  structure of lists, tuples, and/or named tuples. Singleton lists and tuples
+  form the only exceptions to this: when returned by a callable, they are
+  implicitly unpacked to single values. This behavior is disabled by passing
+  `strict=True`.
+
+  @compatibility(v2)
+  `pred_fn_pairs` could be a dictionary in v1. However, tf.Tensor and
+  tf.Variable are no longer hashable in v2, so cannot be used as a key for a
+  dictionary.  Please use a list or a tuple instead.
+  @end_compatibility
+
+
+  **Example 1:**
+
+  Pseudocode:
+
+  ```
+  if (x < y) return 17;
+  else return 23;
+  ```
+
+  Expressions:
+
+  ```python
+  f1 = lambda: tf.constant(17)
+  f2 = lambda: tf.constant(23)
+  r = tf.case([(tf.less(x, y), f1)], default=f2)
+  ```
+
+  **Example 2:**
+
+  Pseudocode:
+
+  ```
+  if (x < y && x > z) raise OpError("Only one predicate may evaluate to True");
+  if (x < y) return 17;
+  else if (x > z) return 23;
+  else return -1;
+  ```
+
+  Expressions:
+
+  ```python
+  def f1(): return tf.constant(17)
+  def f2(): return tf.constant(23)
+  def f3(): return tf.constant(-1)
+  r = tf.case([(tf.less(x, y), f1), (tf.greater(x, z), f2)],
+           default=f3, exclusive=True)
+  ```
+
+  Args:
+    pred_fn_pairs: List of pairs of a boolean scalar tensor and a callable which
+      returns a list of tensors.
+    default: Optional callable that returns a list of tensors.
+    exclusive: True iff at most one predicate is allowed to evaluate to `True`.
+    strict: A boolean that enables/disables 'strict' mode; see above.
+    name: A name for this operation (optional).
+
+  Returns:
+    The tensors returned by the first pair whose predicate evaluated to True, or
+    those returned by `default` if none does.
+
+  Raises:
+    TypeError: If `pred_fn_pairs` is not a list/tuple.
+    TypeError: If `pred_fn_pairs` is a list but does not contain 2-tuples.
+    TypeError: If `fns[i]` is not callable for any i, or `default` is not
+               callable.
+  """
+  return _case_helper(
+      cond.cond,
+      pred_fn_pairs,
+      default,
+      exclusive,
+      name,
+      allow_python_preds=False,
+      strict=strict)
+
+
+@tf_export(v1=["case"])
+@dispatch.add_dispatch_support
+def case(pred_fn_pairs,
+         default=None,
+         exclusive=False,
+         strict=False,
+         name="case"):
+  """Create a case operation.
+
+  See also `tf.switch_case`.
+
+  The `pred_fn_pairs` parameter is a dict or list of pairs of size N.
+  Each pair contains a boolean scalar tensor and a python callable that
+  creates the tensors to be returned if the boolean evaluates to True.
+  `default` is a callable generating a list of tensors. All the callables
+  in `pred_fn_pairs` as well as `default` (if provided) should return the same
+  number and types of tensors.
+
+  If `exclusive==True`, all predicates are evaluated, and an exception is
+  thrown if more than one of the predicates evaluates to `True`.
+  If `exclusive==False`, execution stops at the first predicate which
+  evaluates to True, and the tensors generated by the corresponding function
+  are returned immediately. If none of the predicates evaluate to True, this
+  operation returns the tensors generated by `default`.
+
+  `tf.case` supports nested structures as implemented in
+  `tf.nest`. All of the callables must return the same (possibly nested) value
+  structure of lists, tuples, and/or named tuples. Singleton lists and tuples
+  form the only exceptions to this: when returned by a callable, they are
+  implicitly unpacked to single values. This behavior is disabled by passing
+  `strict=True`.
+
+  If an unordered dictionary is used for `pred_fn_pairs`, the order of the
+  conditional tests is not guaranteed. However, the order is guaranteed to be
+  deterministic, so that variables created in conditional branches are created
+  in fixed order across runs.
+
+  @compatibility(eager)
+  Unordered dictionaries are not supported in eager mode when `exclusive=False`.
+  Use a list of tuples instead.
+  @end_compatibility
+
+
+  **Example 1:**
+
+  Pseudocode:
+
+  ```
+  if (x < y) return 17;
+  else return 23;
+  ```
+
+  Expressions:
+
+  ```python
+  f1 = lambda: tf.constant(17)
+  f2 = lambda: tf.constant(23)
+  r = tf.case([(tf.less(x, y), f1)], default=f2)
+  ```
+
+  **Example 2:**
+
+  Pseudocode:
+
+  ```
+  if (x < y && x > z) raise OpError("Only one predicate may evaluate to True");
+  if (x < y) return 17;
+  else if (x > z) return 23;
+  else return -1;
+  ```
+
+  Expressions:
+
+  ```python
+  def f1(): return tf.constant(17)
+  def f2(): return tf.constant(23)
+  def f3(): return tf.constant(-1)
+  r = tf.case({tf.less(x, y): f1, tf.greater(x, z): f2},
+           default=f3, exclusive=True)
+  ```
+
+  Args:
+    pred_fn_pairs: Dict or list of pairs of a boolean scalar tensor and a
+      callable which returns a list of tensors.
+    default: Optional callable that returns a list of tensors.
+    exclusive: True iff at most one predicate is allowed to evaluate to `True`.
+    strict: A boolean that enables/disables 'strict' mode; see above.
+    name: A name for this operation (optional).
+
+  Returns:
+    The tensors returned by the first pair whose predicate evaluated to True, or
+    those returned by `default` if none does.
+
+  Raises:
+    TypeError: If `pred_fn_pairs` is not a list/dictionary.
+    TypeError: If `pred_fn_pairs` is a list but does not contain 2-tuples.
+    TypeError: If `fns[i]` is not callable for any i, or `default` is not
+               callable.
+  """
+  return _case_helper(
+      cond.cond,
+      pred_fn_pairs,
+      default,
+      exclusive,
+      name,
+      allow_python_preds=False,
+      strict=strict)
+
+
+def _assert_at_most_n_true(predicates, n, msg):
+  """Returns an Assert op that checks that at most n predicates are True.
+
+  Args:
+    predicates: list of bool scalar tensors.
+    n: maximum number of true predicates allowed.
+    msg: Error message.
+  """
+  preds_c = array_ops_stack.stack(predicates, name="preds_c")
+  num_true_conditions = math_ops.reduce_sum(
+      math_ops.cast(preds_c, dtypes.int32), name="num_true_conds")
+  condition = math_ops.less_equal(num_true_conditions,
+                                  constant_op.constant(n, name="n_true_conds"))
+  preds_names = ", ".join(getattr(p, "name", "?") for p in predicates)
+  error_msg = [
+      "%s: more than %d conditions (%s) evaluated as True:" %
+      (msg, n, preds_names), preds_c
+  ]
+  return control_flow_assert.Assert(
+      condition, data=error_msg, summarize=len(predicates))
+
+
+def _case_create_default_action(predicates, actions):
+  """Creates default action for a list of actions and their predicates.
+
+  It uses the input actions to select an arbitrary as default and makes sure
+  that corresponding predicates have valid values.
+
+  Args:
+    predicates: a list of bool scalar tensors
+    actions: a list of callable objects which return tensors.
+
+  Returns:
+    a callable
+  """
+  k = len(predicates) - 1  # could pick any
+  predicate, action = predicates[k], actions[k]
+  other_predicates, other_actions = predicates[:k], actions[:k]
+
+  def default_action():
+    others_msg = ("Implementation error: "
+                  "selected default action #%d was called, but some of other "
+                  "predicates are True: " % k)
+    default_msg = ("Input error: "
+                   "None of conditions evaluated as True:",
+                   array_ops_stack.stack(predicates, name="preds_c"))
+    with ops.control_dependencies([
+        _assert_at_most_n_true(  # pylint: disable=protected-access
+            other_predicates, n=0, msg=others_msg),
+        control_flow_assert.Assert(predicate, data=default_msg)
+    ]):
+      return action()
+
+  return default_action, other_predicates, other_actions
+
+
+def _case_helper(cond_fn,
+                 pred_fn_pairs,
+                 default,
+                 exclusive,
+                 name,
+                 allow_python_preds=False,
+                 **cond_kwargs):
+  """Implementation of case that allows for different cond functions.
+
+  Args:
+    cond_fn: method that has signature and semantics of `cond` above.
+    pred_fn_pairs: Dict or list of pairs of a boolean scalar tensor, and a
+      callable which returns a list of tensors.
+    default: Optional callable that returns a list of tensors.
+    exclusive: True iff at most one predicate is allowed to evaluate to `True`.
+    name: A name for this operation (optional).
+    allow_python_preds: if true, pred_fn_pairs may contain Python bools in
+      addition to boolean Tensors
+    **cond_kwargs: keyword arguments that will be passed to `cond_fn`.
+
+  Returns:
+    The tensors returned by the first pair whose predicate evaluated to True, or
+    those returned by `default` if none does.
+
+  Raises:
+    TypeError: If `pred_fn_pairs` is not a list/dictionary.
+    TypeError: If `pred_fn_pairs` is a list but does not contain 2-tuples.
+    TypeError: If `fns[i]` is not callable for any i, or `default` is not
+               callable.
+  """
+  predicates, actions = _case_verify_and_canonicalize_args(
+      pred_fn_pairs, exclusive, name, allow_python_preds)
+  with ops.name_scope(name, "case", [predicates]):
+    if default is None:
+      default, predicates, actions = _case_create_default_action(
+          predicates, actions)
+    fn = default
+    # To eval conditions in direct order we create nested conditions in reverse:
+    #   cond_fn(c[0], true_fn=.., false_fn=cond_fn(c[1], ...))
+    for predicate, action in reversed(list(zip(predicates, actions))):
+      fn = functools.partial(
+          cond_fn, predicate, true_fn=action, false_fn=fn, **cond_kwargs)
+    if exclusive:
+      with ops.control_dependencies([
+          _assert_at_most_n_true(  # pylint: disable=protected-access
+              predicates, n=1, msg="Input error: exclusive=True")
+      ]):
+        return fn()
+    else:
+      return fn()
+
+
+def _case_verify_and_canonicalize_args(pred_fn_pairs, exclusive, name,
+                                       allow_python_preds):
+  """Verifies input arguments for the case function.
+
+  Args:
+    pred_fn_pairs: Dict or list of pairs of a boolean scalar tensor, and a
+      callable which returns a list of tensors.
+    exclusive: True iff at most one predicate is allowed to evaluate to `True`.
+    name: A name for the case operation.
+    allow_python_preds: if true, pred_fn_pairs may contain Python bools in
+      addition to boolean Tensors
+
+  Raises:
+    TypeError: If `pred_fn_pairs` is not a list/dictionary.
+    TypeError: If `pred_fn_pairs` is a list but does not contain 2-tuples.
+    TypeError: If `fns[i]` is not callable for any i, or `default` is not
+               callable.
+
+  Returns:
+    a tuple <list of scalar bool tensors, list of callables>.
+  """
+  if not isinstance(pred_fn_pairs, (list, tuple, dict)):
+    raise TypeError("'pred_fn_pairs' must be a list, tuple, or dict. "
+                    f"Received: {type(pred_fn_pairs)}")
+
+  if isinstance(pred_fn_pairs, collections.OrderedDict):
+    pred_fn_pairs = pred_fn_pairs.items()
+  elif isinstance(pred_fn_pairs, dict):
+    if context.executing_eagerly():
+      # No name to sort on in eager mode. Use dictionary traversal order,
+      # which is nondeterministic in versions of Python < 3.6
+      if not exclusive:
+        raise ValueError("Unordered dictionaries are not supported for the "
+                         "'pred_fn_pairs' argument when `exclusive=False` and "
+                         "eager mode is enabled.")
+      pred_fn_pairs = list(pred_fn_pairs.items())
+    else:
+      pred_fn_pairs = sorted(
+          pred_fn_pairs.items(), key=lambda item: item[0].name)
+      if not exclusive:
+        logging.warn(
+            "%s: An unordered dictionary of predicate/fn pairs was "
+            "provided, but exclusive=False. The order of conditional "
+            "tests is deterministic but not guaranteed.", name)
+  for pred_fn_pair in pred_fn_pairs:
+    if not isinstance(pred_fn_pair, tuple) or len(pred_fn_pair) != 2:
+      raise TypeError("Each entry in 'pred_fn_pairs' must be a 2-tuple. "
+                      f"Received {pred_fn_pair}.")
+    pred, fn = pred_fn_pair
+
+    if isinstance(pred, tensor.Tensor):
+      if pred.dtype != dtypes.bool:
+        raise TypeError("pred must be Tensor of type bool: %s" % pred.name)
+    elif not allow_python_preds:
+      raise TypeError("pred must be a Tensor, got: %s" % pred)
+    elif not isinstance(pred, bool):
+      raise TypeError("pred must be a Tensor or bool, got: %s" % pred)
+
+    if not callable(fn):
+      raise TypeError("fn for pred %s must be callable." % pred.name)
+
+  predicates, actions = zip(*pred_fn_pairs)
+  return predicates, actions

tensorflow/python/ops/control_flow_grad.py

@@ -0,0 +1,247 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Gradients for operators defined in control_flow_ops.py."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import control_flow_util
+from tensorflow.python.ops import math_ops
+# go/tf-wildcard-import
+# pylint: disable=wildcard-import,undefined-variable,redefined-builtin
+from tensorflow.python.ops.control_flow_ops import *
+# pylint: enable=wildcard-import
+
+
+def _SwitchGrad(op, *grad):
+  """Gradients for a Switch op is calculated using a Merge op.
+
+  If the switch is a loop switch, it will be visited twice. We create
+  the merge on the first visit, and update the other input of the merge
+  on the second visit. A next_iteration is also added on second visit.
+  """
+  graph = ops.get_default_graph()
+  # pylint: disable=protected-access
+  op_ctxt = op._get_control_flow_context()
+  grad_ctxt = graph._get_control_flow_context()
+  # pylint: enable=protected-access
+  if isinstance(op_ctxt, WhileContext):
+    merge_grad = grad_ctxt.grad_state.switch_map.get(op)
+    if merge_grad is not None:
+      # This is the second time this Switch is visited. It comes from
+      # the non-exit branch of the Switch, so update the second input
+      # to the Merge.
+      # TODO(yuanbyu): Perform shape inference with this new input.
+      if grad[1] is not None:
+        # pylint: disable=protected-access
+        control_flow_ops._AddNextAndBackEdge(merge_grad, grad[1],
+                                             enforce_shape_invariant=False)
+        # pylint: enable=protected-access
+      return None, None
+    elif grad[0] is not None:
+      # This is the first time this Switch is visited. It comes from
+      # the Exit branch, which is grad[0]. grad[1] is empty at this point.
+      # Use grad[0] for both inputs to merge for now, but update the second
+      # input of merge when we see this Switch the second time.
+      merge_grad = merge([grad[0], grad[0]], name="b_switch")[0]
+      grad_ctxt.grad_state.switch_map[op] = merge_grad
+      return merge_grad, None
+    else:
+      # This is the first time this Switch is visited. It comes from the
+      # Identity branch. Such a Switch has `None` gradient for the Exit branch,
+      # meaning the output is not differentiable.
+      return None, None
+  elif isinstance(op_ctxt, CondContext):
+    zero_grad = grad[1 - op_ctxt.branch]
+    # At this point, we have created zero_grad guarded by the right switch.
+    # Unfortunately, we may still get None here for not trainable data types.
+    if zero_grad is None:
+      # For resource variables we get None always on the other branch, so bypass
+      # this.
+      if op.inputs[0].dtype == dtypes.resource:
+        return merge(
+            [grad[op_ctxt.branch]] * 2, name="cond_resource_grad")[0], None
+      return None, None
+    return merge(grad, name="cond_grad")[0], None
+  else:
+    false_grad = switch(grad[0], op.inputs[1])[0]
+    true_grad = switch(grad[1], op.inputs[1])[1]
+    return merge([false_grad, true_grad])[0], None
+
+
+ops.RegisterGradient("Switch")(_SwitchGrad)
+ops.RegisterGradient("RefSwitch")(_SwitchGrad)
+
+
+@ops.RegisterGradient("Merge")
+def _MergeGrad(op, grad, _):
+  """Gradients for a Merge op are calculated using a Switch op."""
+  input_op = op.inputs[0].op
+  graph = ops.get_default_graph()
+  # pylint: disable=protected-access
+  op_ctxt = control_flow_util.GetOutputContext(input_op)
+  grad_ctxt = graph._get_control_flow_context()
+  # pylint: enable=protected-access
+  if isinstance(op_ctxt, WhileContext):
+    # pylint: disable=protected-access
+    return control_flow_ops._SwitchRefOrTensor(grad, grad_ctxt.pivot)
+    # pylint: enable=protected-access
+  elif isinstance(op_ctxt, CondContext):
+    pred = op_ctxt.pred
+    if grad_ctxt and grad_ctxt.grad_state:
+      # This Merge node is part of a cond within a loop.
+      # The backprop needs to have the value of this predicate for every
+      # iteration. So we must have its values accumulated in the forward, and
+      # use the accumulated values as the predicate for this backprop switch.
+      grad_state = grad_ctxt.grad_state
+      real_pred = grad_state.history_map.get(pred.name)
+      if real_pred is None:
+        # Remember the value of pred for every iteration.
+        grad_ctxt = grad_state.grad_context
+        grad_ctxt.Exit()
+        history_pred = grad_state.AddForwardAccumulator(pred)
+        grad_ctxt.Enter()
+
+        # Add the stack pop op. If pred.op is in a (outer) CondContext,
+        # the stack pop will be guarded with a switch.
+        real_pred = grad_state.AddBackpropAccumulatedValue(history_pred, pred)
+        grad_state.history_map[pred.name] = real_pred
+      pred = real_pred
+    # pylint: disable=protected-access
+    return control_flow_ops._SwitchRefOrTensor(grad, pred, name="cond_grad")
+    # pylint: enable=protected-access
+  else:
+    num_inputs = len(op.inputs)
+    cond = [math_ops.equal(op.outputs[1], i) for i in range(num_inputs)]
+    # pylint: disable=protected-access
+    return [
+        control_flow_ops._SwitchRefOrTensor(grad, cond[i])[1]
+        for i in range(num_inputs)
+    ]
+    # pylint: enable=protected-access
+
+
+@ops.RegisterGradient("RefMerge")
+def _RefMergeGrad(op, grad, _):
+  return _MergeGrad(op, grad, _)
+
+
+@ops.RegisterGradient("Exit")
+def _ExitGrad(op, grad):
+  """Gradients for an exit op are calculated using an Enter op."""
+  graph = ops.get_default_graph()
+  # pylint: disable=protected-access
+  op_ctxt = op._get_control_flow_context()
+  grad_ctxt = graph._get_control_flow_context()
+  # pylint: enable=protected-access
+  if not grad_ctxt.back_prop:
+    # The flag `back_prop` is set by users to suppress gradient
+    # computation for this loop. If the attribute `back_prop` is false,
+    # no gradient computation.
+    return None
+
+  if op_ctxt.grad_state:
+    raise TypeError("Second-order gradient for while loops not supported.")
+
+  if isinstance(grad, tensor.Tensor):
+    grad_ctxt.AddName(grad.name)
+  else:
+    if not isinstance(
+        grad, (indexed_slices.IndexedSlices, sparse_tensor.SparseTensor)):
+      raise TypeError(f"Type {type(grad)} not supported, must be either"
+                      "`indexed_slices.IndexedSlices` or `SparseTensor`.")
+    grad_ctxt.AddName(grad.values.name)
+    grad_ctxt.AddName(grad.indices.name)
+    dense_shape = grad.dense_shape
+    if dense_shape is not None:
+      grad_ctxt.AddName(dense_shape.name)
+  grad_ctxt.Enter()
+  # pylint: disable=protected-access
+  result = control_flow_ops._Enter(
+      grad, grad_ctxt.name, is_constant=False,
+      parallel_iterations=grad_ctxt.parallel_iterations,
+      name="b_exit")
+  # pylint: enable=protected-access
+  grad_ctxt.loop_enters.append(result)
+  grad_ctxt.Exit()
+  return result
+
+
+ops.RegisterGradient("RefExit")(_ExitGrad)
+
+
+@ops.RegisterGradient("NextIteration")
+def _NextIterationGrad(_, grad):
+  """A forward next_iteration is translated into a backprop identity.
+
+  Note that the backprop next_iteration is added in switch grad.
+  """
+  return grad
+
+
+@ops.RegisterGradient("RefNextIteration")
+def _RefNextIterationGrad(_, grad):
+  return _NextIterationGrad(_, grad)
+
+
+@ops.RegisterGradient("Enter")
+def _EnterGrad(op, grad):
+  """Gradients for an Enter are calculated using an Exit op.
+
+  For loop variables, grad is the gradient so just add an exit.
+  For loop invariants, we need to add an accumulator loop.
+  """
+  graph = ops.get_default_graph()
+  # pylint: disable=protected-access
+  grad_ctxt = graph._get_control_flow_context()
+  # pylint: enable=protected-access
+  if grad_ctxt is None:
+    return grad
+  if not grad_ctxt.back_prop:
+    # Skip gradient computation, if the attribute `back_prop` is false.
+    return grad
+  if grad_ctxt.grad_state is None:
+    # Pass the gradient through if we are not in a gradient while context.
+    return grad
+  if op.get_attr("is_constant"):
+    # Add a gradient accumulator for each loop invariant.
+    if isinstance(grad, tensor.Tensor):
+      result = grad_ctxt.AddBackpropAccumulator(op, grad)
+    elif isinstance(grad, indexed_slices.IndexedSlices):
+      result = grad_ctxt.AddBackpropIndexedSlicesAccumulator(op, grad)
+    else:
+      # TODO(yuanbyu, lukasr): Add support for SparseTensor.
+      raise TypeError(f"Type {type(grad)} not supported,"
+                      "must be Tensor or Indexed Slices")
+  else:
+    result = exit(grad)
+    grad_ctxt.loop_exits.append(result)
+    grad_ctxt.ExitResult([result])
+  return result
+
+
+@ops.RegisterGradient("RefEnter")
+def _RefEnterGrad(op, grad):
+  return _EnterGrad(op, grad)
+
+
+@ops.RegisterGradient("LoopCond")
+def _LoopCondGrad(_):
+  """Stop backprop for the predicate of a while loop."""
+  return None

tensorflow/python/ops/control_flow_ops_benchmark.py

@@ -0,0 +1,119 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Benchmark for control flow ops."""
+
+import time
+
+from tensorflow.python.client import session
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import control_flow_util
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.platform import test
+
+
+class CondWithManyIntermediatesBenchmark(test.Benchmark):
+  """Checks the runtime performance of outputting all intermediates."""
+
+  NUM_INTERMEDIATES = 1000
+  NUM_ITERS = 500
+  NUM_WARM_UP_ITERS = 50
+
+  def _create_cond(self, x):
+
+    def branch_fn():
+      # Use a random value so the adds can't be constant folded.
+      return x + sum(random_ops.random_normal([])
+                     for _ in range(self.NUM_INTERMEDIATES))
+
+    # Use a dynamic predicate to make sure the cond isn't constant folded.
+    return cond.cond(math_ops.not_equal(x, -1),
+                     branch_fn, lambda: 0.0)
+
+  def _benchmark_defun(self):
+    """Benchmarks cond in a defun."""
+
+    @def_function.function
+    def cond_fn(x):
+      return self._create_cond(x)
+
+    # Warm up
+    for _ in range(self.NUM_WARM_UP_ITERS):
+      cond_fn(0.0)
+
+    start_time = time.time()
+
+    for _ in range(self.NUM_ITERS):
+      cond_fn(0.0)
+
+    self.report_benchmark(
+        wall_time=time.time() - start_time,
+        iters=self.NUM_ITERS)
+
+  def _benchmark_graph(self):
+    """Benchmarks cond in legacy graph mode."""
+    with context.graph_mode():
+      with ops.Graph().as_default():
+        x = array_ops.placeholder(dtypes.float32)
+        cond_val = self._create_cond(x)
+
+        with session.Session() as sess:
+          cond_fn = sess.make_callable(cond_val, [x])
+
+          # Warm up
+          for _ in range(self.NUM_WARM_UP_ITERS):
+            cond_fn(0.0)
+
+          start_time = time.time()
+
+          for _ in range(self.NUM_ITERS):
+            cond_fn(0.0)
+
+          self.report_benchmark(
+              wall_time=time.time() - start_time,
+              iters=self.NUM_ITERS)
+
+  def benchmark_cond_v1_defun(self):
+    old_val = control_flow_util.ENABLE_CONTROL_FLOW_V2
+    control_flow_util.ENABLE_CONTROL_FLOW_V2 = False
+    self._benchmark_defun()
+    control_flow_util.ENABLE_CONTROL_FLOW_V2 = old_val
+
+  def benchmark_cond_v2_defun(self):
+    old_val = control_flow_util.ENABLE_CONTROL_FLOW_V2
+    control_flow_util.ENABLE_CONTROL_FLOW_V2 = True
+    self._benchmark_defun()
+    control_flow_util.ENABLE_CONTROL_FLOW_V2 = old_val
+
+  def benchmark_cond_v1_graph(self):
+    old_val = control_flow_util.ENABLE_CONTROL_FLOW_V2
+    control_flow_util.ENABLE_CONTROL_FLOW_V2 = False
+    self._benchmark_graph()
+    control_flow_util.ENABLE_CONTROL_FLOW_V2 = old_val
+
+  def benchmark_cond_v2_graph(self):
+    old_val = control_flow_util.ENABLE_CONTROL_FLOW_V2
+    control_flow_util.ENABLE_CONTROL_FLOW_V2 = True
+    self._benchmark_graph()
+    control_flow_util.ENABLE_CONTROL_FLOW_V2 = old_val
+
+if __name__ == "__main__":
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/ops/control_flow_state.py

@@ -0,0 +1,841 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utilities for managing state of v1 control flow for computing gradients."""
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import control_flow_util as util
+from tensorflow.python.ops import control_flow_v2_func_graphs
+from tensorflow.python.ops import default_gradient
+from tensorflow.python.ops import gen_data_flow_ops
+from tensorflow.python.ops import gen_resource_variable_ops
+from tensorflow.python.ops import resource_variable_ops
+
+# pylint: disable=protected-access
+
+
+def _GetMaxSizeFromNestedMaximumIterations(value, while_ctxt):
+  """Calculate a max_size for use by stack ops inside an XLA while_loop.
+
+  Args:
+    value: The value inside the while_loop forward context.  Used for printing
+      error messages.
+    while_ctxt: The forward context inside which value resides.  This does not
+      always match the value's immediate context, as `value` may be inside e.g.
+      a cond context inside the while_loop.
+
+  Returns:
+    A tensor containing the `max_size` to feed to a Stack initializer.
+
+  Raises:
+    ValueError: If `value` is nested inside a `while_loop` that either
+      lacks a `maximum_iterations` parameter, or the `maximum_iterations`
+      parameter:
+
+        - is inside a `while_loop` that is a parent of the calling context, and
+        - cannot be evaluated at graph build time to a constant.
+  """
+  value_name = value.name
+  # curr_ctxt is the context that tf.gradients was called in.
+  curr_ctxt = ops.get_default_graph()._get_control_flow_context()  # pylint: disable=protected-access
+
+  curr_ctxt_name = curr_ctxt.name if curr_ctxt is not None else ""
+  max_size = constant_op.constant(1)
+
+  # Loop through all containing while contexts between value and the
+  # current context, multiplying together each context's
+  # max_iterations to get the maximum stack size.
+  while while_ctxt not in (None, curr_ctxt):
+    max_iter = while_ctxt.maximum_iterations
+    if max_iter is None:
+      raise ValueError(
+          "Cannot create a gradient accumulator for tensor '%s' inside "
+          "XLA while_loop because maximum_iterations was not passed to "
+          "the tf.while_loop call ('%s')." % (value_name, while_ctxt.name))
+
+    # pylint: disable=protected-access
+    max_iter_ctxt = max_iter.op._get_control_flow_context()
+    # pylint: enable=protected-access
+
+    # If max_iter_ctxt (non-strictly) contains curr_ctxt, then it's OK to use.
+    if util.IsContainingContext(curr_ctxt, max_iter_ctxt):
+      max_size *= max_iter
+    else:
+      # We cannot use max_iter because it's defined in a nested while
+      # or cond context, so will fail if we try to use it as input to
+      # any ops in curr_ctxt (e.g. max_size or the final accumulator
+      # stack). Attempt to get a constant value out to use instead.
+      const_max_iter = tensor_util.constant_value(max_iter)
+      if const_max_iter is None:
+        raise ValueError(
+            "Cannot create a gradient accumulator for tensor '%s' inside XLA "
+            "while_loop. maximum_iterations tensor '%s' for while_loop context "
+            "'%s' must be statically known (e.g. a constant value or known "
+            "shape dimension), or be defined at or outside the while loop "
+            "context '%s' (currently defined in '%s')." %
+            (value_name, max_iter.name, while_ctxt.name, curr_ctxt_name,
+             max_iter_ctxt.name))
+      max_size *= const_max_iter
+
+    # Find the next outer WhileContext (or stop if we reach the
+    # tf.gradient's context).
+    while_ctxt = util.GetContainingWhileContext(
+        while_ctxt.outer_context, stop_ctxt=curr_ctxt)
+
+  return max_size
+
+
+class _GradLoopState:
+  """The state used for constructing the gradient graph for a while loop.
+
+  We create a _GradLoopState for each while loop in forward and its
+  corresponding while loop in backprop. This gives us access to both
+  the forward and the backprop WhileContexts.
+
+  During the construction of gradient graph, any time when we detect
+  a forward value that is needed for backprop, we create a history
+  accumulator and add it to `history_map`. Any time when we backprop
+  a loop switch op (in _SwitchGrad), we add the grad merge op in
+  `switch_map`.
+  """
+
+  def __init__(self, forward_ctxt, outer_grad_state):
+    # The grad loop state for the outer while loop.
+    self._outer_grad_state = None
+
+    # The while loop context for forward.
+    self._forward_context = None
+
+    # The loop counter added by AddForwardLoopCounter. It is the value
+    # of the loop counter for the next iteration.
+    self._forward_index = None
+
+    # A sync op for forward.
+    self._forward_sync = None
+
+    # The while loop context for backprop.
+    self._grad_context = None
+
+    # The loop counter added by AddBackpropLoopCounter. It is the value
+    # of the loop counter for the current iteration.
+    self._grad_index = None
+
+    # A sync op for backprop.
+    self._grad_sync = None
+
+    # Information needed by backprop.
+    self._history_map = {}
+    self._switch_map = {}
+    self._unused_exits = []
+    self._deferred_exits = []
+    self._forward_loop_exits = list(forward_ctxt.loop_exits)
+    self._pending_exits_count = len(forward_ctxt.loop_exits)
+
+    self._outer_grad_state = outer_grad_state
+    if outer_grad_state:
+      outer_forward_ctxt = outer_grad_state.forward_context
+    else:
+      if not hasattr(forward_ctxt, "outer_context"):
+        raise ValueError("Failed to call gradients on a while loop without"
+                         "properly serializing graph via MetaGraphDef")
+      outer_forward_ctxt = forward_ctxt.outer_context
+
+    # Add the forward loop counter.
+    with forward_ctxt._graph.as_default():  # pylint: disable=protected-access
+      if outer_forward_ctxt:
+        outer_forward_ctxt.Enter()
+      cnt, forward_index = forward_ctxt.AddForwardLoopCounter(outer_grad_state)
+      if outer_forward_ctxt:
+        outer_forward_ctxt.Exit()
+    self._forward_context = forward_ctxt
+    self._forward_index = forward_index
+
+    # Add the backprop WhileContext, and the backprop loop counter.
+    if outer_grad_state:
+      # This is a nested loop. Remember the iteration counts for each
+      # execution of this inner loop.
+      outer_forward_ctxt.AddName(cnt.name)
+      history_cnt = outer_grad_state.AddForwardAccumulator(cnt)
+
+      outer_grad_ctxt = outer_grad_state.grad_context
+      outer_grad_ctxt.Enter()
+      self._grad_context = control_flow_ops.WhileContext(
+          maximum_iterations=forward_ctxt.maximum_iterations,
+          parallel_iterations=forward_ctxt.parallel_iterations,
+          back_prop=forward_ctxt.back_prop,
+          swap_memory=forward_ctxt.swap_memory,
+          name=forward_ctxt.name,
+          grad_state=self)
+      real_cnt = outer_grad_state.AddBackpropAccumulatedValue(history_cnt, cnt)
+      self._grad_index = self._grad_context.AddBackpropLoopCounter(
+          real_cnt, outer_grad_state)
+      outer_grad_ctxt.Exit()
+    else:
+      if outer_forward_ctxt:
+        outer_forward_ctxt.Enter()
+      self._grad_context = control_flow_ops.WhileContext(
+          maximum_iterations=forward_ctxt.maximum_iterations,
+          parallel_iterations=forward_ctxt.parallel_iterations,
+          back_prop=forward_ctxt.back_prop,
+          swap_memory=forward_ctxt.swap_memory,
+          name=forward_ctxt.name,
+          grad_state=self)
+      self._grad_index = self._grad_context.AddBackpropLoopCounter(
+          cnt, outer_grad_state)
+      if outer_forward_ctxt:
+        outer_forward_ctxt.Exit()
+
+  @property
+  def outer_grad_state(self):
+    """The grad loop state for outer loop."""
+    return self._outer_grad_state
+
+  @property
+  def forward_context(self):
+    """The while loop context for forward."""
+    return self._forward_context
+
+  @property
+  def forward_index(self):
+    """The loop index of forward loop."""
+    return self._forward_index
+
+  @property
+  def forward_sync(self):
+    """A control trigger node for synchronization in the forward loop.
+
+    One main use is to keep the push ops of a stack executed in the
+    iteration order.
+    """
+    if self._forward_sync is None:
+      with ops.control_dependencies(None):
+        self._forward_sync = control_flow_ops.control_trigger(name="f_sync")
+      self._forward_sync._set_control_flow_context(self._forward_context)
+      self._forward_index.op._add_control_input(self._forward_sync)
+    return self._forward_sync
+
+  @property
+  def grad_context(self):
+    """The corresponding WhileContext for gradient."""
+    return self._grad_context
+
+  @property
+  def grad_index(self):
+    """The loop index of backprop loop."""
+    return self._grad_index
+
+  @property
+  def grad_sync(self):
+    """A control trigger node for synchronization in the grad loop.
+
+    One main use is to keep the pop ops of a stack executed in the
+    iteration order.
+    """
+    if self._grad_sync is None:
+      with ops.control_dependencies(None):
+        self._grad_sync = control_flow_ops.control_trigger(name="b_sync")
+      self._grad_sync._set_control_flow_context(self._grad_context)
+      self._grad_index.op._add_control_input(self._grad_sync)
+      if self._grad_context.outer_context:
+        self._grad_context.outer_context.AddInnerOp(self._grad_sync)
+    return self._grad_sync
+
+  @property
+  def history_map(self):
+    """The map that records all the tensors needed for backprop."""
+    return self._history_map
+
+  @property
+  def switch_map(self):
+    """The map that records all the Switch ops for the while loop."""
+    return self._switch_map
+
+  @property
+  def unused_exits(self):
+    """The list of "unused" exits."""
+    return self._unused_exits
+
+  @property
+  def deferred_exits(self):
+    """The list of "deferred" exits."""
+    return self._deferred_exits
+
+  @property
+  def forward_loop_exits(self):
+    """The list of exits of the forward loop."""
+    return self._forward_loop_exits
+
+  @property
+  def pending_exits_count(self):
+    """The number of exits we expect to see but haven't."""
+    return self._pending_exits_count
+
+  @pending_exits_count.setter
+  def pending_exits_count(self, cnt):
+    """Set the pending count to cnt."""
+    self._pending_exits_count = cnt
+
+  def AddForwardAccumulator(self, value, dead_branch=False):
+    """Add an accumulator for each forward tensor that is needed in backprop.
+
+    This is added to the forward loop at the first time when a tensor
+    in the forward loop is used by backprop gradient computation loop.
+    We create an accumulator that accumulates the value of tensor at each
+    iteration. Called in the control flow context where gradients() is called.
+
+    The pseudocode is:
+    ```
+      acc = stack();
+      while (_pivot) {
+        acc = stack_push(acc, value);
+      }
+    ```
+
+    We make sure that the stack push op in one iteration is executed before
+    next iteration. This is achieved by adding a control edge from
+    `forward_index.op.inputs[0].op` to the push op, and another control
+    edge from the push op to either `forward_index.op` or `forward_sync`.
+
+    Args:
+      value: The source tensor in forward that is to be accumulated.
+      dead_branch: True iff the tensor is on a dead branch of a cond.
+
+    Returns:
+      The stack that contains the accumulated history of the tensor.
+
+    Raises:
+      TypeError: For internal errors involving the value condition context.
+      ValueError: If `value` is inside a XLA scope and a valid max size
+        for the stack can't be found.
+    """
+    # curr_ctxt is the context that tf.gradients was called in.
+    with self._forward_index.graph.as_default():
+      curr_ctxt = ops.get_default_graph()._get_control_flow_context()  # pylint: disable=protected-access
+      with ops.control_dependencies(None):
+        if curr_ctxt:
+          curr_ctxt.Enter()
+        with ops.colocate_with(value):
+          # We only need to pass maximum_iterations to the stack if
+          # we're inside an XLA context.
+          if not util.IsInXLAContext(value.op):
+            max_size = constant_op.constant(-1, dtypes.int32)
+          else:
+            max_size = _GetMaxSizeFromNestedMaximumIterations(
+                value, self.forward_context)
+          acc = gen_data_flow_ops.stack_v2(
+              max_size=max_size, elem_type=value.dtype.base_dtype, name="f_acc")
+        if curr_ctxt:
+          curr_ctxt.Exit()
+
+        # Make acc available in the forward context.
+        enter_acc = self.forward_context.AddValue(acc)
+
+        # Add the stack_push op in the context of value.op.
+        swap_enabled = self.forward_context.swap_memory
+        value_ctxt = util.GetOutputContext(value.op)
+        if value_ctxt == self.forward_context:
+          # value is not nested in the forward context.
+          self.forward_context.Enter()
+          push = gen_data_flow_ops.stack_push_v2(
+              enter_acc, value, swap_memory=swap_enabled)
+          self.forward_context.Exit()
+          # Protect stack push and order it before forward_index.
+          self.forward_index.op._add_control_input(push.op)
+        else:
+          # value is in a cond context within the forward context.
+          if not isinstance(value_ctxt, control_flow_ops.CondContext):
+            raise TypeError("value_ctxt is not a CondContext: %s" % value_ctxt)
+          if dead_branch:
+            # The special case for creating a zero tensor for a dead
+            # branch of a switch. See _ControlFlowState.ZerosLikeV1WhileLoop().
+            value_ctxt.outer_context.Enter()
+            push = gen_data_flow_ops.stack_push_v2(
+                enter_acc, value, swap_memory=swap_enabled)
+            value_ctxt.outer_context.Exit()
+            push.op._set_control_flow_context(value_ctxt)
+          else:
+            value_ctxt.Enter()
+            push = gen_data_flow_ops.stack_push_v2(
+                enter_acc, value, swap_memory=swap_enabled)
+            value_ctxt.Exit()
+          # Protect stack push and order it before forward_sync.
+          self.forward_sync._add_control_input(push.op)
+        # Order stack push after the successor of forward_index
+        add_op = self.forward_index.op.inputs[0].op
+        push.op._add_control_input(add_op)
+        return acc
+
+  def AddBackpropAccumulatedValue(self, history_value, value,
+                                  dead_branch=False):
+    """Add the getter for an accumulated value in the grad context.
+
+    This is added to the backprop loop. Called in the grad context to
+    get the value of an accumulated value. The stack pop op must be guarded
+    by the pred of the controlling cond.
+
+    Args:
+      history_value: The history (a stack) of a value.
+      value: The value that is pushed onto the stack.
+      dead_branch: True iff the tensor is on a dead branch of a cond.
+
+    Returns:
+      The current value (the top of the stack).
+    """
+    history_ctxt = history_value.op._get_control_flow_context()
+    # Find the cond context that controls history_value if any.
+    cond_ctxt = None
+    value_ctxt = value.op._get_control_flow_context()
+    while value_ctxt and value_ctxt != history_ctxt:
+      if isinstance(value_ctxt, control_flow_ops.CondContext):
+        cond_ctxt = value_ctxt
+        break
+      value_ctxt = value_ctxt.outer_context
+    with ops.control_dependencies(None):
+      self.grad_context.Enter()
+      if cond_ctxt:
+        # Guard stack pop with a switch if it is controlled by a cond.
+        grad_state = self
+        pred = None
+        while pred is None and grad_state:
+          pred = grad_state.history_map.get(cond_ctxt.pred.name)
+          grad_state = grad_state.outer_grad_state
+        if pred is None:
+          pred = cond_ctxt.pred
+        branch = (1 - cond_ctxt.branch) if dead_branch else cond_ctxt.branch
+        history_value = control_flow_ops._SwitchRefOrTensor(
+            history_value, pred)[branch]
+      pop = gen_data_flow_ops.stack_pop_v2(history_value,
+                                           value.dtype.base_dtype)
+      pop.set_shape(value.get_shape())
+      self.grad_context.Exit()
+    parallel_iterations = self.grad_context.parallel_iterations
+    if parallel_iterations > 1:
+      # All pops are ordered after pivot_for_body and before grad_sync.
+      self.grad_sync._add_control_input(pop.op)
+    return pop
+
+  def GetRealValue(self, value):
+    """Get the real value of `value`.
+
+    If backprop "uses" a value produced by forward inference, an accumulator
+    is added in the forward loop to accumulate its values.  We use the
+    accumulated value. This method must be called in the grad loop context.
+    `value` must be in forward and needed for backprop.
+
+    Args:
+      value: A tensor to be captured.
+
+    Returns:
+      The same tensor obtained from the saved history.
+    """
+    assert value.op.type not in ["Variable", "VariableV2"]
+    real_value = self._history_map.get(value.name)
+    if real_value is None:
+      cur_value = value
+      cur_grad_state = self
+      while True:
+        enter_op = util.GetLoopConstantEnter(cur_value)
+        if enter_op:
+          # Special case: cur_value comes from a constant Enter node.
+          cur_value = enter_op.inputs[0]
+          cur_grad_state = cur_grad_state.outer_grad_state
+          if cur_grad_state is None:
+            # We are now outside all nested loops for this gradient(),
+            # so `value` is a loop invariant and there is no need to
+            # save the history of value. Just make cur_value to enter
+            # the right control flow context.
+            real_value = self._grad_context.AddValue(cur_value)
+            break
+        elif constant_op.is_constant(cur_value):
+          # If the value to be forwarded is a constant, clone the constant in
+          # the gradient loop rather than using a stack.
+          # TODO(phawkins): consider hoisting the constant out of the loop
+          # instead.
+          real_value = constant_op.constant(
+              tensor_util.constant_value(cur_value), dtype=cur_value.dtype)
+          break
+        else:
+          # Record the history of this value in forward_ctxt.
+          self._grad_context.Exit()
+          history_value = cur_grad_state.AddForwardAccumulator(cur_value)
+          self._grad_context.Enter()
+          break
+
+      if real_value is None:
+        # Add the stack pop op in the grad context.
+        real_value = cur_grad_state.AddBackpropAccumulatedValue(
+            history_value, cur_value)
+        if cur_grad_state != self:
+          real_value = self._grad_context.AddValue(real_value)
+      self._history_map[value.name] = real_value
+    return real_value
+
+
+class _ControlFlowState:
+  """Maintain the mapping from the loops to their grad states."""
+
+  def __init__(self):
+    self._map = {}  # maps forward loop context to _GradLoopState
+
+  def GetGradState(self, op: ops.Operation, before):
+    """Return the grad state for this op if it's in a forward loop context."""
+    if before and util.IsLoopExit(op):
+      forward_ctxt = op._get_control_flow_context()  # pylint: disable=protected-access
+      forward_ctxt = forward_ctxt.outer_context
+      if forward_ctxt:
+        forward_ctxt = forward_ctxt.GetWhileContext()
+    else:
+      forward_ctxt = util.GetWhileContext(op)
+    if forward_ctxt:
+      return self._map.get(forward_ctxt)
+    return None
+
+  def ProcessUnusedLoopExits(self, pending_count, to_ops_set):
+    """Process all the "unused" loop exits.
+
+    The "unused" exits of the loops are added to `unused_exits`. An exit is
+    unused if its pending_count is 0. If there is an exit with real gradient,
+    all these deferred exits will enter the backprop loop with zero gradient.
+    Otherwise, they will enter the backprop loop with None. As an example,
+    people often write:
+
+    ```python
+    v1, _ = tf.while_loop(p, b, [x1, x2])
+    result = gradients(v1, x1)
+    ```
+
+    The exit node for x2 is not included by the betweenness analysis. But we
+    need to backprop x2 if x2 is involved in computing v1.
+
+    Args:
+      pending_count: The number of backprop inputs for every op.
+      to_ops_set: The set of ops for ys in gradients(ys, xs)
+
+    Returns:
+      The set of unused loop exits that we know at this point we need
+      to backprop.
+    """
+    loop_exits = []
+    for grad_state in self._map.values():
+      for y in grad_state.forward_loop_exits:
+        if pending_count[y.op] == 0:
+          grad_state.pending_exits_count -= 1
+          if y.op not in to_ops_set:
+            grad_state.unused_exits.append(y)
+          if grad_state.pending_exits_count == 0:
+            loop_exits.extend(grad_state.unused_exits)
+      # Need to include Enters in backprop for higher-order gradients.
+      for y in grad_state.forward_context.loop_enters:
+        if pending_count[y.op] == 0:
+          pending_count[y.op] = 1
+    return loop_exits
+
+  def EnterGradWhileContext(self, op, before):
+    """Enter the WhileContext for gradient computation."""
+    grad_state = self.GetGradState(op, before)
+    if grad_state:
+      grad_state.grad_context.Enter()
+
+  def ExitGradWhileContext(self, op, before):
+    """Exit the WhileContext for gradient computation."""
+    grad_state = self.GetGradState(op, before)
+    if grad_state:
+      grad_state.grad_context.Exit()
+
+  def AddWhileContext(self, op, between_op_list, between_ops):
+    """Add the grad state for the while loop that op belongs to.
+
+    Note that op is an Exit, and this method must be called in
+    the control flow context where gradients() is called.
+
+    Note that this method modifies `between_op_list` and `between_ops`.
+    """
+    forward_ctxt = util.GetWhileContext(op)
+    grad_state = self._map.get(forward_ctxt)
+    if grad_state is None:
+      # This is a new while loop so create a grad state for it.
+      outer_forward_ctxt = forward_ctxt.outer_context
+      if outer_forward_ctxt:
+        outer_forward_ctxt = outer_forward_ctxt.GetWhileContext()
+      outer_grad_state = None
+      if outer_forward_ctxt:
+        outer_grad_state = self._map.get(outer_forward_ctxt)
+      grad_state = _GradLoopState(forward_ctxt, outer_grad_state)
+      self._map[forward_ctxt] = grad_state
+
+      # We need to include all exits of a loop for backprop.
+      for loop_exit in grad_state.forward_loop_exits:
+        if loop_exit.op not in between_ops:
+          between_ops.add(loop_exit.op)
+          between_op_list.append(loop_exit.op)
+
+  def ZerosLikeForExit(self, val):
+    """Create zeros_like gradient for a loop exit.
+
+    If the result of a loop variable is not used but is involved in
+    computing the result of some needed loop variable, we create a
+    zero-valued tensor that is fed as gradient for the Exit node of that
+    loop variable. Note that val.op is an Exit, and this method must be
+    called in the control flow context where gradients() is called.
+
+    Args:
+      val: The output tensor of an Exit op.
+
+    Returns:
+      A zero tensor of the same shape of val.
+    """
+    val_shape = val.get_shape()
+    forward_ctxt = val.op._get_control_flow_context()
+    outer_forward_ctxt = forward_ctxt.outer_context
+    if outer_forward_ctxt:
+      outer_forward_ctxt = outer_forward_ctxt.GetWhileContext()
+    outer_grad_state = None
+    if outer_forward_ctxt:
+      outer_grad_state = self._map.get(outer_forward_ctxt)
+    if outer_grad_state:
+      # This is a nested loop.
+      if val_shape.is_fully_defined():
+        # If the shape is known statically, just create a zero tensor
+        # with the right shape in the right context.
+        outer_grad_state.grad_context.Enter()
+        result = array_ops.zeros(val_shape.dims, val.dtype)
+        outer_grad_state.grad_context.Exit()
+      else:
+        # Only the shape of value is needed for backprop.
+        forward_ctxt.outer_context.Enter()
+        shape = array_ops.shape_internal(val, optimize=False)
+        forward_ctxt.outer_context.Exit()
+        # Save the shape to a stack.
+        history_shape = outer_grad_state.AddForwardAccumulator(shape)
+        # Get the shape back from the stack.
+        outer_grad_ctxt = outer_grad_state.grad_context
+        outer_grad_ctxt.Enter()
+        real_shape = outer_grad_state.AddBackpropAccumulatedValue(
+            history_shape, shape)
+        result = array_ops.zeros(real_shape, val.dtype)
+        outer_grad_ctxt.Exit()
+    else:
+      # This is not a nested loop.
+      if val_shape.is_fully_defined():
+        # If the shape is known statically, just create a zero tensor
+        # with the right shape.
+        result = array_ops.zeros(val_shape.dims, val.dtype)
+      else:
+        result = array_ops.zeros_like(val, optimize=False)
+    return result
+
+  def ZerosLikeV1WhileLoop(self, op, index):
+    """Create zeros_like for the specified output of an op.
+
+    If op is in a while loop that is part of gradients(), this method
+    must be called in its grad loop context.
+
+    Args:
+      op: A tensorflow operation.
+      index: the index for a specific output of the op.
+
+    Returns:
+      A zero tensor of the same shape of op.outputs[index].
+    """
+    if util.IsLoopSwitch(op):
+      return None
+    if op.graph.building_function:
+      # The optimization here is tricky to apply to functions
+      return array_ops.zeros_like(op.outputs[index])
+    dead_branch = util.IsSwitch(op)
+    forward_ctxt = util.GetWhileContext(op)
+    grad_state = self._map.get(forward_ctxt)
+    if grad_state is None:
+      # op is not in a while loop that is part of gradients().
+      return ZerosLike(op, index)
+    op_ctxt = op._get_control_flow_context()
+    val = ops.convert_to_tensor(op.outputs[index], name="tensor")
+    shape = val.get_shape()
+    if shape.is_fully_defined():
+      # If the shape is known statically, just create a zero tensor with
+      # the right shape in the grad loop context.
+      if val.dtype == dtypes.resource:
+        result = array_ops.zeros(
+            resource_variable_ops.variable_shape(val),
+            dtype=default_gradient.get_zeros_dtype(val))
+      else:
+        result = constant_op.constant(0, shape=shape.dims, dtype=val.dtype)
+      if dead_branch:
+        # op is a cond switch. Guard the zero tensor with a switch.
+        pred = grad_state.history_map.get(op_ctxt.pred.name)
+        branch = op_ctxt.branch
+        result = control_flow_ops._SwitchRefOrTensor(result, pred)[1 - branch]
+    else:
+      # Unknown shape so keep a history of the shape at runtime.
+      if dead_branch:
+        # Need to add a special switch to guard the value.
+        pred = op_ctxt.pred
+        branch = op_ctxt.branch
+        op_ctxt.outer_context.Enter()
+        val = control_flow_ops._SwitchRefOrTensor(op.inputs[0],
+                                                  pred)[1 - branch]
+        zeros_shape = array_ops.shape_internal(val, optimize=False)
+        op_ctxt.outer_context.Exit()
+        val.op._set_control_flow_context(op_ctxt)
+        zeros_shape.op._set_control_flow_context(op_ctxt)
+      else:
+        op_ctxt.Enter()
+        zeros_shape = array_ops.shape_internal(val, optimize=False)
+        op_ctxt.Exit()
+
+      # Add forward accumulator for shape.
+      grad_state.grad_context.Exit()
+      history_zeros_shape = grad_state.AddForwardAccumulator(
+          zeros_shape, dead_branch=dead_branch)
+      grad_state.grad_context.Enter()
+
+      # Create a zero tensor with the right shape.
+      shape = grad_state.AddBackpropAccumulatedValue(history_zeros_shape,
+                                                     zeros_shape, dead_branch)
+      result = array_ops.zeros(shape, val.dtype)
+    return result
+
+  def PostProcessing(self):
+    """Perform postprocessing at the end of gradients().
+
+    We have created the gradient graph at this point. So this function
+    can be used to perform any postprocessing on the gradient graph.
+    We currently perform the following postprocessing:
+      1. Patch the gradient graph if the output of a loop variable
+         doesn't depend on its input.
+    """
+    for _, grad_state in self._map.items():
+      for _, b_merge in grad_state.switch_map.items():
+        if b_merge.op.inputs[0] == b_merge.op.inputs[1]:
+          # The value of this loop variable at iteration i+1 doesn't
+          # depend on its value at iteration i. So use zeros as the
+          # gradients for all iterations > 0.
+          dtype = b_merge.op.inputs[0].dtype
+          shape = b_merge.op.inputs[0].get_shape()
+          # pylint: disable=protected-access
+          if shape.is_fully_defined():
+            grad_state.grad_context.Enter()
+            # Create a zeros and use it for iterations > 0.
+            grad_val = constant_op.constant(0, dtype=dtype, shape=shape)
+            next_grad_val = control_flow_ops._NextIteration(grad_val)
+            grad_state.grad_context.Exit()
+          else:
+            # Create a zeros in the outer grad context.
+            outer_grad_ctxt = grad_state.grad_context.outer_context
+            if outer_grad_ctxt:
+              outer_grad_ctxt.Enter()
+            enter_grad_op = b_merge.op.inputs[0].op
+            enter_grad = enter_grad_op.inputs[0]
+            grad_shape = array_ops.shape_internal(enter_grad, optimize=False)
+            grad_val = array_ops.zeros(grad_shape)
+            if outer_grad_ctxt:
+              outer_grad_ctxt.Exit()
+            # Use the zeros for iterations > 0.
+            grad_state.grad_context.Enter()
+            next_grad_val = control_flow_ops._NextIteration(grad_val)
+            grad_state.grad_context.Exit()
+          b_merge.op._update_input(1, next_grad_val)
+          # pylint: enable=protected-access
+
+
+def MaybeCreateControlFlowState(between_op_list, between_ops,
+                                colocate_gradients_with_ops):
+  """Create the state for all the while loops involved in one gradients().
+
+  We create a _ControlFlowState when there are while loops involved in
+  gradients(). In gradients(), control flow logic is only invoked when
+  the _ControlFlowState is not None.
+
+  Note that this method modifies `between_op_list` and `between_ops`.
+  """
+  loop_state = None
+  for op in between_op_list:
+    if util.IsLoopExit(op):
+      if loop_state is None:
+        loop_state = _ControlFlowState()
+      if colocate_gradients_with_ops:
+        with ops.colocate_with(op):
+          loop_state.AddWhileContext(op, between_op_list, between_ops)
+      else:
+        loop_state.AddWhileContext(op, between_op_list, between_ops)
+  return loop_state
+
+
+def _ZerosLikeV1(op, index):
+  """Branch of ZerosLike for TF1."""
+  val = op.outputs[index]
+  op_ctxt = op._get_control_flow_context()  # pylint: disable=protected-access
+  if op_ctxt:
+    # We are in a cond context. Use a switch to create zeros only when needed.
+    pred = op_ctxt.pred
+    branch = op_ctxt.branch
+    switch_val = control_flow_ops.switch(op.inputs[0], pred)[1 - branch]
+    # A op is created along the branch taken as control dependencies are on
+    # the whole op and not on the tensor output.
+    pivot = array_ops.identity(switch_val)
+    if val.dtype == dtypes.resource:
+      with ops.control_dependencies([pivot]):
+        return array_ops.zeros(
+            gen_resource_variable_ops.variable_shape(switch_val),
+            dtype=default_gradient.get_zeros_dtype(val))
+    zeros_shape = array_ops.shape_internal(switch_val, optimize=False)
+    # Ensure ops created within array_ops.zeros are dominated by switch in
+    # cond context.
+    with ops.control_dependencies([pivot]):
+      return array_ops.zeros(zeros_shape, dtype=val.dtype)
+  else:
+    return array_ops.zeros_like(val, optimize=False)
+
+
+@array_ops._tag_zeros_tensor  # pylint: disable=protected-access
+def _ConstantZeros(shape, dtype):
+  """Create a constant zero tensor."""
+  return constant_op.constant(0, shape=shape, dtype=dtype)
+
+
+def _ZerosLikeV2(op, index):
+  """Branch of ZerosLike for TF2."""
+  val = op.outputs[index]
+  if val.dtype == dtypes.resource:
+    return array_ops.zeros(
+        gen_resource_variable_ops.variable_shape(val),
+        dtype=default_gradient.get_zeros_dtype(val))
+  if (isinstance(val.op.graph, control_flow_v2_func_graphs.WhileBodyFuncGraph)
+      and val.dtype != dtypes.variant):
+    # In while_v2 we do not want to add a `ZerosLike` op because that will
+    # trigger accumulation of `val`. Normally `ZerosLike` is preferred because
+    # it helps avoid creating extra nodes(possibly Consts) for the shape.
+    # For variants, we must use ZerosLike.
+    if val.shape.is_fully_defined():
+      return _ConstantZeros(val.shape.dims, val.dtype)
+    else:
+      # Note: Even though we add `Shape` in the default graph, while_v2 is smart
+      # enough to place it in the forward graph i.e. `val.graph`.
+      zeros_shape = array_ops.shape_internal(val, optimize=False)
+      return array_ops.zeros(zeros_shape, val.dtype)
+  else:
+    return array_ops.zeros_like(val, optimize=False)
+
+
+def ZerosLike(op, index):
+  """Create zeros_like for the specified output of an op."""
+  if not util.IsSwitch(op):
+    return _ZerosLikeV2(op, index)
+  else:
+    return _ZerosLikeV1(op, index)

tensorflow/python/ops/control_flow_switch_case.py

@@ -0,0 +1,253 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Switch case for Control Flow Operations."""
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond_v2
+from tensorflow.python.ops import control_flow_util as util
+from tensorflow.python.ops import gen_functional_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.util.tf_export import tf_export
+
+
+def _indexed_case_verify_and_canonicalize_args(branch_fns, default,
+                                               branch_index):
+  """Verifies input arguments for the case function.
+
+  Args:
+    branch_fns: Dict or list of pairs of an `int` and a callable which returns a
+      list of tensors.
+    default: Optional callable that returns a list of tensors.
+    branch_index: Optional int `Tensor`, which selects for the corresponding
+      pred_fn_pair.
+
+  Raises:
+    TypeError: If `branch_fns` is not a list/dictionary.
+    TypeError: If `branch_fns` is a list but does not contain 2-tuples or
+               callables.
+    TypeError: If `fns[i]` is not callable for any i, or `default` is not
+               callable.
+
+  Returns:
+    branch_fns: validated list of callables for each branch (default last).
+  """
+  if not isinstance(branch_index, tensor.Tensor):
+    raise TypeError("'branch_index' must be a Tensor, got {}".format(
+        type(branch_index)))
+  if not branch_index.dtype.is_integer:
+    raise TypeError("'branch_index' must be an integer Tensor, got {}".format(
+        branch_index.dtype))
+
+  if not branch_fns:
+    raise ValueError("Must provide at least one item in 'branch_fns'")
+  if not isinstance(branch_fns, (list, tuple, dict)):
+    raise TypeError("'branch_fns' must be a list, tuple, or dict")
+
+  if isinstance(branch_fns, dict):
+    branch_fns = branch_fns.items()
+
+  if all(callable(fn) for fn in branch_fns):
+    branch_fns = list(enumerate(branch_fns))
+
+  for key_fn_pair in branch_fns:
+    if not isinstance(key_fn_pair, tuple) or len(key_fn_pair) != 2:
+      raise TypeError("Each entry in 'branch_fns' must be a 2-tuple. "
+                      f"Received {key_fn_pair}.")
+    key, branch_fn = key_fn_pair
+
+    if not isinstance(key, int):
+      raise TypeError("key must be a Python `int`, got {}".format(type(key)))
+
+    if not callable(branch_fn):
+      raise TypeError("fn for key {} must be callable.".format(key))
+
+  keys = [p[0] for p in branch_fns]
+  if min(keys) < 0 or max(keys) >= len(keys) or len(set(keys)) != len(keys):
+    raise ValueError(
+        "branch indices (keys) must form contiguous range of [0 to {}) but "
+        "found {{{}}}".format(len(keys), ",".join(map(str, sorted(keys)))))
+  actions = [p[1] for p in sorted(branch_fns)]
+  if default is not None:
+    actions.append(default)
+  return actions
+
+
+def _indexed_case_helper(branch_fns,
+                         default,
+                         branch_index,
+                         name,
+                         lower_using_switch_merge=None):
+  """Implementation of case that emits the n-way indexed Case op.
+
+  Args:
+    branch_fns: Dict or list of pairs of a boolean scalar tensor, and a callable
+      which returns a list of tensors.
+    default: Optional callable that returns a list of tensors.
+    branch_index: Optional int `Tensor`, which selects for the corresponding
+      pred_fn_pair.
+    name: A name for this operation (optional).
+    lower_using_switch_merge: Lower this op using switch merge ops (optional).
+
+  Returns:
+    The tensors returned by the pair whose key matched branch_index, or
+    those returned by `default` if none does.
+
+  Raises:
+    TypeError: If `branch_fns` is not a list/dictionary.
+    TypeError: If `branch_fns` is a list but does not contain 2-tuples or
+               callables.
+    TypeError: If `fns[i]` is not callable for any i, or `default` is not
+               callable.
+  """
+  branch_fns = _indexed_case_verify_and_canonicalize_args(
+      branch_fns, default, branch_index)
+  with ops.name_scope(name, "case", [branch_index]):
+    if context.executing_eagerly() and not hasattr(branch_index, "graph"):
+      branch_index = array_ops.where(
+          math_ops.less(branch_index, 0)
+          | math_ops.greater_equal(branch_index, len(branch_fns)),
+          len(branch_fns) - 1, branch_index)
+      return branch_fns[int(branch_index)]()
+    return cond_v2.indexed_case(
+        branch_index,
+        branch_fns,
+        lower_using_switch_merge=lower_using_switch_merge)
+
+
+@tf_export("__internal__.execute_fn_for_device", v1=[])
+def execute_fn_for_device(device_branch_fns, default_fn, name="execute_fn"):
+  """Executes one of the provided callables based on the device placement.
+
+  This API is used when the implementations for high level function depend on
+  the underlying device placement. It takes a dictionary of device type to
+  callables. The device type includes "CPU", "GPU", "TPU", etc. When the type of
+  the device where to run this op matches the key in 'device_branch_fns',
+  the corresponding callable is executed, falling back to 'default_fn' if none
+  matches.
+
+  **Example:**
+  ```python
+  def f1(): return tf.constant(1)
+  def f2(): return tf.constant(2)
+  r = tf.execute_fn_for_device({"CPU": f1, "GPU": f2}, default_fn=f1)
+  ```
+  'r' is evaluated as 1 when it runs on CPU, 2 running on GPU, 1 running on
+  any other device types.
+
+
+  Args:
+    device_branch_fns: a dictionary of device types to the callables. Each
+      callable must return a matching structure of tensors.
+    default_fn: fallback callable when the underlying device does not match any
+      key in the 'device_branch_fns'.
+    name: A name for this operation (optional).
+
+  Returns:
+    The tensors returned by the callable identified by device type during
+    execution, or those returned by 'default_fn' if no key matches.
+  """
+  # Always execute the default fn for XLA to avoid complicated graph by case op.
+  # see more discussions in b/167276293.
+  is_in_xla = util.GraphOrParentsInXlaContext(ops.get_default_graph())
+  if is_in_xla:
+    return default_fn()
+  device_branch_fns_upper = {k.upper(): v for k, v in device_branch_fns.items()}
+  branch_fns = list(device_branch_fns_upper.values())
+  devices = list(device_branch_fns_upper.keys())
+  device_index = gen_functional_ops.device_index(device_names=devices)
+  return _indexed_case_helper(
+      branch_fns,
+      default_fn,
+      device_index,
+      name,
+      lower_using_switch_merge=False)
+
+
+@tf_export("switch_case")
+def switch_case(branch_index, branch_fns, default=None, name="switch_case"):
+  """Create a switch/case operation, i.e.
+
+  an integer-indexed conditional.
+
+  See also `tf.case`.
+
+  This op can be substantially more efficient than `tf.case` when exactly one
+  branch will be selected. `tf.switch_case` is more like a C++ switch/case
+  statement than `tf.case`, which is more like an if/elif/elif/else chain.
+
+  The `branch_fns` parameter is either a dict from `int` to callables, or list
+  of (`int`, callable) pairs, or simply a list of callables (in which case the
+  index is implicitly the key). The `branch_index` `Tensor` is used to select an
+  element in `branch_fns` with matching `int` key, falling back to `default`
+  if none match, or `max(keys)` if no `default` is provided. The keys must form
+  a contiguous set from `0` to `len(branch_fns) - 1`.
+
+  `tf.switch_case` supports nested structures as implemented in `tf.nest`. All
+  callables must return the same (possibly nested) value structure of lists,
+  tuples, and/or named tuples.
+
+  **Example:**
+
+  Pseudocode:
+
+  ```c++
+  switch (branch_index) {  // c-style switch
+    case 0: return 17;
+    case 1: return 31;
+    default: return -1;
+  }
+  ```
+  or
+  ```python
+  branches = {0: lambda: 17, 1: lambda: 31}
+  branches.get(branch_index, lambda: -1)()
+  ```
+
+  Expressions:
+
+  ```python
+  def f1(): return tf.constant(17)
+  def f2(): return tf.constant(31)
+  def f3(): return tf.constant(-1)
+  r = tf.switch_case(branch_index, branch_fns={0: f1, 1: f2}, default=f3)
+  # Equivalent: tf.switch_case(branch_index, branch_fns={0: f1, 1: f2, 2: f3})
+  ```
+
+  Args:
+    branch_index: An int Tensor specifying which of `branch_fns` should be
+      executed.
+    branch_fns: A `dict` mapping `int`s to callables, or a `list` of (`int`,
+      callable) pairs, or simply a list of callables (in which case the index
+      serves as the key). Each callable must return a matching structure of
+      tensors.
+    default: Optional callable that returns a structure of tensors.
+    name: A name for this operation (optional).
+
+  Returns:
+    The tensors returned by the callable identified by `branch_index`, or those
+    returned by `default` if no key matches and `default` was provided, or those
+    returned by the max-keyed `branch_fn` if no `default` is provided.
+
+  Raises:
+    TypeError: If `branch_fns` is not a list/dictionary.
+    TypeError: If `branch_fns` is a list but does not contain 2-tuples or
+               callables.
+    TypeError: If `fns[i]` is not callable for any i, or `default` is not
+               callable.
+  """
+  return _indexed_case_helper(branch_fns, default, branch_index, name)

tensorflow/python/ops/control_flow_util.py

@@ -0,0 +1,367 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Utility functions for control flow.
+
+This file is necessary to avoid cyclic dependencies between ops.py and
+control_flow_ops.py.
+"""
+
+import os
+import traceback
+
+from tensorflow.python import tf2
+from tensorflow.python.platform import tf_logging as logging
+
+ENABLE_CONTROL_FLOW_V2 = ((tf2.enabled() and
+                           os.getenv("TF_ENABLE_CONTROL_FLOW_V2") != "0") or
+                          os.getenv("TF_ENABLE_CONTROL_FLOW_V2", "0") != "0" or
+                          os.getenv("TF_ENABLE_COND_V2", "0") != "0" or
+                          os.getenv("TF_ENABLE_WHILE_V2", "0") != "0" or
+                          os.getenv("TF_ENABLE_TENSOR_ARRAY_V2", "0") != "0")
+
+
+# TODO(b/137793122): Remove this.
+def enable_control_flow_v2():  # pylint: disable=invalid-name
+  """Use control flow v2.
+
+  Do not use this symbol. This will be removed.
+  """
+  global ENABLE_CONTROL_FLOW_V2
+  ENABLE_CONTROL_FLOW_V2 = True
+
+
+def EnableControlFlowV2(graph):
+  """Returns whether control flow v2 should be used in `graph`."""
+  # Enable new control flow in FuncGraphs (but not legacy _FuncGraphs).
+  # TODO(skyewm): do something better than hasattr without messing up imports.
+  return ENABLE_CONTROL_FLOW_V2 or (
+      graph.building_function and not hasattr(graph, "_captured"))
+
+
+def IsInXLAContext(op):
+  try:
+    xla_compile = op.get_attr("_XlaCompile")
+    if xla_compile: return True
+  except ValueError:
+    pass
+  ctxt = op._get_control_flow_context()  # pylint: disable=protected-access
+  return GetContainingXLAContext(ctxt) is not None
+
+
+def InXlaContext(graph):
+  ctxt = graph._get_control_flow_context()  # pylint: disable=protected-access
+  return GetContainingXLAContext(ctxt) is not None
+
+
+def GraphOrParentsInXlaContext(graph):
+  while True:
+    if InXlaContext(graph): return True
+    try:
+      graph = graph.outer_graph
+    except AttributeError:
+      return False
+
+
+def IsInWhileLoop(op):
+  ctxt = op._get_control_flow_context()  # pylint: disable=protected-access
+  return GetContainingWhileContext(ctxt) is not None
+
+
+def IsInCond(op):
+  ctxt = op._get_control_flow_context()  # pylint: disable=protected-access
+  return GetContainingCondContext(ctxt) is not None
+
+
+def IsSwitch(op):
+  """Return true if `op` is a Switch."""
+  return op.type == "Switch" or op.type == "RefSwitch"
+
+
+def IsMerge(op):
+  """Return true if `op` is a Merge."""
+  return op.type == "Merge" or op.type == "RefMerge"
+
+
+def IsLoopEnter(op):
+  """Returns true if `op` is an Enter."""
+  return op.type == "Enter" or op.type == "RefEnter"
+
+
+def IsLoopExit(op):
+  """Return true if `op` is an Exit."""
+  return op.type == "Exit" or op.type == "RefExit"
+
+
+def IsCondSwitch(op):
+  """Return true if `op` is the Switch for a conditional."""
+  if not IsSwitch(op):
+    return False
+  if not op.outputs:
+    return False
+  # Switch nodes are not part of the cond control flow context that they
+  # represent, so consider the consumers of its outputs to determine if it is
+  # cond switch or not. A switch is a cond switch iff all its consumers are in
+  # cond contexts.
+  is_cond_switch = True
+  for o in op.outputs:
+    for c in o.consumers():
+      ctxt = c._get_control_flow_context()  # pylint: disable=protected-access
+      if IsLoopEnter(c):
+        ctxt = ctxt.outer_context
+      is_cond_switch = is_cond_switch and (ctxt is not None and
+                                           ctxt.IsCondContext())
+  return is_cond_switch
+
+
+def IsCondMerge(op):
+  """Return true if `op` is the Merge for a conditional."""
+  if not IsMerge(op):
+    return False
+  if not op.inputs:
+    return False
+  # Merge nodes are not part of the cond control flow context that they
+  # represent, so consider the inputs to the merge of to determine if it is
+  # cond merge or not: A merge is a cond merge iff all its inputs are in
+  # cond contexts.
+  is_cond_merge = True
+  for i in op.inputs:
+    ctxt = GetOutputContext(i.op)
+    is_cond_merge = is_cond_merge and ctxt is not None and ctxt.IsCondContext()
+  return is_cond_merge
+
+
+def IsLoopSwitch(op):
+  """Return true if `op` is the Switch for a while loop."""
+  if IsSwitch(op):
+    ctxt = op._get_control_flow_context()  # pylint: disable=protected-access
+    return ctxt is not None and ctxt.IsWhileContext() and not IsCondSwitch(op)
+  return False
+
+
+def IsLoopMerge(op):
+  """Return true if `op` is the Merge for a while loop."""
+  if IsMerge(op):
+    ctxt = op._get_control_flow_context()  # pylint: disable=protected-access
+    return ctxt is not None and ctxt.IsWhileContext() and not IsCondMerge(op)
+  return False
+
+
+def IsLoopConstantEnter(op):
+  """Return true iff op is a loop invariant."""
+  return IsLoopEnter(op) and op.get_attr("is_constant")
+
+
+def GetLoopConstantEnter(value):
+  """Return the enter op if we can infer `value` to be a loop invariant."""
+  id_ops = {"Switch", "RefSwitch", "Identity", "RefIdentity"}
+  op = value.op
+  while op.type in id_ops:
+    op = op.inputs[0].op
+  return op if IsLoopConstantEnter(op) else None
+
+
+def GetOutputContext(op):
+  """Return the control flow context for the output of an op."""
+  ctxt = op._get_control_flow_context()  # pylint: disable=protected-access
+  # Exit nodes usually have a control flow context, except in the case where the
+  # exit node was imported via import_graph_def (in which case no nodes have
+  # control flow contexts).
+  if ctxt is not None and IsLoopExit(op):
+    ctxt = ctxt.outer_context
+  return ctxt
+
+
+def GetContainingWhileContext(ctxt, stop_ctxt=None):
+  """Returns the first ancestor WhileContext of `ctxt`.
+
+  Returns `ctxt` if `ctxt` is a WhileContext, or None if `ctxt` is not in a
+  while loop.
+
+  Args:
+    ctxt: ControlFlowContext
+    stop_ctxt: ControlFlowContext, optional. If provided, the search will end
+      if it sees stop_ctxt.
+
+  Returns:
+    `ctxt` if `ctxt` is a WhileContext, the most nested WhileContext containing
+    `ctxt`, or None if `ctxt` is not in a while loop.  If `stop_ctxt` is not
+    `None`, this returns `ctxt` if it matches `stop_ctxt` in its traversal.
+  """
+  while ctxt:
+    if ctxt.IsWhileContext() or ctxt == stop_ctxt: return ctxt
+    ctxt = ctxt.outer_context
+  return None
+
+
+def GetContainingXLAContext(ctxt):
+  """Returns the first ancestor XLAContext of `ctxt`.
+
+  Returns `ctxt` if `ctxt` is a XLAContext, or None if `ctxt` is not in a
+  while loop.
+
+  Args:
+    ctxt: ControlFlowContext
+
+  Returns:
+    `ctxt` if `ctxt` is a XLAContext, the most nested XLAContext containing
+    `ctxt`, or None if `ctxt` is not in a while loop.
+  """
+  while ctxt:
+    if ctxt.IsXLAContext(): return ctxt
+    ctxt = ctxt.outer_context
+  return None
+
+
+def GetContainingCondContext(ctxt):
+  """Returns the first ancestor CondContext of `ctxt`.
+
+  Returns `ctxt` if `ctxt` is a CondContext, or None if `ctxt` is not in a cond.
+
+  Args:
+    ctxt: ControlFlowContext
+
+  Returns:
+    `ctxt` if `ctxt` is a CondContext, the most nested CondContext containing
+    `ctxt`, or None if `ctxt` is not in a cond.
+  """
+  while ctxt:
+    if ctxt.IsCondContext(): return ctxt
+    ctxt = ctxt.outer_context
+  return None
+
+
+def IsContainingContext(ctxt, maybe_containing_ctxt):
+  """Returns true if `maybe_containing_ctxt` is or contains `ctxt`."""
+  while ctxt is not maybe_containing_ctxt:
+    if ctxt is None: return False
+    ctxt = ctxt.outer_context
+  return True
+
+
+def OpInContext(op, ctxt):
+  return IsContainingContext(op._get_control_flow_context(), ctxt)  # pylint: disable=protected-access
+
+
+def TensorInContext(tensor, ctxt):
+  return OpInContext(tensor.op, ctxt)
+
+
+def CheckInputFromValidContext(op, input_op):
+  """Returns whether `input_op` can be used from `op`s context.
+
+  Conceptually, only inputs from op's while context or any ancestor while
+  context (including outside of any context) are valid. In practice, there are
+  many other edge cases as well.
+
+  Args:
+    op: Operation
+    input_op: Operation
+
+  Raises:
+    ValueError: if input_op is from an invalid context.
+  """
+  op_ctxt = op._get_control_flow_context()  # pylint: disable=protected-access
+  input_ctxt = GetOutputContext(input_op)
+  valid = False
+
+  if not input_ctxt:
+    # input_op isn't in a control flow context.
+    valid = True
+  elif op_ctxt is input_ctxt:
+    # input_op is in the same context as op.
+    valid = True
+  else:
+    while_ctxt = GetContainingWhileContext(op_ctxt)
+    input_while_ctxt = GetContainingWhileContext(input_ctxt)
+
+    if while_ctxt is None:
+      if input_while_ctxt is None:
+        # Neither op nor input_op is in a while loop, but one or both are in
+        # conds. We allow this, although execution will fail if the branch
+        # corresponding to input_op's cond context isn't taken.
+        valid = True
+      # Invalid if op isn't in a while loop and input_op is. Unless...
+      if IsLoopEnter(op):
+        # WhileContext._BuildLoop clears context for Enter nodes.
+        valid = True
+      if IsSwitch(op):
+        # CondContext.AddValue clears context for Switch nodes.
+        valid = True
+    elif IsContainingContext(while_ctxt, input_while_ctxt):
+      # input_op is in a while loop which contains op's while loop (or not in a
+      # while loop at all).
+      valid = True
+    elif (while_ctxt.grad_state and
+          IsContainingContext(while_ctxt.grad_state.forward_context,
+                              input_while_ctxt)):
+      # op is in a gradient context and input_op is in the associated forward
+      # pass context or an ancestor thereof. This case is need to build while
+      # loop gradients.
+      # NOTE(skyewm): we theoretically also need this case for custom gradient
+      # functions that close over tensors from ancestor contexts, but I haven't
+      # verified this.
+      valid = True
+    elif (while_ctxt.grad_state and
+          while_ctxt.grad_state.forward_context is
+          input_while_ctxt._outer_context):  # pylint: disable=protected-access
+      # op is in a gradient context and input_op is in a child of the associated
+      # forward pass context. This case is needed for the gradients of while
+      # loops with conds.
+      valid = True
+    elif (input_while_ctxt.grad_state and
+          input_while_ctxt.grad_state.forward_context is while_ctxt):
+      # input_op is in the gradient context of op's context. This case is needed
+      # when the gradient of a while loop gradient is requested (this will
+      # eventually fail unless there is a stop_gradient() or similar).
+      valid = True
+    elif (input_while_ctxt.grad_state and
+          input_ctxt.grad_state.forward_context.grad_state and
+          input_ctxt.grad_state.forward_context.grad_state.forward_context is
+          while_ctxt):
+      # input_op is in the grad grad context of op's context. This case is
+      # needed when the gradient of a while loop gradient is requested (this
+      # will eventually fail unless there is a stop_gradient() or similar).
+      valid = True
+
+  if not valid:
+    if while_ctxt:
+      error_msg = (
+          f"Cannot use '{input_op.name}' as input to '{op.name}' because they "
+          "are in different while loops.")
+    else:
+      error_msg = (
+          f"Cannot use '{input_op.name}' as input to '{op.name}' because "
+          f"'{input_op.name}' is in a while loop.")
+
+    # Log the error message plus the relevant stack traces. The stacks may be
+    # useful for debugging this error, but we don't want to raise an
+    # unreadable exception.
+    log_msg = error_msg
+    log_msg += "\n\n%s while context: %s" % (op.name, while_ctxt)
+    log_msg += "\n%s while context: %s" % (input_op.name, input_while_ctxt)
+    log_msg += "\n\nTraceback for %s:\n%s\nTraceback for %s:\n%s\n" % (
+        op.name, "".join(traceback.format_list(op.traceback)),
+        input_op.name, "".join(traceback.format_list(input_op.traceback)))
+    logging.info(log_msg)
+    raise ValueError(error_msg + " See info log for more details.")
+
+
+def GetWhileContext(op):
+  """Get the WhileContext to which this op belongs."""
+  ctxt = op._get_control_flow_context()  # pylint: disable=protected-access
+  if ctxt:
+    ctxt = ctxt.GetWhileContext()
+  return ctxt

tensorflow/python/ops/control_flow_util_v2.py

@@ -0,0 +1,429 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Utilities for V2 control flow."""
+
+from tensorflow.core.framework import attr_value_pb2
+from tensorflow.python.eager import context
+from tensorflow.python.eager.polymorphic_function import atomic_function
+from tensorflow.python.eager.polymorphic_function import concrete_function
+from tensorflow.python.eager.polymorphic_function import tracing_compilation
+from tensorflow.python.eager.polymorphic_function import transform
+from tensorflow.python.framework import function_def_to_graph
+from tensorflow.python.framework import ops
+from tensorflow.python.framework.func_graph import FuncGraph
+from tensorflow.python.ops import control_flow_util
+from tensorflow.python.ops import control_flow_v2_func_graphs
+from tensorflow.python.ops import gradients_util
+from tensorflow.python.util import keras_deps
+from tensorflow.python.util import tf_contextlib
+from tensorflow.python.util.tf_export import tf_export
+
+_EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE = None
+_DISABLE_LOWER_USING_SWITCH_MERGE = False
+
+
+CondBranchFuncGraph = control_flow_v2_func_graphs.CondBranchFuncGraph
+WhileCondFuncGraph = control_flow_v2_func_graphs.WhileCondFuncGraph
+WhileBodyFuncGraph = control_flow_v2_func_graphs.WhileBodyFuncGraph
+
+
+def in_defun():
+  """Returns if the current graph is, or is nested in, a defun."""
+  if context.executing_eagerly(): return False
+
+  graph = ops.get_default_graph()
+  while (isinstance(graph, CondBranchFuncGraph) or
+         isinstance(graph, WhileBodyFuncGraph) or
+         isinstance(graph, WhileCondFuncGraph)):
+    graph = graph.outer_graph
+  return isinstance(graph, FuncGraph)
+
+
+def in_while_loop_defun(graph):
+  """Returns if the graph is a while loop FuncGraph."""
+  if context.executing_eagerly(): return False
+  return (isinstance(graph, WhileCondFuncGraph) or
+          isinstance(graph, WhileBodyFuncGraph))
+
+
+def create_new_tf_function(func_graph):
+  """Converts func_graph to a TF_Function and adds it to the current graph.
+
+  Args:
+    func_graph: FuncGraph
+
+  Returns:
+    The name of the new TF_Function.
+  """
+  transform.apply_func_graph_transforms(func_graph)
+  func = atomic_function.from_func_graph(func_graph.name, func_graph, {})
+
+  func_graph.outer_graph._add_function_recursive(func)  # pylint: disable=protected-access
+  return func_graph.name
+
+
+def unique_fn_name(scope, name):
+  """Returns a unique name to use for a control flow function.
+
+  Args:
+    scope: A name scope string.
+    name: An identifier for this function (e.g. "true", "body").
+
+  Returns:
+    A string, the name to use for the function.
+  """
+  return ("%s%s_%s" % (scope, name, ops.uid())).replace("/", "_")
+
+
+def unique_grad_fn_name(forward_name):
+  return "%s_grad_%s" % (forward_name, ops.uid())
+
+
+def maybe_set_lowering_attr(op, lower_using_switch_merge=None):
+  """Sets the flag to enable lowering on `op` if necessary.
+
+  Lowering allows cond_v2 and while_v2 to avoid some of the limitations of
+  Functions, allowing users to specify devices & colocation inside of cond_v2
+  and while_v2 input functions, and enabling non-strict evaluation & partial
+  pruning. This brings v2 control flow closer to feature parity with v1 control
+  flow.
+
+  However, we do not lower in the following cases:
+    - When the `If` or `While` ops are in the XLA context. Because it is easier
+      for XLA to apply its own optimizations when dealing with un-lowered
+      control flow operators than with low-level control flow primitives.
+    - When the eager execution context specifies the executor of functions to
+      be the single threaded executor (see context.function_executor_type()).
+      Because the single threaded executor does not support v1 control flow ops.
+    - When 'lower_using_switch_merge' is explicitly set to False.
+
+  Args:
+    op: An `If` or `While` Operation.
+    lower_using_switch_merge: Explicit value to lower or not (optional).
+  """
+  if lower_using_switch_merge is not None:
+    # pylint: disable=protected-access
+    op._set_attr("_lower_using_switch_merge",
+                 attr_value_pb2.AttrValue(b=lower_using_switch_merge))
+    # pylint: enable=protected-access
+  elif (not _DISABLE_LOWER_USING_SWITCH_MERGE and
+        not control_flow_util.GraphOrParentsInXlaContext(op.graph) and
+        context.context().function_call_options.executor_type !=
+        "SINGLE_THREADED_EXECUTOR"):
+    # pylint: disable=protected-access
+    op._set_attr("_lower_using_switch_merge", attr_value_pb2.AttrValue(b=True))
+    # pylint: enable=protected-access
+
+
+def maybe_propagate_compile_time_consts_in_xla(op):
+  """Tells XLA whether to propagate compile-time consts in the loop body.
+
+  This is needed to make compile time constants available to ops, for example
+  `max_num_elements` in `EmptyTensorList`, inside the loop body. Ideally this
+  would always be turned on, but that doesn't work with legacy functionalized
+  while_loops.
+
+  Args:
+    op: A `While` Operation.
+  """
+  if control_flow_util.GraphOrParentsInXlaContext(op.graph):
+    # pylint: disable=protected-access
+    op._set_attr("_xla_propagate_compile_time_consts",
+                 attr_value_pb2.AttrValue(b=True))
+    # pylint: enable=protected-access
+
+
+def resource_input_index(tensor_name, input_names, node_defs, functions):
+  """Returns the index of the input corresponding to `tensor_name`.
+
+  This method is used to find the corresponding index of an arbitrary resource
+  tensor in a function (the function could be a loop body). We assume that
+  resource handles are never created in functions, so that every resource
+  tensor can be traced back to a function input.
+
+  The awkward signature of this method is to make it work with both FuncGraphs
+  and FunctionDefs. This is so we can recurse on function call ops without
+  building the corresponding FuncGraph (note that even if a FuncGraph for a
+  FunctionDef already exists, the input/output/node names may have been
+  changed when the FuncGraph was serialized to the FunctionDef, which makes it
+  unusable with this algorithm).
+
+  Args:
+    tensor_name: the name of the resource tensor to be resolved to an input.
+    input_names: a list of the names of all inputs to the function.
+    node_defs: a dict mapping op name -> NodeDef for every op in the function.
+    functions: a dict mapping function name -> AtomicFunction.
+
+  Returns:
+    The index into input_names corresponding to `tensor_name`.
+  """
+  while tensor_name not in input_names:
+    # FunctionDefs and graphs use different tensor naming conventions.
+    parts = tensor_name.split(":")
+    if len(parts) == 3:
+      op_name, _, output_idx = parts
+    elif len(parts) == 2:
+      op_name, output_idx = parts
+    else:
+      assert len(parts) == 1
+      op_name = parts[0]
+      output_idx = 0
+      tensor_name = "%s:%d" % (tensor_name, output_idx)
+      # Check again for cases where the tensor suffix (":0") is stripped out.
+      if tensor_name in input_names:
+        break
+    output_idx = int(output_idx)
+    node_def = node_defs[op_name]
+
+    def _extract_input_index(function_attribute_name):
+      func_name = node_def.attr[function_attribute_name].func.name
+      fdef = functions[func_name].cached_definition
+      output_arg_name = fdef.signature.output_arg[output_idx].name
+      output_tensor_name = fdef.ret[output_arg_name]
+      return resource_input_index(
+          output_tensor_name, [arg.name for arg in fdef.signature.input_arg],
+          {ndef.name: ndef for ndef in fdef.node_def}, functions)
+
+    if node_def.op in ("Identity", "While"):
+      # Captured resources occur at the same index in the lists of inputs and
+      # outputs of a while or identity op. So we lookup the input of `tensor.op`
+      # at the same index as the index of `tensor` in the `tensor.op.outputs`.
+      tensor_name = node_def.input[output_idx]
+    elif node_def.op in ("PartitionedCall", "StatefulPartitionedCall"):
+      # Functions output any captured resource tensors used by their
+      # gradients.  `tensor_name` is one of these outputs from a nested
+      # function call, so recursively find the corresponding input in the
+      # nested FunctionDef.
+      tensor_name = node_def.input[_extract_input_index("f")]
+    elif node_def.op in ("If", "StatelessIf"):
+      input_index = _extract_input_index("then_branch")
+      if input_index != _extract_input_index("else_branch"):
+        raise AssertionError(
+            ("Expected cond branches ({} op) to each have the same "
+             "input->output mapping of resources.").format(node_def.op))
+      tensor_name = node_def.input[
+          # Ignore the `cond` input; the function inputs come after.
+          input_index + 1]
+    else:
+      # We assume there are no other ops types that will "forward" resource
+      # handles like this, so all other handles must have been created by the
+      # op. (Note that cond_v2 wraps resource handle outputs in optionals,
+      # which we'll end up accumulating).
+      raise ValueError("Taking gradient of a while loop which creates "
+                       "a resource in its body is not supported: %s (%s)"
+                       % (op_name, node_def.op))
+
+  return input_names.index(tensor_name)
+
+
+@tf_contextlib.contextmanager
+def clear_control_inputs():
+  """Clears the control inputs but preserves the ControlFlowContext.
+
+  This is needed to preserve the XLAControlFlowControl when clearing
+  control inputs for the gradient accumulators in while_v2.
+  `ops.control_dependencies` does not allow that.
+
+  Yields:
+    A context manager in which the ops created will not have any control inputs
+    by default but the control flow context is the same.
+  """
+  # pylint: disable=protected-access
+  control_flow_context = ops.get_default_graph()._get_control_flow_context()
+  with ops.control_dependencies(None):
+    ops.get_default_graph()._set_control_flow_context(control_flow_context)
+    yield
+  # pylint: enable=protected-access
+
+
+def _is_tpu_strategy(strategy):
+  return (strategy is not None and
+          strategy.__class__.__name__.startswith("TPUStrategy"))
+
+
+def _is_building_keras_layer():
+  # TODO(srbs): Remove this function when we no long support session with Keras.
+  keras_call_context_function = keras_deps.get_call_context_function()
+  if keras_call_context_function:
+    return keras_call_context_function().layer is not None
+  else:
+    return False
+
+
+def output_all_intermediates():
+  """Whether to output all intermediates of a functional control flow op.
+
+  The default behavior is to output intermediates only when building a Keras
+  Layer in graph mode and that too when certain other conditions are met:
+  1. We do not output intermediates if the functional control flow op
+     is being built inside a FuncGraph which is not a If/While graph. This
+     guards against outputting intermediates in eager mode since keras adds
+     tensors to a FuncGraph named "keras_graph" in that case. Also because we
+     do not output intermediates of tf.function (since this feature is only for
+     backwards compatibility) outputting intermediates of functional control
+     flow ops built inside tf.function is of no value.
+  2. We do not output intermediates when the compilation is using XLA or for a
+     TPU.
+  3. We do not output intermediates when a single threaded executor is used
+     since that does not perform inlining and pruning.
+
+  Returns:
+    A bool telling whether to output all intermediates.
+  """
+  if _EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE is not None:
+    return _EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE
+  if in_defun():
+    return False
+  if control_flow_util.GraphOrParentsInXlaContext(ops.get_default_graph()):
+    return False
+  if (context.context().function_call_options.executor_type ==
+      "SINGLE_THREADED_EXECUTOR"):
+    return False
+  return _is_building_keras_layer()
+
+
+def get_func_graph(op, input_shapes, func_name):
+  """Generates and returns a FuncGraph for the given op and input_shapes."""
+  fdef = None
+  graph = op.graph
+  # Recursively search the func in graphs.
+  while graph is not None:
+    func = graph._get_function(func_name)  # pylint: disable=protected-access
+    if func is not None:
+      fdef = func.cached_definition
+      break
+    if hasattr(graph, "outer_graph"):
+      graph = graph.outer_graph
+    else:
+      break
+
+  if fdef is None:
+    raise KeyError("%s cannot be found in the graph" % func_name)
+
+  # `op.graph` may not be the same as `ops.get_default_graph()` e.g.
+  # in the case of nested if ops or when the gradient is being computed
+  # from inside a Defun. We build the `func_graph` with `op.graph` as its
+  # `outer_graph`. This resembles how the `FuncGraph` was built in the
+  # forward pass. We need this so that we can resolve references to tensors
+  # in `func_graph` from its gradient graph in `_resolve_grad_inputs`.
+  with op.graph.as_default():
+    func_graph = function_def_to_graph.function_def_to_graph(
+        fdef, input_shapes=input_shapes)
+
+  # TODO(xjun): Ideally we want to retrieve the gradient functions instead of
+  # re-create them. But the lifetime of gradient functions of PartitionedCall
+  # ops is attached to ParitionedCall ops in the original func_graph and
+  # when we are inside this function we don't have access to the original
+  # func_graph or PartitionedCall ops. See cl/499362867 and cl/273858076 for
+  # more context.
+  for operation in func_graph.get_operations():
+    if operation.type in ["PartitionedCall", "StatefulPartitionedCall"]:
+      f = graph._get_function(operation.get_attr("f").name)  # pylint: disable=protected-access
+      try:
+        cf = concrete_function.ConcreteFunction.from_func_graph(
+            f.graph,
+            f.function_type,
+            attrs=f.cached_definition.attr,
+        )
+      except AttributeError:
+        # f is not found or f is a _DefinedFunction that doesn't have a graph.
+        continue
+      operation._gradient_function = cf._get_gradient_function()  # pylint: disable=protected-access
+
+  return func_graph
+
+
+def get_op_and_outputs(op_or_outputs):
+  if isinstance(op_or_outputs, ops.Operation):
+    return op_or_outputs, []
+  elif not op_or_outputs:  # Empty list.
+    return None, []
+  else:
+    return op_or_outputs[0].op, op_or_outputs
+
+
+def graph_wrapped_for_higher_order_tape_gradients(graph):
+  """Check if `graph` is wrapped by `run_as_function_for_tape_gradients`."""
+  while graph is not None:
+    if "cflow_gradient_wrapper" in getattr(graph, "name", ""):
+      return True
+    graph = getattr(graph, "outer_graph", None)
+  return False
+
+
+def run_as_function_for_tape_gradients(make_op, inputs):
+  """Fix higher-order tape gradients by wrapping `make_op` in a function.
+
+  Args:
+    make_op: A function that takes a list of inputs and returns a list of output
+      tensors. This function should set any handle data relevant to its outputs
+      before returning.
+    inputs: A list of tensors to check for tape gradients and pass to
+      `make_op`. These should include all tensors used in `make_op`.
+
+  Returns:
+    Tensors corresponding to `make_op`'s output.
+  """
+  # GradientTapes created inside a function currently don't work well with
+  # un-wrapped control flow ops in that same function. Wrapping in an extra
+  # layer of intermediate function means we run extra logic in the function
+  # gradient code to record the correct intermediates on the tape.
+  #
+  # The function attribute inputs to control flow ops are not hashable, so we
+  # pass everything as a capture to bypass defun's caching.
+  if (gradients_util.PossibleTapeGradientTypes(inputs)
+      == gradients_util.POSSIBLE_GRADIENT_TYPES_HIGHER_ORDER
+      # We only need one function between the tape and the op; if we've already
+      # wrapped once, we stop wrapping to avoid infinite recursion.
+      and not (ops.get_default_graph().building_function
+               and "cflow_gradient_wrapper" in ops.get_default_graph().name)):
+    results = tracing_compilation.call_function(
+        (inputs,),
+        tracing_options=tracing_compilation.TracingOptions(
+            make_op, "cflow_gradient_wrapper", autograph=False
+        ),
+    )
+    return results
+  else:
+    return make_op(inputs)
+
+
+@tf_export(v1=["experimental.output_all_intermediates"])
+def set_output_all_intermediates(state):  # pylint: disable=invalid-name
+  """Whether to output all intermediates from functional control flow ops.
+
+  The "default" behavior to is to output all intermediates when using v2 control
+  flow inside Keras models in graph mode. This is needed to support taking
+  gradients of v2 control flow. In graph mode, Keras can sometimes freeze the
+  forward graph before the gradient computation which does not work for v2
+  control flow since it requires updating the forward ops to output the needed
+  intermediates. We work around this by proactively outputting the needed
+  intermediates when building the forward pass itself. Ideally any such extra
+  tensors should be pruned out at runtime. However, if for any reason this
+  doesn't work for you or if you have an inference-only model you can turn this
+  behavior off using
+  `tf.compat.v1.experimental.output_all_intermediates(False)`.
+
+  If with the default behavior you are still seeing errors of the form
+  "Connecting to invalid output X of source node Y which has Z outputs" try
+  setting `tf.compat.v1.experimental.output_all_intermediates(True)` and
+  please file an issue at https://github.com/tensorflow/tensorflow/issues.
+
+  Args:
+    state: True, False or None. None restores the default behavior.
+  """
+  global _EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE
+  _EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE = state  # pylint: disable=protected-access

tensorflow/python/ops/control_flow_v2_disable_test.py

@@ -0,0 +1,35 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests that TF2_BEHAVIOR=1 and TF_ENABLE_CONTROL_FLOW_V2=0 disables cfv2."""
+
+import os
+os.environ["TF2_BEHAVIOR"] = "1"
+os.environ["TF_ENABLE_CONTROL_FLOW_V2"] = "0"
+
+from tensorflow.python import tf2  # pylint: disable=g-import-not-at-top
+from tensorflow.python.ops import control_flow_util
+from tensorflow.python.platform import googletest
+from tensorflow.python.platform import test
+
+
+class ControlFlowV2DisableTest(test.TestCase):
+
+  def testIsDisabled(self):
+    self.assertTrue(tf2.enabled())
+    self.assertFalse(control_flow_util.ENABLE_CONTROL_FLOW_V2)
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/python/ops/control_flow_v2_enable_test.py

@@ -0,0 +1,34 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests that TF2_BEHAVIOR=1 enables cfv2."""
+
+import os
+os.environ["TF2_BEHAVIOR"] = "1"
+
+from tensorflow.python import tf2  # pylint: disable=g-import-not-at-top
+from tensorflow.python.ops import control_flow_util
+from tensorflow.python.platform import googletest
+from tensorflow.python.platform import test
+
+
+class ControlFlowV2EnableTest(test.TestCase):
+
+  def testIsEnabled(self):
+    self.assertTrue(tf2.enabled())
+    self.assertTrue(control_flow_util.ENABLE_CONTROL_FLOW_V2)
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/python/ops/control_flow_v2_func_graphs.py

@@ -0,0 +1,56 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""FuncGraphs for V2 control flow."""
+
+from tensorflow.python.framework import func_graph
+from tensorflow.python.framework import ops
+
+
+class ControlFlowFuncGraph(func_graph.FuncGraph):
+  """Contains control flow-specific FuncGraph logic."""
+
+  def __init__(self, *args, **kwargs):
+    super(ControlFlowFuncGraph, self).__init__(*args, **kwargs)
+    outer_graph = self.outer_graph
+    # Unlike tf.function, control flow FuncGraphs are generally created one per
+    # op. This means hard-coding any outer device scopes in the body (rather
+    # than inspecting the call-time placement of the control flow op) makes
+    # sense.
+    self._device_function_stack = outer_graph._device_function_stack.copy()  # pylint: disable=protected-access
+    self.is_control_flow_graph = True
+    if ops.executing_eagerly_outside_functions():
+      func_graph.override_func_graph_name_scope(
+          self, self.outer_graph.get_name_scope())
+
+
+class CondBranchFuncGraph(ControlFlowFuncGraph):
+  """FuncGraph for branches of tf.cond().
+
+  This is used to distinguish cond branches from other functions.
+  """
+
+
+class WhileCondFuncGraph(ControlFlowFuncGraph):
+  """FuncGraph for the condition of tf.while_loop().
+
+  This is used to distinguish while conditions from other functions.
+  """
+
+
+class WhileBodyFuncGraph(ControlFlowFuncGraph):
+  """FuncGraph for the body of tf.while_loop().
+
+  This is used to distinguish while bodies from other functions.
+  """

tensorflow/python/ops/control_flow_v2_toggles.py

@@ -0,0 +1,69 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""API for enabling v2 control flow."""
+
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import control_flow_util
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["enable_control_flow_v2"])
+def enable_control_flow_v2():  # pylint: disable=invalid-name
+  """Use control flow v2.
+
+  control flow v2 (cfv2) is an improved version of control flow in TensorFlow
+  with support for higher order derivatives. Enabling cfv2 will change the
+  graph/function representation of control flow, e.g., `tf.while_loop` and
+  `tf.cond` will generate functional `While` and `If` ops instead of low-level
+  `Switch`, `Merge` etc. ops. Note: Importing and running graphs exported
+  with old control flow will still be supported.
+
+  Calling tf.enable_control_flow_v2() lets you opt-in to this TensorFlow 2.0
+  feature.
+
+  Note: v2 control flow is always enabled inside of tf.function. Calling this
+  function is not required.
+  """
+  # pylint: disable=protected-access
+  logging.vlog(1, "Enabling control flow v2")
+  ops._control_flow_api_gauge.get_cell().set(True)
+  control_flow_util.ENABLE_CONTROL_FLOW_V2 = True
+
+
+@tf_export(v1=["disable_control_flow_v2"])
+def disable_control_flow_v2():  # pylint: disable=invalid-name
+  """Opts out of control flow v2.
+
+  Note: v2 control flow is always enabled inside of tf.function. Calling this
+  function has no effect in that case.
+
+  If your code needs tf.disable_control_flow_v2() to be called to work
+  properly please file a bug.
+  """
+  # pylint: disable=protected-access
+  logging.vlog(1, "Disabling control flow v2")
+  ops._control_flow_api_gauge.get_cell().set(False)
+  control_flow_util.ENABLE_CONTROL_FLOW_V2 = False
+
+
+@tf_export(v1=["control_flow_v2_enabled"])
+def control_flow_v2_enabled():  # pylint: disable=invalid-name
+  """Returns `True` if v2 control flow is enabled.
+
+  Note: v2 control flow is always enabled inside of tf.function.
+  """
+  return control_flow_util.EnableControlFlowV2(ops.get_default_graph())

tensorflow/python/ops/control_flow_v2_toggles_test.py

@@ -0,0 +1,39 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for control_flow_v2_toggles.py."""
+
+from tensorflow.python.ops import control_flow_util_v2
+from tensorflow.python.platform import googletest
+from tensorflow.python.platform import test
+
+
+class ControlFlowV2TogglesTest(test.TestCase):
+
+  def testOutputAllIntermediates(self):
+    self.assertIsNone(
+        control_flow_util_v2._EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE)
+    control_flow_util_v2.set_output_all_intermediates(True)
+    self.assertTrue(
+        control_flow_util_v2._EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE)
+    control_flow_util_v2.set_output_all_intermediates(False)
+    self.assertFalse(
+        control_flow_util_v2._EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE)
+    control_flow_util_v2.set_output_all_intermediates(None)
+    self.assertIsNone(
+        control_flow_util_v2._EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/conv2d_benchmark.py

@@ -0,0 +1,208 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Benchmark for Conv2D op."""
+
+import itertools
+import time
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.core.protobuf import rewriter_config_pb2
+from tensorflow.python.client import session as session_lib
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import flags
+from tensorflow.python.platform import test
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_boolean(
+    "enable_layout_optimizer", False,
+    "If true, enables layout optimizer to update input data format for faster "
+    "execution of convolution ops.")
+
+
+def build_graph(device, dtype, data_format, input_shape, filter_shape, strides,
+                padding, num_iters, warmup_iters):
+  """builds a graph containing a sequence of conv2d operations.
+
+  Args:
+    device: String, the device to run on.
+    dtype: Data type for the convolution.
+    data_format: A string from: "NHWC" or "NCHW". Data format for input and
+                 output data.
+    input_shape: Shape of the input tensor.
+    filter_shape: Shape of the filter tensor.
+    strides: A list of ints. 1-D of length 4. The stride of sliding
+             window for each dimension of input.
+    padding: A string from: "SAME", "VALID". The type of padding
+             algorithm to use.
+    num_iters: number of iterations to run conv2d.
+    warmup_iters: number of iterations for warmup runs.
+
+  Returns:
+    An array of tensors to run()
+  """
+  with ops.device("/%s:0" % device):
+    inp = variable_v1.VariableV1(
+        random_ops.truncated_normal(input_shape, dtype=dtype))
+    filt = variable_v1.VariableV1(
+        random_ops.truncated_normal(filter_shape, dtype=dtype))
+
+    outputs = []
+    conv2d_op = nn_ops.conv2d(
+        inp, filt, strides, padding, data_format=data_format)
+    outputs.append(conv2d_op)
+    for _ in range(1, num_iters):
+      with ops.control_dependencies([conv2d_op]):
+        conv2d_op = nn_ops.conv2d(
+            inp, filt, strides, padding, data_format=data_format)
+        outputs.append(conv2d_op)
+
+    warmup_groups = []
+    warmup_conv2d_op = nn_ops.conv2d(
+        inp, filt, strides, padding, data_format=data_format)
+    warmup_groups.append(warmup_conv2d_op)
+    for _ in range(1, warmup_iters):
+      with ops.control_dependencies([warmup_conv2d_op]):
+        warmup_conv2d_op = nn_ops.conv2d(
+            inp, filt, strides, padding, data_format=data_format)
+        warmup_groups.append(warmup_conv2d_op)
+    return control_flow_ops.group(*warmup_groups), control_flow_ops.group(
+        *outputs)
+
+
+class Conv2DBenchmark(test.Benchmark):
+  """Benchmark conv2d!"""
+
+  def _run_graph(self, device, dtype, data_format, input_shape, filter_shape,
+                 strides, padding, num_iters, warmup_iters):
+    """runs the graph and print its execution time.
+
+    Args:
+      device: String, the device to run on.
+      dtype: Data type for the convolution.
+      data_format: A string from: "NHWC" or "NCHW". Data format for input and
+                   output data.
+      input_shape: Shape of the input tensor.
+      filter_shape: Shape of the filter tensor.
+      strides: A list of ints. 1-D of length 4. The stride of sliding
+               window for each dimension of input.
+      padding: A string from: "SAME", "VALID". The type of padding
+               algorithm to use.  num_iters: Number of iterations to run the
+                 benchmark.
+      num_iters: number of iterations to run conv2d.
+      warmup_iters: number of iterations for warmup runs.
+
+    Returns:
+      The duration of the run in seconds.
+    """
+    graph = ops.Graph()
+    with graph.as_default():
+      warmup_outputs, outputs = build_graph(device, dtype, data_format,
+                                            input_shape, filter_shape, strides,
+                                            padding, num_iters, warmup_iters)
+
+      config = config_pb2.ConfigProto()
+      config.graph_options.optimizer_options.opt_level = -1
+      rewrite_options = config.graph_options.rewrite_options
+
+      # Disable layout optimizer to not change input data_format.
+      rewrite_options.layout_optimizer = (
+          rewriter_config_pb2.RewriterConfig.ON if FLAGS.enable_layout_optimizer
+          else rewriter_config_pb2.RewriterConfig.OFF)
+      # Convolution ops are effectively noop in the test graph as we are not
+      # fetching the convolution outputs. Disable dependency optimizer to not
+      # remove the conv ops.
+      rewrite_options.dependency_optimization = (
+          rewriter_config_pb2.RewriterConfig.OFF)
+
+      with session_lib.Session(graph=graph, config=config) as session:
+        # TODO(hinsu): Use run_op_benchmark method from test.Benchmark to run
+        # benchmark along with warmup.
+        variables.global_variables_initializer().run()
+        # warmup runs
+        session.run(warmup_outputs)
+
+        start_time = time.time()
+        session.run(outputs)
+        duration = (time.time() - start_time) / num_iters
+        print("%s %s %s inputshape:%s filtershape:%s strides:%s padding:%s "
+              "%d iters: %.8f sec" %
+              (device, str(dtype), data_format, str(input_shape).replace(
+                  " ", ""), str(filter_shape).replace(" ", ""),
+               str(strides).replace(" ", ""), padding, num_iters, duration))
+
+    name_template = (
+        "conv2d_{device}_{datatype}_{data_format}_input_shape_{inputshape}_"
+        "filter_shape_{filtershape}_strides_{strides}_padding_{padding}")
+
+    self.report_benchmark(
+        name=name_template.format(
+            device=device,
+            datatype=str(dtype),
+            data_format=str(data_format),
+            inputshape=str(input_shape).replace(" ", ""),
+            filtershape=str(filter_shape).replace(" ", ""),
+            strides=str(strides).replace(" ", ""),
+            padding=padding).replace(" ", ""),
+        iters=num_iters,
+        wall_time=duration)
+
+    return duration
+
+  def benchmark_conv2d(self):
+    print("conv2d benchmark:")
+
+    data_types = [dtypes.float32, dtypes.float16]
+    data_formats = ["NHWC", "NCHW"]
+    in_channels = list(range(1, 10)) + list(range(10, 20, 2)) + list(
+        range(20, 33, 4))
+    out_channels = [4, 16, 32]
+    hw_strides = [[2, 2]]
+    paddings = ["VALID", "SAME"]
+
+    args_lists = [
+        data_types, data_formats, in_channels, out_channels, hw_strides,
+        paddings
+    ]
+    for args in itertools.product(*args_lists):
+      dtype, data_format, in_channel, out_channel, hw_stride, padding = args
+
+      # Keep batch size same as out channels just to reduce the number of
+      # different configurations to benchmark.
+      batch_size = out_channel
+      h, w, fh, fw = 500, 500, 3, 3
+      if data_format == "NHWC":
+        ishape = [batch_size, h, w, in_channel]
+        stride = [1] + hw_stride + [1]
+      elif data_format == "NCHW":
+        ishape = [batch_size, in_channel, h, w]
+        stride = [1, 1] + hw_stride
+      else:
+        raise ValueError("Unknown data_format: " + str(data_format))
+      fshape = [fh, fw, in_channel, out_channel]
+      num_iters = 80
+      warmup_iters = 2
+      self._run_graph("gpu", dtype, data_format, ishape, fshape, stride,
+                      padding, num_iters, warmup_iters)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/ops/critical_section_ops.py

@@ -0,0 +1,419 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Critical Section object and execution logic."""
+
+import collections
+import contextlib
+import threading
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import gen_resource_variable_ops
+from tensorflow.python.ops import tensor_array_ops
+from tensorflow.python.util import nest
+from tensorflow.python.util import object_identity
+from tensorflow.python.util.tf_export import tf_export
+
+
+__all__ = ["CriticalSection"]
+
+
+# Graph Keys
+CRITICAL_SECTIONS = "critical_sections"
+CRITICAL_SECTION_EXECUTIONS = "critical_section_executions"
+
+
+class _ExecutionSignature(
+    collections.namedtuple("_ExecutionSignature",
+                           ("op", "handle",
+                            "resources", "exclusive_resource_access"))):
+  """A class storing an `ExecuteInCriticalResource` op and associated attrs."""
+  pass
+
+
+def _identity(x):
+  """Identity op that recognizes `TensorArray`, `Operation`, and `Tensor`."""
+  if isinstance(x, tensor_array_ops.TensorArray):
+    return x.identity()
+  elif isinstance(x, ops.Operation):
+    return control_flow_ops.group(x)
+  elif context.executing_eagerly() and x is None:
+    return None
+  else:
+    return array_ops.identity(x)
+
+
+def _get_device_or_colocation(op):
+  return op.device or _get_colocation(op)
+
+
+def _get_colocation(op):
+  """Get colocation symbol from op, if any."""
+  try:
+    return op.get_attr("_class")
+  except (ValueError, AttributeError):
+    return None
+
+
+_CRITICAL_SECTION_STACK = threading.local()
+
+
+def _get_critical_section_stack():
+  try:
+    return _CRITICAL_SECTION_STACK.value
+  except AttributeError:
+    _CRITICAL_SECTION_STACK.value = []
+    return _CRITICAL_SECTION_STACK.value
+
+
+@contextlib.contextmanager
+def _push_critical_section_stack(signature):
+  """Push a CriticalSection._signature to the thread-local stack.
+
+  If the signature is already on the stack, raise an error because it means
+  we're trying to execute inside the same locked CriticalSection, which
+  will create a deadlock.
+
+  Args:
+    signature: Tuple of the type `CriticalSection._signature`.  Uniquely
+      identifies a CriticalSection by its `shared_name`, `container`,
+      and device.
+
+  Yields:
+    An empty value.  The context is guaranteed to run without deadlock.
+
+  Raises:
+    ValueError: If the signature is already on the stack.
+    RuntimeError: If another thread or function modifies the current stack
+      entry during the yield.
+  """
+  stack = _get_critical_section_stack()
+  if signature in stack:
+    raise ValueError(
+        f"Attempting to lock a CriticalSection (signature={signature}) in which"
+        " we are already running. This is illegal and may cause deadlocks.")
+  stack.append(signature)
+  try:
+    yield
+  finally:
+    received_signature = stack.pop()
+    if received_signature != signature:
+      raise RuntimeError(
+          "CriticalSection stack inconsistency: expected signature "
+          f"{signature} but received {received_signature}")
+
+
+@tf_export("CriticalSection")
+class CriticalSection:
+  """Critical section.
+
+  A `CriticalSection` object is a resource in the graph which executes subgraphs
+  in **serial** order.  A common example of a subgraph one may wish to run
+  exclusively is the one given by the following function:
+
+  ```python
+  v = resource_variable_ops.ResourceVariable(0.0, name="v")
+
+  def count():
+    value = v.read_value()
+    with tf.control_dependencies([value]):
+      with tf.control_dependencies([v.assign_add(1)]):
+        return tf.identity(value)
+  ```
+
+  Here, a snapshot of `v` is captured in `value`; and then `v` is updated.
+  The snapshot value is returned.
+
+  If multiple workers or threads all execute `count` in parallel, there is no
+  guarantee that access to the variable `v` is atomic at any point within
+  any thread's calculation of `count`.  In fact, even implementing an atomic
+  counter that guarantees that the user will see each value `0, 1, ...,` is
+  currently impossible.
+
+  The solution is to ensure any access to the underlying resource `v` is
+  only processed through a critical section:
+
+  ```python
+  cs = CriticalSection()
+  f1 = cs.execute(count)
+  f2 = cs.execute(count)
+  output = f1 + f2
+  session.run(output)
+  ```
+  The functions `f1` and `f2` will be executed serially, and updates to `v`
+  will be atomic.
+
+  **NOTES**
+
+  All resource objects, including the critical section and any captured
+  variables of functions executed on that critical section, will be
+  colocated to the same device (host and cpu/gpu).
+
+  When using multiple critical sections on the same resources, there is no
+  guarantee of exclusive access to those resources.  This behavior is disallowed
+  by default (but see the kwarg `exclusive_resource_access`).
+
+  For example, running the same function in two separate critical sections
+  will not ensure serial execution:
+
+  ```python
+  v = tf.compat.v1.get_variable("v", initializer=0.0, use_resource=True)
+  def accumulate(up):
+    x = v.read_value()
+    with tf.control_dependencies([x]):
+      with tf.control_dependencies([v.assign_add(up)]):
+        return tf.identity(x)
+  ex1 = CriticalSection().execute(
+    accumulate, 1.0, exclusive_resource_access=False)
+  ex2 = CriticalSection().execute(
+    accumulate, 1.0, exclusive_resource_access=False)
+  bad_sum = ex1 + ex2
+  sess.run(v.initializer)
+  sess.run(bad_sum)  # May return 0.0
+  ```
+  """
+
+  def __init__(self, name=None, shared_name=None,
+               critical_section_def=None, import_scope=None):
+    """Creates a critical section."""
+    context.ensure_initialized()
+    if critical_section_def and name is not None:
+      raise ValueError(f"Arguments critical_section_def={critical_section_def} "
+                       f"and shared_name={shared_name} are mutually exclusive. "
+                       "Please only specify one of them.")
+    if critical_section_def:
+      raise ValueError("Argument `critical_section_def` is not supported.")
+    else:
+      self._init_from_args(name, shared_name)
+
+  def _init_from_args(self, name, shared_name):  # pylint: disable=invalid-name
+    """Initialize the CriticalSection from constructor arguments."""
+    with ops.name_scope(name, "CriticalSection", []) as name:
+      with ops.init_scope():
+        # pylint: disable=protected-access
+        container = ops.get_default_graph()._container
+        # pylint: enable=protected-access
+        if shared_name is None:
+          shared_name = name
+        if container is None:
+          container = ""
+        self._handle = gen_resource_variable_ops.mutex_v2(
+            shared_name=shared_name, container=container, name=name)
+        # Get a uniquely identifying signature for the handle.
+        self._signature = (
+            container,
+            # If shared_name is empty, a unique CriticalSection is created.
+            shared_name or id(self._handle),
+            _get_device_or_colocation(self._handle))
+
+    if not context.executing_eagerly():
+      ops.add_to_collections(CRITICAL_SECTIONS, self)
+
+  @property
+  def name(self):
+    return self._handle.op.name
+
+  def execute(self, fn, exclusive_resource_access=True, name=None):
+    """Execute function `fn()` inside the critical section.
+
+    `fn` should not accept any arguments.  To add extra arguments to when
+    calling `fn` in the critical section, create a lambda:
+
+    ```python
+    critical_section.execute(lambda: fn(*my_args, **my_kwargs))
+    ```
+
+    Args:
+      fn: The function to execute.  Must return at least one tensor.
+      exclusive_resource_access: Whether the resources required by
+        `fn` should be exclusive to this `CriticalSection`.  Default: `True`.
+        You may want to set this to `False` if you will be accessing a
+        resource in read-only mode in two different CriticalSections.
+      name: The name to use when creating the execute operation.
+
+    Returns:
+      The tensors returned from `fn()`.
+
+    Raises:
+      ValueError: If `fn` attempts to lock this `CriticalSection` in any nested
+        or lazy way that may cause a deadlock.
+      ValueError: If `exclusive_resource_access == True` and
+        another `CriticalSection` has an execution requesting the same
+        resources as `fn``.  Note, even if `exclusive_resource_access` is
+        `True`, if another execution in another `CriticalSection` was created
+        without `exclusive_resource_access=True`, a `ValueError` will be raised.
+    """
+    with ops.name_scope(name, "critical_section_execute", []):
+      # Ensure that mutex locking only happens *after* all args and
+      # kwargs have been executed.  This avoids certain types of deadlocks.
+      with _push_critical_section_stack(self._signature):
+        lock = gen_resource_variable_ops.mutex_lock(self._handle)
+
+        if not context.executing_eagerly():
+          # NOTE(ebrevdo): This is to ensure we don't pick up spurious
+          # Operations created by other threads.
+          with ops.get_default_graph()._lock:  # pylint: disable=protected-access
+            existing_ops = ops.get_default_graph().get_operations()
+            with ops.control_dependencies([lock]):
+              r = fn()
+            # TODO(ebrevdo): If creating critical sections in a python loop,
+            # this makes graph creation time quadratic.  Revisit if this
+            # becomes a problem.
+            created_ops = (set(ops.get_default_graph().get_operations())
+                           .difference(existing_ops))
+        else:
+          with ops.control_dependencies([lock]):
+            r = fn()
+
+      if not context.executing_eagerly():
+        self._add_control_dependencies_to_lock(created_ops, lock.op)
+
+        # captured_resources is a list of resources that are directly
+        # accessed only by ops created during fn(), not by any
+        # ancestors of those ops in the graph.
+        captured_resources = object_identity.ObjectIdentitySet([
+            input_ for op in created_ops
+            for input_ in op.inputs
+            if input_.dtype == dtypes.resource
+        ])
+
+        # NOTE(ebrevdo): The only time self._is_self_handle() is True
+        # in this call is if one of the recently created ops, within
+        # the execute(), themselves attempt to access the
+        # CriticalSection.  This will cause a deadlock.
+        if any(self._is_self_handle(x) for x in captured_resources):
+          raise ValueError(
+              "Attempting to lock a CriticalSection in which we are "
+              f"already running (signature={self._signature}). This is illegal "
+              "and may cause deadlocks.")
+
+        self._check_multiple_access_to_resources(
+            captured_resources, exclusive_resource_access)
+
+      r_flat = [_identity(x) for x in nest.flatten(r)]
+
+      with ops.control_dependencies(r_flat):
+        # The identity must run on the same machine as self._handle
+        with ops.colocate_with(self._handle):
+          # Do not use array_ops.identity as there are special
+          # optimizations within TensorFlow which seem to elide it
+          # even when optimizations are disabled(!).
+          ensure_lock_exists = gen_resource_variable_ops.consume_mutex_lock(
+              lock)
+
+        # Make sure that if any element of r is accessed, all of
+        # them are executed together.
+        r = nest.pack_sequence_as(r, control_flow_ops.tuple(nest.flatten(r)))
+
+      with ops.control_dependencies([ensure_lock_exists]):
+        outputs = nest.map_structure(_identity, r)
+
+      if not context.executing_eagerly():
+        signature = _ExecutionSignature(
+            op=lock.op,
+            handle=self._handle,
+            resources=list(captured_resources),
+            exclusive_resource_access=exclusive_resource_access)
+        ops.add_to_collections(
+            CRITICAL_SECTION_EXECUTIONS, signature)
+
+      return outputs
+
+  def _add_control_dependencies_to_lock(self, created_ops, lock_op):
+    """To avoid deadlocks, all args must be executed before lock_op."""
+    # Get all arguments (explicit and captured) of all ops created by fn().
+    all_args = set([input_.op for op in created_ops for input_ in op.inputs])
+    all_args.update(
+        input_op for op in created_ops for input_op in op.control_inputs)
+    # Unfortunately, we can't use sets throughout because TF seems to
+    # create new Operation objects for the same op sometimes; and we
+    # can't rely on id(op).
+
+    # pylint: disable=protected-access
+    all_args_dict = dict((op._id, op) for op in all_args)
+
+    # Remove ops created within fn, or that lock_op already has a
+    # control dependency on.  Also remove a possible self-loop.
+    for op in created_ops:
+      all_args_dict.pop(op._id, None)
+    for op in lock_op.control_inputs:
+      all_args_dict.pop(op._id, None)
+    for input_ in lock_op.inputs:
+      all_args_dict.pop(input_.op._id, None)
+    all_args_dict.pop(lock_op._id, None)
+
+    all_args = all_args_dict.values()
+
+    if not all_args:
+      # No control dependencies to add; return early.
+      return
+
+    # This group is important: it ensures that any ops in all_args
+    # outside the control context of the lock_op (and this fn, which
+    # runs in the same context) are added to this context before
+    # being added to the control dependencies of lock_op.
+    all_args = control_flow_ops.group(*all_args)
+
+    lock_op._add_control_input(all_args)
+    # pylint: enable=protected-access
+
+  def _is_self_handle(self, x):
+    """Check if the tensor `x` is the same Mutex as `self._handle`."""
+    if isinstance(x, ops.EagerTensor):
+      return x is self._handle
+    return (x.op.type == "MutexV2"
+            # blank shared_name means the op will create a unique one.
+            and x.op.get_attr("shared_name")
+            and (x.op.get_attr("shared_name") ==
+                 self._handle.op.get_attr("shared_name"))
+            and (x.op.device == self._handle.op.device
+                 or _get_colocation(x.op) == _get_colocation(self._handle.op)))
+
+  def _check_multiple_access_to_resources(
+      self, captured_resources, exclusive_resource_access):
+    """Raise if captured_resources are accessed by another CriticalSection.
+
+    Args:
+      captured_resources: Set of tensors of type resource.
+      exclusive_resource_access: Whether this execution requires exclusive
+        resource access.
+
+    Raises:
+      ValueError: If any tensors in `captured_resources` are also accessed
+        by another `CriticalSection`, and at least one of them requires
+        exclusive resource access.
+    """
+    # Collections and op introspection does not work in eager
+    # mode.  This is generally ok; since eager mode (as of
+    # writing) executes sequentially anyway.
+    for sg in ops.get_collection(CRITICAL_SECTION_EXECUTIONS):
+      if self._is_self_handle(sg.handle):
+        # Other executions in the same critical section are allowed.
+        continue
+      if not (exclusive_resource_access or sg.exclusive_resource_access):
+        # Neither execution requested exclusive access.
+        continue
+      resource_intersection = captured_resources.intersection(sg.resources)
+      if resource_intersection:
+        raise ValueError(
+            "This execution would access resources: "
+            f"{list(resource_intersection)}. Either this lock "
+            f"(CriticalSection: {self._handle}) or lock '{sg}' "
+            f"(CriticalSection: {sg.handle}) requested exclusive resource "
+            "access of this resource. Did you mean to call execute with "
+            "keyword argument exclusive_resource_access=False?")

tensorflow/python/ops/cudnn_rnn_grad.py

@@ -0,0 +1,99 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Gradients for CuudnnRNN operators."""
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_cudnn_rnn_ops
+
+
+@ops.RegisterGradient("CudnnRNN")
+def _cudnn_rnn_backward(op: ops.Operation, *grads):
+  """Gradients for the CudnnRNN op."""
+  if not op.get_attr("is_training"):
+    raise ValueError(
+        "To use CudnnRNN in gradients, is_training must be set to True.")
+  return gen_cudnn_rnn_ops.cudnn_rnn_backprop(
+      input=op.inputs[0],
+      input_h=op.inputs[1],
+      input_c=op.inputs[2],
+      params=op.inputs[3],
+      output=op.outputs[0],
+      output_h=op.outputs[1],
+      output_c=op.outputs[2],
+      output_backprop=grads[0],
+      output_h_backprop=grads[1],
+      output_c_backprop=grads[2],
+      reserve_space=op.outputs[3],
+      dropout=op.get_attr("dropout"),
+      seed=op.get_attr("seed"),
+      seed2=op.get_attr("seed2"),
+      rnn_mode=op.get_attr("rnn_mode"),
+      input_mode=op.get_attr("input_mode"),
+      direction=op.get_attr("direction"))
+
+
+@ops.RegisterGradient("CudnnRNNV2")
+def _cudnn_rnn_backward_v2(op: ops.Operation, *grad):
+  if not op.get_attr("is_training"):
+    raise ValueError(
+        "To use CudnnRNNV2 in gradients, is_training must be set to True.")
+  return gen_cudnn_rnn_ops.cudnn_rnn_backprop_v2(
+      input=op.inputs[0],
+      input_h=op.inputs[1],
+      input_c=op.inputs[2],
+      params=op.inputs[3],
+      output=op.outputs[0],
+      output_h=op.outputs[1],
+      output_c=op.outputs[2],
+      output_backprop=grad[0],
+      output_h_backprop=grad[1],
+      output_c_backprop=grad[2],
+      reserve_space=op.outputs[3],
+      host_reserved=op.outputs[4],
+      dropout=op.get_attr("dropout"),
+      seed=op.get_attr("seed"),
+      seed2=op.get_attr("seed2"),
+      rnn_mode=op.get_attr("rnn_mode"),
+      input_mode=op.get_attr("input_mode"),
+      direction=op.get_attr("direction"))
+
+
+@ops.RegisterGradient("CudnnRNNV3")
+def _cudnn_rnn_backwardv3(op: ops.Operation, *grads):
+  """Gradients for the CudnnRNNV3 op."""
+  if not op.get_attr("is_training"):
+    raise ValueError(
+        "To use CudnnRNNV3 in gradients, is_training must be set to True.")
+  return gen_cudnn_rnn_ops.cudnn_rnn_backprop_v3(
+      input=op.inputs[0],
+      input_h=op.inputs[1],
+      input_c=op.inputs[2],
+      params=op.inputs[3],
+      sequence_lengths=op.inputs[4],
+      output=op.outputs[0],
+      output_h=op.outputs[1],
+      output_c=op.outputs[2],
+      output_backprop=grads[0],
+      output_h_backprop=grads[1],
+      output_c_backprop=grads[2],
+      reserve_space=op.outputs[3],
+      host_reserved=op.outputs[4],
+      dropout=op.get_attr("dropout"),
+      seed=op.get_attr("seed"),
+      seed2=op.get_attr("seed2"),
+      time_major=op.get_attr("time_major"),
+      num_proj=op.get_attr("num_proj"),
+      rnn_mode=op.get_attr("rnn_mode"),
+      input_mode=op.get_attr("input_mode"),
+      direction=op.get_attr("direction")) + (None,)

tensorflow/python/ops/custom_gradient.py

@@ -0,0 +1,823 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Decorator to overrides the gradient for a function."""
+
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import context
+from tensorflow.python.eager import record
+from tensorflow.python.framework import composite_tensor_gradient
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.ops import handle_data_util
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import op_selector
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops.unconnected_gradients import UnconnectedGradients
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util import nest
+from tensorflow.python.util import tf_decorator
+from tensorflow.python.util import tf_inspect
+from tensorflow.python.util import variable_utils
+from tensorflow.python.util.tf_export import tf_export
+
+
+VAR_OP_TYPES = [
+    "VariableV2",
+    "VarHandleOp",
+]
+
+
+@tf_export("custom_gradient")
+def custom_gradient(f=None):
+  """Decorator to define a function with a custom gradient.
+
+  This decorator allows fine grained control over the gradients of a sequence
+  for operations.  This may be useful for multiple reasons, including providing
+  a more efficient or numerically stable gradient for a sequence of operations.
+
+  For example, consider the following function that commonly occurs in the
+  computation of cross entropy and log likelihoods:
+
+  ```python
+  def log1pexp(x):
+    return tf.math.log(1 + tf.exp(x))
+  ```
+
+  Due to numerical instability, the gradient of this function evaluated at x=100
+  is NaN.  For example:
+
+  ```python
+  with tf.GradientTape() as tape:
+    tape.watch(x)
+    y=log1pexp(x)
+  dy_dx = tape.gradient(y, x) # Will be NaN when evaluated.
+  ```
+
+  The gradient expression can be analytically simplified to provide numerical
+  stability:
+
+  ```python
+  @tf.custom_gradient
+  def log1pexp(x):
+    e = tf.exp(x)
+    def grad(upstream):
+      return upstream * (1 - 1 / (1 + e))
+    return tf.math.log(1 + e), grad
+  ```
+
+  With this definition, the gradient `dy_dx` at `x = 100` will be correctly
+  evaluated as 1.0.
+
+  The variable `upstream` is defined as the upstream gradient. i.e. the gradient
+  from all the layers or functions originating from this layer. The above
+  example has no upstream functions, therefore `upstream = dy/dy = 1.0`.
+
+  Assume that `x_i` is `log1pexp` in the forward pass `x_1 = x_1(x_0)`,
+  `x_2 = x_2(x_1)`, ..., `x_i = x_i(x_i-1)`, ..., `x_n = x_n(x_n-1)`. By
+  chain rule we know that `dx_n/dx_0 = dx_n/dx_n-1 * dx_n-1/dx_n-2 * ... *
+  dx_i/dx_i-1 * ... * dx_1/dx_0`.
+
+  In this case the gradient of our current function defined as
+  `dx_i/dx_i-1 = (exp(x_i) / (1 + exp(x_i))) = (1 - 1 / (1 + exp(x_i)))`. The
+  upstream gradient `upstream` would be `dx_n/dx_n-1 * dx_n-1/dx_n-2 * ... *
+  dx_i+1/dx_i`. The upstream gradient multiplied by the current gradient is
+  then passed downstream.
+
+  In case the function takes multiple variables as input, the `grad`
+  function must also return  the same number of variables.
+  We take the function `z = x * y` as an example.
+
+  >>> @tf.custom_gradient
+  ... def bar(x, y):
+  ...   def grad(upstream):
+  ...     dz_dx = y
+  ...     dz_dy = x
+  ...     return upstream * dz_dx, upstream * dz_dy
+  ...   z = x * y
+  ...   return z, grad
+  >>> x = tf.constant(2.0, dtype=tf.float32)
+  >>> y = tf.constant(3.0, dtype=tf.float32)
+  >>> with tf.GradientTape(persistent=True) as tape:
+  ...   tape.watch(x)
+  ...   tape.watch(y)
+  ...   z = bar(x, y)
+  >>> z
+  <tf.Tensor: shape=(), dtype=float32, numpy=6.0>
+  >>> tape.gradient(z, x)
+  <tf.Tensor: shape=(), dtype=float32, numpy=3.0>
+  >>> tape.gradient(z, y)
+  <tf.Tensor: shape=(), dtype=float32, numpy=2.0>
+
+  Nesting custom gradients can lead to unintuitive results. The default
+  behavior does not correspond to n-th order derivatives. For example
+
+  ```python
+  @tf.custom_gradient
+  def op(x):
+    y = op1(x)
+    @tf.custom_gradient
+    def grad_fn(dy):
+      gdy = op2(x, y, dy)
+      def grad_grad_fn(ddy):  # Not the 2nd order gradient of op w.r.t. x.
+        return op3(x, y, dy, ddy)
+      return gdy, grad_grad_fn
+    return y, grad_fn
+  ```
+
+  The function `grad_grad_fn` will be calculating the first order gradient
+  of `grad_fn` with respect to `dy`, which is used to generate forward-mode
+  gradient graphs from backward-mode gradient graphs, but is not the same as
+  the second order gradient of `op` with respect to `x`.
+
+  Instead, wrap nested `@tf.custom_gradients` in another function:
+
+  ```python
+  @tf.custom_gradient
+  def op_with_fused_backprop(x):
+    y, x_grad = fused_op(x)
+    def first_order_gradient(dy):
+      @tf.custom_gradient
+      def first_order_custom(unused_x):
+        def second_order_and_transpose(ddy):
+          return second_order_for_x(...), gradient_wrt_dy(...)
+        return x_grad, second_order_and_transpose
+      return dy * first_order_custom(x)
+    return y, first_order_gradient
+  ```
+
+  Additional arguments to the inner `@tf.custom_gradient`-decorated function
+  control the expected return values of the innermost function.
+
+  The examples above illustrate how to specify custom gradients for functions
+  which do not read from variables. The following example uses variables, which
+  require special handling because they are effectively inputs of the forward
+  function.
+
+  >>> weights = tf.Variable(tf.ones([2]))  # Trainable variable weights
+  >>> @tf.custom_gradient
+  ... def linear_poly(x):
+  ...   # Creating polynomial
+  ...   poly = weights[1] * x + weights[0]
+  ...
+  ...   def grad_fn(dpoly, variables):
+  ...     # dy/dx = weights[1] and we need to left multiply dpoly
+  ...     grad_xs = dpoly * weights[1]  # Scalar gradient
+  ...
+  ...     grad_vars = []  # To store gradients of passed variables
+  ...     assert variables is not None
+  ...     assert len(variables) == 1
+  ...     assert variables[0] is weights
+  ...     # Manually computing dy/dweights
+  ...     dy_dw = dpoly * tf.stack([x ** 1, x ** 0])
+  ...     grad_vars.append(
+  ...         tf.reduce_sum(tf.reshape(dy_dw, [2, -1]), axis=1)
+  ...     )
+  ...     return grad_xs, grad_vars
+  ...   return poly, grad_fn
+  >>> x = tf.constant([1., 2., 3.])
+  >>> with tf.GradientTape(persistent=True) as tape:
+  ...   tape.watch(x)
+  ...   poly = linear_poly(x)
+  >>> poly # poly = x + 1
+  <tf.Tensor: shape=(3,),
+    dtype=float32,
+    numpy=array([2., 3., 4.], dtype=float32)>
+  >>> tape.gradient(poly, x)  # conventional scalar gradient dy/dx
+  <tf.Tensor: shape=(3,),
+    dtype=float32,
+    numpy=array([1., 1., 1.], dtype=float32)>
+  >>> tape.gradient(poly, weights)
+  <tf.Tensor: shape=(2,), dtype=float32, numpy=array([6., 3.], dtype=float32)>
+
+  Above example illustrates usage of trainable variable `weights`.
+  In the example, the inner `grad_fn` accepts an extra `variables` input
+  parameter and also returns an extra `grad_vars` output. That extra argument
+  is passed if the forward function reads any variables. You need to
+  compute the gradient w.r.t. each of those `variables` and output it as a list
+  of `grad_vars`. Note here that default value of `variables` is set to `None`
+  when no variables are used in the forward function.
+
+  It should be noted `tf.GradientTape` is still watching the forward pass of a
+  `tf.custom_gradient`, and will use the ops it watches. As a consequence,
+  calling `tf.function` while the tape is still watching leads
+  to a gradient graph being built. If an op is used in `tf.function` without
+  registered gradient, a `LookupError` will be raised.
+
+  Users can insert `tf.stop_gradient` to customize this behavior. This
+  is demonstrated in the example below. `tf.random.shuffle` does not have a
+  registered gradient. As a result `tf.stop_gradient` is used to avoid the
+  `LookupError`.
+
+  ```python
+  x = tf.constant([0.3, 0.5], dtype=tf.float32)
+
+  @tf.custom_gradient
+  def test_func_with_stop_grad(x):
+    @tf.function
+    def _inner_func():
+      # Avoid exception during the forward pass
+      return tf.stop_gradient(tf.random.shuffle(x))
+      # return tf.random.shuffle(x)  # This will raise
+
+    res = _inner_func()
+    def grad(upstream):
+      return upstream  # Arbitrarily defined custom gradient
+    return res, grad
+
+  with tf.GradientTape() as g:
+    g.watch(x)
+    res = test_func_with_stop_grad(x)
+
+  g.gradient(res, x)
+  ```
+
+  See also `tf.RegisterGradient` which registers a gradient function for a
+  primitive TensorFlow operation. `tf.custom_gradient` on the other hand allows
+  for fine grained control over the gradient computation of a sequence of
+  operations.
+
+  Note that if the decorated function uses `Variable`s, the enclosing variable
+  scope must be using
+  [ResourceVariables](https://www.tensorflow.org/guide/migrate/tf1_vs_tf2#resourcevariables_instead_of_referencevariables).
+
+  Args:
+    f: function `f(*x)` that returns a tuple `(y, grad_fn)` where: - `x` is a
+      sequence of (nested structures of) `Tensor` inputs to the function. - `y`
+      is a (nested structure of) `Tensor` outputs of applying TensorFlow
+      operations in `f` to `x`. - `grad_fn` is a function with the signature
+      `g(*grad_ys)` which returns a list of `Tensor`s the same size as
+      (flattened) `x` - the derivatives of `Tensor`s in `y` with respect to the
+      `Tensor`s in `x`.  `grad_ys` is a sequence of `Tensor`s the same size as
+      (flattened) `y` holding the initial value gradients for each `Tensor` in
+      `y`.  In a pure mathematical sense, a vector-argument vector-valued
+      function `f`'s derivatives should be its Jacobian matrix `J`. Here we are
+      expressing the Jacobian `J` as a function `grad_fn` which defines how `J`
+      will transform a vector `grad_ys` when left-multiplied with it (`grad_ys *
+      J`, the vector-Jacobian product, or VJP). This functional representation
+      of a matrix is convenient to use for chain-rule calculation (in e.g. the
+      back-propagation algorithm).  If `f` uses `Variable`s (that are not part
+      of the inputs), i.e. through `get_variable`, then `grad_fn` should have
+      signature `g(*grad_ys, variables=None)`, where `variables` is a list of
+      the `Variable`s, and return a 2-tuple `(grad_xs, grad_vars)`, where
+      `grad_xs` is the same as above, and `grad_vars` is a `list<Tensor>` with
+      the derivatives of `Tensor`s in `y` with respect to the variables (that
+      is, grad_vars has one Tensor per variable in variables).
+
+  Returns:
+    A function `h(x)` which returns the same value as `f(x)[0]` and whose
+    gradient (as calculated by `tf.gradients`) is determined by `f(x)[1]`.
+  """
+
+  if f is None:
+    return lambda f: custom_gradient(f=f)
+
+  @Bind.decorator
+  def decorated(wrapped, args, kwargs):
+    """Decorated function with custom gradient."""
+    if context.executing_eagerly():
+      return _eager_mode_decorator(wrapped, args, kwargs)
+    else:
+      return _graph_mode_decorator(wrapped, args, kwargs)
+
+  return tf_decorator.make_decorator(f, decorated(f))  # pylint: disable=no-value-for-parameter
+
+
+class Bind:
+  """When called evaluates `d(f, args, kwargs)` but supports binding `f`.
+
+  >>> @Bind.decorator
+  ... def my_decorator(f, args, kwargs):
+  ...   print("my_decorator called with", args, kwargs)
+  ...   return f(*args, **kwargs)
+
+  >>> class Foo:
+  ...   @my_decorator
+  ...   def bar(self, a, b, c):
+  ...     return a * b * c
+
+  >>> Foo.bar(None, 1, 2, c=3)
+  my_decorator called with (None, 1, 2) {'c': 3}
+  6
+
+  >>> foo = Foo()
+  >>> foo.bar(1, 2, c=3)
+  my_decorator called with (1, 2) {'c': 3}
+  6
+  """
+
+  @classmethod
+  def decorator(cls, d):
+    return lambda f: Bind(f, d)
+
+  def __init__(self, f, d):
+    self._f = f
+    self._d = d
+
+  def __get__(self, instance, owner):
+    if instance is not None:
+      f = self._f.__get__(instance, owner)
+      return tf_decorator.make_decorator(f, Bind(f, self._d))
+    else:
+      return self
+
+  def __call__(self, *a, **k):
+    return self._d(self._f, a, k)
+
+
+def get_variable_by_name(var_name):
+  """Given a variable name, retrieves a handle on the tensorflow Variable."""
+  global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
+
+  def _filter_fn(item):
+    try:
+      return var_name == item.op.name
+    except AttributeError:
+      # Collection items without operation are ignored.
+      return False
+
+  candidate_vars = list(filter(_filter_fn, global_vars))
+
+  if len(candidate_vars) >= 1:
+    # Filter out non-trainable variables.
+    candidate_vars = [v for v in candidate_vars if v.trainable]
+  else:
+    raise ValueError("Unsuccessful at finding variable {}.".format(var_name))
+
+  if len(candidate_vars) == 1:
+    return candidate_vars[0]
+  elif len(candidate_vars) > 1:
+    raise ValueError(
+        "Unsuccessful at finding trainable variable {}. "
+        "Number of candidates: {}. "
+        "Candidates: {}".format(var_name, len(candidate_vars), candidate_vars))
+  else:
+    # The variable is not trainable.
+    return None
+
+
+def _get_dependent_variables(input_ops, output_ops):
+  """Finds variables involved in the subgraph between input_ops and output_ops.
+
+  Args:
+    input_ops: Flattened list of input ops
+    output_ops: Flattened list of output ops
+
+  Returns:
+    A list of variables
+  """
+
+  # avoids the edge-case when input_ops == output_ops.
+  output_ops = nest.map_structure(gen_array_ops.identity, output_ops)
+  inbetween_ops = op_selector.get_backward_walk_ops(
+      seed_ops=output_ops,
+      stop_at_ts=input_ops,
+      inclusive=False,
+      only_differentiable=True)
+  var_ops = (op for op in inbetween_ops if op.type in VAR_OP_TYPES)
+  var_names = (op.name for op in var_ops)
+  tf_vars = (get_variable_by_name(var_name) for var_name in var_names)
+  tf_vars = [v for v in tf_vars if v is not None]
+  return tf_vars
+
+
+def generate_name():
+  return "CustomGradient-%s" % ops.uid()
+
+
+def _graph_mode_decorator(f, args, kwargs):
+  """Implement custom gradient decorator for graph mode."""
+  # TODO(rsepassi): Add support for kwargs
+  if kwargs:
+    raise ValueError(
+        "The custom_gradient decorator currently supports keywords "
+        "arguments only when eager execution is enabled.")
+  name = generate_name()
+  args = variable_utils.convert_variables_to_tensors(args)
+  args = nest.map_structure(ops.convert_to_tensor, args, expand_composites=True)
+
+  # Checking global and local variables attempts to ensure that no non-resource
+  # Variables are added to the graph.
+  current_var_scope = variable_scope.get_variable_scope()
+  before_vars = set([
+      v.ref() for v in current_var_scope.global_variables() +
+      current_var_scope.local_variables()
+  ])
+  with record.VariableWatcher() as variable_watcher:
+    result, grad_fn = f(*args)
+
+  flat_args = composite_tensor_gradient.get_flat_tensors_for_gradients(
+      nest.flatten(args))
+  flat_result = composite_tensor_gradient.get_flat_tensors_for_gradients(
+      nest.flatten(result))
+  flat_result_len = len(flat_result)
+
+  after_vars = set([
+      v.ref() for v in current_var_scope.global_variables() +
+      current_var_scope.local_variables()
+  ])
+  new_vars = after_vars - before_vars
+  new_vars_list = [v.deref() for v in new_vars]
+  for v in new_vars_list:
+    if not resource_variable_ops.is_resource_variable(v):
+      raise TypeError(
+          "All variables used by a function wrapped with @custom_gradient must "
+          "be `ResourceVariable`s. Ensure that no `variable_scope` is created "
+          "with `use_resource=False`.")
+
+  # The variables that grad_fn needs to return gradients for are the set of
+  # variables used that are *not* part of the inputs.
+  variables_in_tape = frozenset([
+      v.ref() for v in variable_watcher.watched_variables()
+  ])
+
+  graphs = {getattr(o, "graph", None) for o in flat_result}
+  # Not all results may be tensors. However, we want to ensure all tensor
+  # outputs are from the same graph and get a list of captured inputs for
+  # variable search
+  graphs.discard(None)  # Discard non-graph outputs
+  if graphs:
+    if len(graphs) > 1:
+      raise ValueError(
+          "All custom_gradient outputs should be from the same graph")
+    output_graph = graphs.pop()
+    filtered_input_tensors = []
+    for i in flat_args:
+      if i.graph == output_graph:
+        filtered_input_tensors.append(i)
+  else:
+    filtered_input_tensors = flat_args
+
+  variables_in_subgraph = frozenset([
+      v.ref() for v in _get_dependent_variables(
+          input_ops=filtered_input_tensors, output_ops=flat_result)
+  ])
+  variables = sorted(
+      [v.deref() for v in variables_in_subgraph.union(variables_in_tape)],
+      key=lambda v: v.name)
+
+  grad_argspec = tf_inspect.getfullargspec(grad_fn)
+  variables_in_signature = ("variables" in grad_argspec.args or
+                            "variables" in grad_argspec.kwonlyargs or
+                            grad_argspec.varkw)
+  if variables and not variables_in_signature:
+    raise TypeError(
+        "@tf.custom_gradient grad_fn must accept keyword argument 'variables', "
+        "since function uses variables: {}".format(variables))
+  if variables_in_signature and not variables:
+    # User seems to intend to use variables but none were captured.
+    logging.vlog(
+        1, "@custom_gradient grad_fn has 'variables' in signature, "
+        "but no ResourceVariables were used on the forward pass.")
+
+  all_tensors = flat_result + flat_args + variables
+
+  def tape_grad_fn(*result_grad_components):
+    """Custom grad fn wrapper."""
+    result_grads = composite_tensor_gradient.replace_flat_tensors_for_gradients(
+        nest.flatten(result), result_grad_components[:flat_result_len])
+    if not isinstance(result_grads, (list, tuple)):
+      result_grads = [result_grads]
+
+    if variables:
+      input_grads, variable_grads = grad_fn(*result_grads, variables=variables)
+      if len(variable_grads) != len(variables):
+        raise ValueError("Must return gradient for each variable from "
+                         "@custom_gradient grad_fn.")
+    else:
+      input_grads = grad_fn(*result_grads)
+      variable_grads = []
+
+    # Need to return one value per input to the IdentityN, so pad the
+    # gradients of the inputs of the custom_gradient function with the
+    # gradients of the outputs as well.
+    input_grads = composite_tensor_gradient.get_flat_tensors_for_gradients(
+        nest.flatten(input_grads))
+    return ([None] * flat_result_len) + input_grads + variable_grads
+
+  @ops.RegisterGradient(name)
+  def internal_grad_fn(unused_op, *result_grads):  # pylint: disable=unused-variable
+    """Custom grad fn wrapper."""
+    return tape_grad_fn(*result_grads)
+
+  original_tensors = all_tensors
+  with ops.get_default_graph().gradient_override_map({"IdentityN": name}):
+    all_tensors = array_ops.identity_n(all_tensors)
+
+  original_tensors = [ops.convert_to_tensor(x) for x in original_tensors]
+
+  # Propagate handle data for happier shape inference for resource variables.
+  for i, t in enumerate(original_tensors):
+    if t.dtype == dtypes.resource and hasattr(t, "_handle_data"):
+      all_tensors[i]._handle_data = t._handle_data  # pylint: disable=protected-access
+  record.record_operation(
+      f.__name__, all_tensors, original_tensors, tape_grad_fn)
+  for ot, t in zip(original_tensors, all_tensors):
+    handle_data_util.copy_handle_data(ot, t)
+  flat_result = composite_tensor_gradient.replace_flat_tensors_for_gradients(
+      nest.flatten(result), all_tensors[:flat_result_len])
+  return nest.pack_sequence_as(result, flat_result)
+
+
+def _eager_mode_decorator(f, args, kwargs):
+  """Implement custom gradient decorator for eager mode."""
+  with record.VariableWatcher() as variable_watcher:
+    result, grad_fn = f(*args, **kwargs)
+  flat_args = composite_tensor_gradient.get_flat_tensors_for_gradients(
+      nest.flatten(args))
+  flat_kwargs = composite_tensor_gradient.get_flat_tensors_for_gradients(
+      nest.flatten(kwargs))
+  all_inputs = flat_args + flat_kwargs
+  # The variables that grad_fn needs to return gradients for are the set of
+  # variables used that are *not* part of the inputs.
+  variables = [
+      v.deref()  # pylint: disable=g-complex-comprehension
+      for v in set(v.ref() for v in variable_watcher.watched_variables())
+      if all(v.deref() is not i for i in all_inputs)
+  ]
+  grad_argspec = tf_inspect.getfullargspec(grad_fn)
+  if (variables and ("variables" not in grad_argspec.args) and
+      ("variables" not in grad_argspec.kwonlyargs) and
+      not grad_argspec.varkw):
+    raise TypeError(
+        "@tf.custom_gradient grad_fn must accept keyword argument 'variables', "
+        "since function uses variables: {}".format(variables))
+  flat_result = composite_tensor_gradient.get_flat_tensors_for_gradients(
+      nest.flatten(result))
+  # TODO(apassos) consider removing the identity below.
+  flat_result = [gen_array_ops.identity(x) for x in flat_result]
+
+  input_tensors = [
+      ops.convert_to_tensor(x) for x in flat_args + list(variables)]
+
+  recorded_inputs = input_tensors
+  arg_count = len(flat_args)
+
+  def actual_grad_fn(*result_grad_components):
+    """Custom grad fn wrapper."""
+    result_grads = composite_tensor_gradient.replace_flat_tensors_for_gradients(
+        nest.flatten(result), result_grad_components)
+    if not isinstance(result_grads, (list, tuple)):
+      result_grads = [result_grads]
+
+    if variables:
+      input_grads, variable_grads = grad_fn(*result_grads, variables=variables)
+      if len(variable_grads) != len(variables):
+        raise ValueError("Must return gradient for each variable from "
+                         "@custom_gradient grad_fn.")
+    else:
+      input_grads = grad_fn(*result_grads)
+      variable_grads = []
+    flat_grads = composite_tensor_gradient.get_flat_tensors_for_gradients(
+        nest.flatten(input_grads))
+    if len(flat_grads) != arg_count:
+      raise ValueError(
+          f"custom_gradient function expected to return {arg_count} "
+          f"gradients, but returned {len(flat_grads)} instead.")
+    return flat_grads + variable_grads
+
+  record.record_operation(f.__name__, flat_result, recorded_inputs,
+                          actual_grad_fn)
+  flat_result = composite_tensor_gradient.replace_flat_tensors_for_gradients(
+      nest.flatten(result), flat_result)
+  return nest.pack_sequence_as(result, flat_result)
+
+
+@tf_export("recompute_grad")
+def recompute_grad(f):
+  """Defines a function as a recompute-checkpoint for the tape auto-diff.
+
+  Tape checkpointing is a technique to reduce the memory consumption of the
+  auto-diff tape:
+
+  - Without tape checkpointing operations and intermediate values are
+  recorded to the tape for use in the backward pass.
+
+  - With tape checkpointing, only the function call and its inputs are
+  recorded. During back-propagation the `recompute_grad` custom gradient
+  (`tf.custom_gradient`) recomputes the function under a localized Tape object.
+  This recomputation of the function during backpropagation performs redundant
+  calculation, but reduces the overall memory usage of the Tape.
+
+  >>> y = tf.Variable(1.0)
+
+  >>> def my_function(x):
+  ...   tf.print('running')
+  ...   z = x*y
+  ...   return z
+
+  >>> my_function_recompute = tf.recompute_grad(my_function)
+
+  >>> with tf.GradientTape() as tape:
+  ...   r = tf.constant(1.0)
+  ...   for i in range(4):
+  ...     r = my_function_recompute(r)
+  running
+  running
+  running
+  running
+
+  >>> grad = tape.gradient(r, [y])
+  running
+  running
+  running
+  running
+
+  Without `recompute_grad`, the tape contains all intermitate steps, and no
+  recomputation is performed.
+
+  >>> with tf.GradientTape() as tape:
+  ...   r = tf.constant(1.0)
+  ...   for i in range(4):
+  ...     r = my_function(r)
+  running
+  running
+  running
+  running
+
+  >>> grad = tape.gradient(r, [y])
+
+
+  If `f` was a `tf.keras` `Model` or `Layer` object, methods and attributes
+  such as `f.variables` are not available on the returned function `g`.
+  Either keep a reference of `f` , or use `g.__wrapped__` for accessing
+  these variables and methods.
+
+
+  >>> def print_running_and_return(x):
+  ...   tf.print("running")
+  ...   return x
+
+  >>> model = tf.keras.Sequential([
+  ...   tf.keras.layers.Lambda(print_running_and_return),
+  ...   tf.keras.layers.Dense(2)
+  ... ])
+
+  >>> model_recompute = tf.recompute_grad(model)
+
+  >>> with tf.GradientTape(persistent=True) as tape:
+  ...   r = tf.constant([[1,2]])
+  ...   for i in range(4):
+  ...     r = model_recompute(r)
+  running
+  running
+  running
+  running
+
+  >>> grad = tape.gradient(r, model.variables)
+  running
+  running
+  running
+  running
+
+  Alternatively, use the `__wrapped__` attribute to access the original
+  model object.
+
+  >>> grad = tape.gradient(r, model_recompute.__wrapped__.variables)
+  running
+  running
+  running
+  running
+
+
+  Args:
+    f: function `f(*x)` that returns a `Tensor` or sequence of `Tensor` outputs.
+
+  Returns:
+    A function `g` wrapping `f` that defines a custom gradient, which recomputes
+    `f` on the backwards pass of a gradient call.
+  """
+  # TODO(cdfreeman) Add is_recomputing functionality from graph mode version
+
+  @custom_gradient
+  def inner(*args, **kwargs):
+    """Inner function closure for calculating gradients."""
+    current_var_scope = variable_scope.get_variable_scope()
+    with record.stop_recording():
+      result = f(*args, **kwargs)
+
+    def grad_wrapper(*wrapper_args, variables=None):
+      """Wrapper function to accommodate lack of kwargs in graph mode custom_gradient."""
+
+      @custom_gradient
+      def inner_recompute_grad(*dresult):
+        """Nested custom gradient function for computing grads in reverse and forward mode autodiff."""
+        # Gradient calculation for reverse mode autodiff.
+        with backprop.GradientTape() as t:
+          id_args = nest.map_structure(gen_array_ops.identity, args)
+          # Tuple `dresult` should contain at least one tensor.
+          assert len(dresult) >= 1
+
+          if not context.executing_eagerly():
+            # XLA doesn't respect `tf.control_dependencies`. The code block
+            # below manually adds a data dependency to `dresult` to ensure
+            # recomputation of `f(*args, **kwargs)` happens after `dresult`.
+
+            # This works even if `dresult[0]` is a size 0 tensor as reduce_max
+            # of a size 0 tensor returns -inf. Use reshape here to avoid reading
+            # the entire `dresult[0]`.
+            elem = math_ops.reduce_max(array_ops.reshape(dresult[0], [-1])[:1])
+            # Cast elem to bool in case elem is NaN.
+            elem_bool = math_ops.cast(elem, dtypes.bool)
+            dresult_dep = array_ops.where_v2(
+                elem_bool == elem_bool, 0., float("nan"))  # pylint: disable=comparison-with-itself
+            id_args = nest.map_structure(
+                lambda x: x + math_ops.cast(dresult_dep, x.dtype), id_args)
+
+          t.watch(id_args)
+          if variables is not None:
+            t.watch(variables)
+          with variable_scope.variable_scope(current_var_scope):
+            recomputed_result = f(*id_args, **kwargs)
+        kw_vars = []
+        if variables is not None:
+          kw_vars = list(variables)
+        grads = t.gradient(
+            recomputed_result,
+            list(id_args) + kw_vars,
+            output_gradients=dresult,
+            unconnected_gradients=UnconnectedGradients.ZERO)
+
+        def transpose(*t_args, **t_kwargs):
+          """Gradient function calculation for forward mode autodiff."""
+          # Just throw an error since gradients / activations are not stored on
+          # tape for recompute.
+          raise NotImplementedError(
+              "recompute_grad tried to transpose grad of {}. "
+              "Consider not using recompute_grad in forward mode"
+              "autodiff".format(f.__name__))
+
+        return (grads[:len(id_args)], grads[len(id_args):]), transpose
+
+      return inner_recompute_grad(*wrapper_args)
+
+    return result, grad_wrapper
+
+  return tf_decorator.make_decorator(f, inner)
+
+
+@tf_export("grad_pass_through")
+def grad_pass_through(f):
+  """Creates a grad-pass-through op with the forward behavior provided in f.
+
+  Use this function to wrap any op, maintaining its behavior in the forward
+  pass, but replacing the original op in the backward graph with an identity.
+  For example:
+
+  ```python
+  x = tf.Variable(1.0, name="x")
+  z = tf.Variable(3.0, name="z")
+
+  with tf.GradientTape() as tape:
+    # y will evaluate to 9.0
+    y = tf.grad_pass_through(x.assign)(z**2)
+  # grads will evaluate to 6.0
+  grads = tape.gradient(y, z)
+  ```
+
+  Another example is a 'differentiable' moving average approximation, where
+  gradients are allowed to flow into the last value fed to the moving average,
+  but the moving average is still used for the forward pass:
+
+  ```python
+  x = ... # Some scalar value
+  # A moving average object, we don't need to know how this is implemented
+  moving_average = MovingAverage()
+  with backprop.GradientTape() as tape:
+    # mavg_x will evaluate to the current running average value
+    mavg_x = tf.grad_pass_through(moving_average)(x)
+  grads = tape.gradient(mavg_x, x) # grads will evaluate to 1.0
+  ```
+
+  Args:
+    f: function `f(*x)` that returns a `Tensor` or nested structure of `Tensor`
+      outputs.
+
+  Returns:
+    A function `h(x)` which returns the same values as `f(x)` and whose
+    gradients are the same as those of an identity function.
+  """
+  @custom_gradient
+  def _grad_pass_through_op(*args, **kwargs):
+    def grad(*args, **kwargs):
+      variables = kwargs.get("variables")
+      if variables is not None:
+        # Variables involved in the wrapped op will not receive gradients.
+        return args, [None] * len(variables)
+      return args
+    return f(*args, **kwargs), grad
+  return tf_decorator.make_decorator(f, _grad_pass_through_op)

tensorflow/python/ops/data_flow_grad.py

@@ -0,0 +1,115 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Gradients for operators defined in data_flow_ops.py."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import data_flow_ops
+from tensorflow.python.ops import math_ops
+
+
+@ops.RegisterGradient("DynamicPartition")
+def _DynamicPartitionGrads(op, *grads):
+  """Gradients for DynamicPartition."""
+  data = op.inputs[0]
+  indices = op.inputs[1]
+  num_partitions = op.get_attr("num_partitions")
+
+  prefix_shape = array_ops.shape(indices)
+  original_indices = array_ops.reshape(
+      math_ops.range(math_ops.reduce_prod(prefix_shape)), prefix_shape)
+  partitioned_indices = data_flow_ops.dynamic_partition(
+      original_indices, indices, num_partitions)
+  reconstructed = data_flow_ops.parallel_dynamic_stitch(partitioned_indices,
+                                                        grads)
+  reconstructed = array_ops.reshape(reconstructed, array_ops.shape(data))
+  return [reconstructed, None]
+
+
+@ops.RegisterGradient("DynamicStitch")
+@ops.RegisterGradient("ParallelDynamicStitch")
+def _DynamicStitchGrads(op, grad):
+  """Gradients for DynamicStitch and ParallelDynamicStitch."""
+
+  num_values = len(op.inputs) // 2
+  indices_grad = [None] * num_values
+
+  def AsInt32(x):
+    return (x if op.inputs[0].dtype == dtypes.int32 else
+            math_ops.cast(x, dtypes.int32))
+
+  inputs = [AsInt32(op.inputs[i]) for i in range(num_values)]
+  if isinstance(grad, indexed_slices.IndexedSlices):
+    output_shape = array_ops.shape(op.outputs[0])
+    output_rows = output_shape[0]
+    grad = math_ops.unsorted_segment_sum(grad.values, grad.indices, output_rows)
+
+  ids = []
+  current_size = array_ops.zeros([], dtype=dtypes.int32)
+  for inp in inputs:
+    num_elements = math_ops.cast(array_ops.size(inp), current_size.dtype)
+    flat_id = math_ops.range(current_size, current_size + num_elements)
+    ids.append(array_ops.reshape(flat_id, array_ops.shape(inp)))
+    current_size += num_elements
+
+  stitch_op = (
+      data_flow_ops.parallel_dynamic_stitch
+      if op.type == "ParallelDynamicStitch"
+      else data_flow_ops.dynamic_stitch
+  )
+
+  stitched_ids = stitch_op(inputs, ids)
+
+  values_grad = []
+  num_inner_dims = array_ops.rank(grad) - 1
+  for inp, single_id in zip(inputs, ids):
+    value_grad = array_ops.gather(grad, inp)
+    winning_ids = array_ops.gather(stitched_ids, inp)
+    is_winner = math_ops.equal(winning_ids, single_id)
+    mask = math_ops.cast(is_winner, value_grad.dtype)
+    winner_shape = array_ops.shape(is_winner)
+    mask_shape = array_ops.concat(
+        [
+            winner_shape,
+            array_ops.ones([num_inner_dims], dtype=winner_shape.dtype),
+        ],
+        axis=0,
+    )
+    values_grad.append(value_grad * array_ops.reshape(mask, mask_shape))
+
+  return indices_grad + values_grad
+
+
+ops.NotDifferentiable("Queue")
+ops.NotDifferentiable("QueueEnqueue")
+ops.NotDifferentiable("QueueEnqueueMany")
+ops.NotDifferentiable("QueueDequeue")
+ops.NotDifferentiable("QueueDequeueMany")
+ops.NotDifferentiable("QueueDequeueUpTo")
+ops.NotDifferentiable("QueueClose")
+ops.NotDifferentiable("QueueSize")
+
+ops.NotDifferentiable("Stack")
+ops.NotDifferentiable("StackPush")
+ops.NotDifferentiable("StackPop")
+ops.NotDifferentiable("StackClose")
+
+ops.NotDifferentiable("GetSessionHandle")
+ops.NotDifferentiable("GetSessionHandleV2")
+ops.NotDifferentiable("GetSessionTensor")
+ops.NotDifferentiable("DeleteSessionTensor")

tensorflow/python/ops/default_gradient.py

@@ -0,0 +1,80 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utilities for computing default gradients."""
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import resource_variable_ops
+
+
+def get_zeros_dtype(t):
+  """Return the dtype for the default gradient for a Tensor."""
+  if t.dtype == dtypes.resource:
+    handle_data = resource_variable_ops.get_eager_safe_handle_data(t)
+    if (handle_data is None or not handle_data.is_set or
+        len(handle_data.shape_and_type) != 1):
+      raise ValueError("Internal error: Tried to take gradients (or similar) "
+                       "of a variable without handle data:\n%s" % str(t))
+    return handle_data.shape_and_type[0].dtype
+  return t.dtype
+
+
+def shape_and_dtype(t):
+  """Return the shape and dtype for the default gradient for a Tensor."""
+  if t.dtype == dtypes.resource:
+    handle_data = resource_variable_ops.get_eager_safe_handle_data(t)
+    if (handle_data is None or not handle_data.is_set or
+        len(handle_data.shape_and_type) != 1):
+      raise ValueError("Internal error: Tried to take gradients (or similar) "
+                       "of a variable without handle data:\n%s" % str(t))
+    shape_and_type = handle_data.shape_and_type[0]
+    return (tensor_shape.TensorShape(shape_and_type.shape),
+            dtypes.as_dtype(shape_and_type.dtype))
+  return t.shape, t.dtype
+
+
+def zeros_like(t):
+  """Like array_ops.zeros_like, but respects resource handles."""
+  if t.dtype == dtypes.resource:
+    return array_ops.zeros(*shape_and_dtype(t))
+  else:
+    return array_ops.zeros_like(t)
+
+
+def ones_like(t):
+  """Like array_ops.ones_like, but respects resource handles."""
+  if t.dtype == dtypes.resource:
+    return array_ops.ones(*shape_and_dtype(t))
+  else:
+    return array_ops.ones_like(t)
+
+
+def supports_default_grad(t):
+  """Whether tensor `t` supports creating a default gradient.
+
+  This function assumes that `t` is of a trainable type.
+
+  Args:
+    t: Tensor
+
+  Returns:
+    Bool
+  """
+  if t.dtype == dtypes.resource:
+    handle_data = resource_variable_ops.get_eager_safe_handle_data(t)
+    if (handle_data is None or not handle_data.is_set or
+        len(handle_data.shape_and_type) != 1):
+      return False
+  return True

tensorflow/python/ops/dequantize_op_test.py

@@ -0,0 +1,134 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for Dequantize Operations."""
+
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+class DequantizeOpTest(test.TestCase):
+
+  def __init__(self, method_name="runTest"):
+    super(DequantizeOpTest, self).__init__(method_name)
+
+  def _testDequantizeOp(self, inputs, min_range, max_range, dtype,
+                        mode="MIN_COMBINED", narrow_range=False):
+    with self.cached_session():
+      input_op = constant_op.constant(inputs, shape=[len(inputs)], dtype=dtype)
+      dequantized = array_ops.dequantize(input_op, min_range, max_range,
+                                         mode=mode, narrow_range=narrow_range)
+      tf_ans = self.evaluate(dequantized)
+
+    # TODO(vrv): Add support for DT_QINT32 quantization if needed.
+    type_dict = {
+        dtypes.quint8: np.uint8,
+        dtypes.qint8: np.int8,
+        dtypes.quint16: np.uint16,
+        dtypes.qint16: np.int16
+    }
+    self.assertIn(dtype, type_dict.keys())
+    v_max = np.iinfo(type_dict[dtype]).max
+    v_min = np.iinfo(type_dict[dtype]).min
+    self.assertGreaterEqual(min_range, v_min)
+    self.assertLessEqual(max_range, v_max)
+    type_range = v_max - v_min
+
+    if mode == "MIN_COMBINED":
+      if v_min < 0:
+        half_range = (type_range + 1) / 2
+      else:
+        half_range = 0.0
+      np_ans = ((inputs.astype(np.float32) + half_range) *
+                (max_range - min_range) / type_range) + min_range
+    elif mode == "SCALED":
+      if narrow_range:
+        v_min += 1
+      scale_factor = max(min_range / v_min, max_range / v_max)
+      np_ans = inputs.astype(np.float32) * scale_factor
+
+    self.assertAllClose(tf_ans, np_ans, rtol=1e-5, atol=1e-5)
+
+  def testBasicQuint8(self):
+    self._testDequantizeOp(np.array([0, 128, 255]), 0.0, 6.0, dtypes.quint8)
+    self._testDequantizeOp(np.array([0, 128, 255]), 0.0, 123.456, dtypes.quint8)
+    self._testDequantizeOp(
+        np.array([0, 4, 42, 108, 243]), 5.0, 200.2, dtypes.quint8)
+
+  def testBasicQint8(self):
+    self._testDequantizeOp(np.array([-128, 0, 127]), -1.0, 2.0, dtypes.qint8)
+    self._testDequantizeOp(np.array([-2, 4, -17]), -5.0, -3.0, dtypes.qint8)
+    self._testDequantizeOp(np.array([0, -4, 42, -108]), 5.0, 40.0, dtypes.qint8)
+
+  def testScaledMode(self):
+    self._testDequantizeOp(np.array([-128, 0, 127]), -1.0, 2.0, dtypes.qint8,
+                           mode="SCALED")
+    self._testDequantizeOp(np.array([-2, 4, -17]), -5.0, -3.0, dtypes.qint8,
+                           mode="SCALED")
+    self._testDequantizeOp(np.array([0, -4, 42, -108]), 5.0, 40.0, dtypes.qint8,
+                           mode="SCALED")
+
+  def testNarrowRange(self):
+    self._testDequantizeOp(np.array([-128, 0, 127]), -1.0, 2.0, dtypes.qint8,
+                           mode="SCALED", narrow_range=True)
+    self._testDequantizeOp(np.array([-2, 4, -17]), -5.0, -3.0, dtypes.qint8,
+                           mode="SCALED", narrow_range=True)
+    self._testDequantizeOp(np.array([0, -4, 42, -108]), 5.0, 40.0, dtypes.qint8,
+                           mode="SCALED", narrow_range=True)
+
+  def testAxis(self):
+    # Generates a tensor of the specified `shape` using values from `values`
+    # scaled by (slice_idx + 1) along `axis` dimension.
+    def scale_per_slice(shape, axis, values):
+      # Note: repeats the values if the shape is larger than values.
+      out = np.take(values, np.remainder(np.arange(np.prod(shape)),
+                                         len(values))).reshape(shape)
+      if axis is not None:
+        scale_shape = [1] * len(shape)
+        scale_shape[axis] = shape[axis]
+        out *= np.arange(1, shape[axis] + 1).reshape(scale_shape)
+      return out
+
+    shape = np.array([2, 3, 4, 5])
+    values = np.array([-128, -64, 0, 38, 102, 71, 64], dtype=np.int32)
+    dequant_values = np.array([-2, -1.0, 0, 0.59375, 1.59375, 1.109375, 1.0],
+                              dtype=np.float32)
+    for axis in [None, 0, 1, 2, 3]:
+      inputs = constant_op.constant(
+          scale_per_slice(shape, None, values), dtype=dtypes.qint8)
+      expected_dequantized = scale_per_slice(shape, axis, dequant_values)
+      if axis is None:
+        min_range, max_range = -2.0, 1.6
+      else:
+        num_slices = shape[axis]
+        min_range, max_range = [], []
+        for slice_idx in range(num_slices):
+          min_range.append(-2.0 * (slice_idx + 1))
+          max_range.append(1.6 * (slice_idx + 1))
+      dequantized = self.evaluate(
+          array_ops.dequantize(
+              inputs, min_range, max_range, mode="SCALED", axis=axis))
+      self.assertAllEqual(dequantized, expected_dequantized)
+      if axis is not None:
+        dequantized = self.evaluate(
+            array_ops.dequantize(
+                inputs, min_range, max_range, mode="SCALED", axis=(axis - 4)))
+        self.assertAllClose(dequantized, expected_dequantized)
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/ops/distributions/BUILD

@@ -0,0 +1,418 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:internal"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "distributions",
+    srcs = ["__init__.py"],
+    deprecation = ("TensorFlow Distributions has migrated to " +
+                   "TensorFlow Probability " +
+                   "(https://github.com/tensorflow/probability). " +
+                   "Deprecated copies remaining in tf.distributions " +
+                   "will not receive new features, and will be removed by " +
+                   "early 2019. You should update all usage of " +
+                   "`tf.distributions` to `tfp.distributions`."),
+    strict_deps = True,
+    deps = [
+        ":distributions_py",
+    ],
+)
+
+py_library(
+    name = "distributions_py",
+    srcs = ["distributions.py"],
+    strict_deps = True,
+    deps = [
+        ":bernoulli",
+        ":beta",
+        ":categorical",
+        ":dirichlet",
+        ":dirichlet_multinomial",
+        ":distribution",
+        ":exponential",
+        ":gamma",
+        ":kullback_leibler",
+        ":laplace",
+        ":multinomial",
+        ":normal",
+        ":student_t",
+        ":uniform",
+        "//tensorflow/python/util:deprecation",
+    ],
+)
+
+py_library(
+    name = "util",
+    srcs = ["util.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:array_ops_stack",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:linalg_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/util:tf_inspect",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_library(
+    name = "kullback_leibler",
+    srcs = ["kullback_leibler.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:control_flow_assert",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+        "//tensorflow/python/util:tf_inspect",
+    ],
+)
+
+py_library(
+    name = "laplace",
+    srcs = ["laplace.py"],
+    strict_deps = True,
+    deps = [
+        ":distribution",
+        ":special_math",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_library(
+    name = "dirichlet",
+    srcs = ["dirichlet.py"],
+    strict_deps = True,
+    deps = [
+        ":distribution",
+        ":kullback_leibler",
+        ":util",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:special_math_ops",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_library(
+    name = "beta",
+    srcs = ["beta.py"],
+    strict_deps = True,
+    deps = [
+        ":distribution",
+        ":kullback_leibler",
+        ":util",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_library(
+    name = "bernoulli",
+    srcs = ["bernoulli.py"],
+    strict_deps = True,
+    deps = [
+        ":distribution",
+        ":kullback_leibler",
+        ":util",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "uniform",
+    srcs = ["uniform.py"],
+    strict_deps = True,
+    deps = [
+        ":distribution",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "special_math",
+    srcs = ["special_math.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_library(
+    name = "bijector_impl",
+    srcs = ["bijector_impl.py"],
+    strict_deps = True,
+    deps = [
+        ":util",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/util:object_identity",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_library(
+    name = "normal",
+    srcs = ["normal.py"],
+    strict_deps = True,
+    deps = [
+        ":distribution",
+        ":kullback_leibler",
+        ":special_math",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "dirichlet_multinomial",
+    srcs = ["dirichlet_multinomial.py"],
+    strict_deps = True,
+    deps = [
+        ":distribution",
+        ":util",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:special_math_ops",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "distribution",
+    srcs = ["distribution.py"],
+    strict_deps = True,
+    deps = [
+        ":kullback_leibler",
+        ":util",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+        "//tensorflow/python/util:tf_inspect",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_library(
+    name = "identity_bijector",
+    srcs = ["identity_bijector.py"],
+    strict_deps = True,
+    deps = [
+        ":bijector",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/util:deprecation",
+    ],
+)
+
+py_library(
+    name = "categorical",
+    srcs = ["categorical.py"],
+    strict_deps = True,
+    deps = [
+        ":distribution",
+        ":kullback_leibler",
+        ":util",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "bijector_test_util",
+    srcs = ["bijector_test_util.py"],
+    strict_deps = True,
+    deps = [
+        ":uniform",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/ops:math_ops",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_library(
+    name = "bijector",
+    srcs = ["bijector.py"],
+    strict_deps = True,
+    deps = [":bijector_impl"],
+)
+
+py_library(
+    name = "exponential",
+    srcs = ["exponential.py"],
+    strict_deps = True,
+    deps = [
+        ":gamma",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_library(
+    name = "gamma",
+    srcs = ["gamma.py"],
+    strict_deps = True,
+    deps = [
+        ":distribution",
+        ":kullback_leibler",
+        ":util",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_library(
+    name = "multinomial",
+    srcs = ["multinomial.py"],
+    strict_deps = True,
+    deps = [
+        ":distribution",
+        ":util",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:map_fn",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "transformed_distribution",
+    srcs = ["transformed_distribution.py"],
+    strict_deps = True,
+    deps = [
+        ":distribution",
+        ":identity_bijector",
+        ":util",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:array_ops_stack",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/util:deprecation",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_library(
+    name = "student_t",
+    srcs = ["student_t.py"],
+    strict_deps = True,
+    deps = [
+        ":distribution",
+        ":util",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:special_math_ops",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+        "//third_party/py/numpy",
+    ],
+)

tensorflow/python/ops/distributions/__init__.py

@@ -0,0 +1,16 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Core module for TensorFlow distribution objects and helpers."""
+from tensorflow.python.ops.distributions import distributions

tensorflow/python/ops/distributions/bernoulli.py

@@ -0,0 +1,183 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""The Bernoulli distribution class."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops.distributions import distribution
+from tensorflow.python.ops.distributions import kullback_leibler
+from tensorflow.python.ops.distributions import util as distribution_util
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["distributions.Bernoulli"])
+class Bernoulli(distribution.Distribution):
+  """Bernoulli distribution.
+
+  The Bernoulli distribution with `probs` parameter, i.e., the probability of a
+  `1` outcome (vs a `0` outcome).
+  """
+
+  @deprecation.deprecated(
+      "2019-01-01",
+      "The TensorFlow Distributions library has moved to "
+      "TensorFlow Probability "
+      "(https://github.com/tensorflow/probability). You "
+      "should update all references to use `tfp.distributions` "
+      "instead of `tf.distributions`.",
+      warn_once=True)
+  def __init__(self,
+               logits=None,
+               probs=None,
+               dtype=dtypes.int32,
+               validate_args=False,
+               allow_nan_stats=True,
+               name="Bernoulli"):
+    """Construct Bernoulli distributions.
+
+    Args:
+      logits: An N-D `Tensor` representing the log-odds of a `1` event. Each
+        entry in the `Tensor` parametrizes an independent Bernoulli distribution
+        where the probability of an event is sigmoid(logits). Only one of
+        `logits` or `probs` should be passed in.
+      probs: An N-D `Tensor` representing the probability of a `1`
+        event. Each entry in the `Tensor` parameterizes an independent
+        Bernoulli distribution. Only one of `logits` or `probs` should be passed
+        in.
+      dtype: The type of the event samples. Default: `int32`.
+      validate_args: Python `bool`, default `False`. When `True` distribution
+        parameters are checked for validity despite possibly degrading runtime
+        performance. When `False` invalid inputs may silently render incorrect
+        outputs.
+      allow_nan_stats: Python `bool`, default `True`. When `True`,
+        statistics (e.g., mean, mode, variance) use the value "`NaN`" to
+        indicate the result is undefined. When `False`, an exception is raised
+        if one or more of the statistic's batch members are undefined.
+      name: Python `str` name prefixed to Ops created by this class.
+
+    Raises:
+      ValueError: If p and logits are passed, or if neither are passed.
+    """
+    parameters = dict(locals())
+    with ops.name_scope(name) as name:
+      self._logits, self._probs = distribution_util.get_logits_and_probs(
+          logits=logits,
+          probs=probs,
+          validate_args=validate_args,
+          name=name)
+    super(Bernoulli, self).__init__(
+        dtype=dtype,
+        reparameterization_type=distribution.NOT_REPARAMETERIZED,
+        validate_args=validate_args,
+        allow_nan_stats=allow_nan_stats,
+        parameters=parameters,
+        graph_parents=[self._logits, self._probs],
+        name=name)
+
+  @staticmethod
+  def _param_shapes(sample_shape):
+    return {"logits": ops.convert_to_tensor(sample_shape, dtype=dtypes.int32)}
+
+  @property
+  def logits(self):
+    """Log-odds of a `1` outcome (vs `0`)."""
+    return self._logits
+
+  @property
+  def probs(self):
+    """Probability of a `1` outcome (vs `0`)."""
+    return self._probs
+
+  def _batch_shape_tensor(self):
+    return array_ops.shape(self._logits)
+
+  def _batch_shape(self):
+    return self._logits.get_shape()
+
+  def _event_shape_tensor(self):
+    return array_ops.constant([], dtype=dtypes.int32)
+
+  def _event_shape(self):
+    return tensor_shape.TensorShape([])
+
+  def _sample_n(self, n, seed=None):
+    new_shape = array_ops.concat([[n], self.batch_shape_tensor()], 0)
+    uniform = random_ops.random_uniform(
+        new_shape, seed=seed, dtype=self.probs.dtype)
+    sample = math_ops.less(uniform, self.probs)
+    return math_ops.cast(sample, self.dtype)
+
+  def _log_prob(self, event):
+    if self.validate_args:
+      event = distribution_util.embed_check_integer_casting_closed(
+          event, target_dtype=dtypes.bool)
+
+    # TODO(jaana): The current sigmoid_cross_entropy_with_logits has
+    # inconsistent behavior for logits = inf/-inf.
+    event = math_ops.cast(event, self.logits.dtype)
+    logits = self.logits
+    # sigmoid_cross_entropy_with_logits doesn't broadcast shape,
+    # so we do this here.
+
+    def _broadcast(logits, event):
+      return (array_ops.ones_like(event) * logits,
+              array_ops.ones_like(logits) * event)
+
+    if not (event.get_shape().is_fully_defined() and
+            logits.get_shape().is_fully_defined() and
+            event.get_shape() == logits.get_shape()):
+      logits, event = _broadcast(logits, event)
+    return -nn.sigmoid_cross_entropy_with_logits(labels=event, logits=logits)
+
+  def _entropy(self):
+    return (-self.logits * (math_ops.sigmoid(self.logits) - 1) +  # pylint: disable=invalid-unary-operand-type
+            nn.softplus(-self.logits))  # pylint: disable=invalid-unary-operand-type
+
+  def _mean(self):
+    return array_ops.identity(self.probs)
+
+  def _variance(self):
+    return self._mean() * (1. - self.probs)
+
+  def _mode(self):
+    """Returns `1` if `prob > 0.5` and `0` otherwise."""
+    return math_ops.cast(self.probs > 0.5, self.dtype)
+
+
+@kullback_leibler.RegisterKL(Bernoulli, Bernoulli)
+def _kl_bernoulli_bernoulli(a, b, name=None):
+  """Calculate the batched KL divergence KL(a || b) with a and b Bernoulli.
+
+  Args:
+    a: instance of a Bernoulli distribution object.
+    b: instance of a Bernoulli distribution object.
+    name: (optional) Name to use for created operations.
+      default is "kl_bernoulli_bernoulli".
+
+  Returns:
+    Batchwise KL(a || b)
+  """
+  with ops.name_scope(name, "kl_bernoulli_bernoulli",
+                      values=[a.logits, b.logits]):
+    delta_probs0 = nn.softplus(-b.logits) - nn.softplus(-a.logits)
+    delta_probs1 = nn.softplus(b.logits) - nn.softplus(a.logits)
+    return (math_ops.sigmoid(a.logits) * delta_probs0
+            + math_ops.sigmoid(-a.logits) * delta_probs1)

tensorflow/python/ops/distributions/beta.py

@@ -0,0 +1,407 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""The Beta distribution class."""
+
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops.distributions import distribution
+from tensorflow.python.ops.distributions import kullback_leibler
+from tensorflow.python.ops.distributions import util as distribution_util
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+__all__ = [
+    "Beta",
+    "BetaWithSoftplusConcentration",
+]
+
+
+_beta_sample_note = """Note: `x` must have dtype `self.dtype` and be in
+`[0, 1].` It must have a shape compatible with `self.batch_shape()`."""
+
+
+@tf_export(v1=["distributions.Beta"])
+class Beta(distribution.Distribution):
+  """Beta distribution.
+
+  The Beta distribution is defined over the `(0, 1)` interval using parameters
+  `concentration1` (aka "alpha") and `concentration0` (aka "beta").
+
+  #### Mathematical Details
+
+  The probability density function (pdf) is,
+
+  ```none
+  pdf(x; alpha, beta) = x**(alpha - 1) (1 - x)**(beta - 1) / Z
+  Z = Gamma(alpha) Gamma(beta) / Gamma(alpha + beta)
+  ```
+
+  where:
+
+  * `concentration1 = alpha`,
+  * `concentration0 = beta`,
+  * `Z` is the normalization constant, and,
+  * `Gamma` is the [gamma function](
+    https://en.wikipedia.org/wiki/Gamma_function).
+
+  The concentration parameters represent mean total counts of a `1` or a `0`,
+  i.e.,
+
+  ```none
+  concentration1 = alpha = mean * total_concentration
+  concentration0 = beta  = (1. - mean) * total_concentration
+  ```
+
+  where `mean` in `(0, 1)` and `total_concentration` is a positive real number
+  representing a mean `total_count = concentration1 + concentration0`.
+
+  Distribution parameters are automatically broadcast in all functions; see
+  examples for details.
+
+  Warning: The samples can be zero due to finite precision.
+  This happens more often when some of the concentrations are very small.
+  Make sure to round the samples to `np.finfo(dtype).tiny` before computing the
+  density.
+
+  Samples of this distribution are reparameterized (pathwise differentiable).
+  The derivatives are computed using the approach described in
+  (Figurnov et al., 2018).
+
+  #### Examples
+
+  ```python
+  import tensorflow_probability as tfp
+  tfd = tfp.distributions
+
+  # Create a batch of three Beta distributions.
+  alpha = [1, 2, 3]
+  beta = [1, 2, 3]
+  dist = tfd.Beta(alpha, beta)
+
+  dist.sample([4, 5])  # Shape [4, 5, 3]
+
+  # `x` has three batch entries, each with two samples.
+  x = [[.1, .4, .5],
+       [.2, .3, .5]]
+  # Calculate the probability of each pair of samples under the corresponding
+  # distribution in `dist`.
+  dist.prob(x)         # Shape [2, 3]
+  ```
+
+  ```python
+  # Create batch_shape=[2, 3] via parameter broadcast:
+  alpha = [[1.], [2]]      # Shape [2, 1]
+  beta = [3., 4, 5]        # Shape [3]
+  dist = tfd.Beta(alpha, beta)
+
+  # alpha broadcast as: [[1., 1, 1,],
+  #                      [2, 2, 2]]
+  # beta broadcast as:  [[3., 4, 5],
+  #                      [3, 4, 5]]
+  # batch_Shape [2, 3]
+  dist.sample([4, 5])  # Shape [4, 5, 2, 3]
+
+  x = [.2, .3, .5]
+  # x will be broadcast as [[.2, .3, .5],
+  #                         [.2, .3, .5]],
+  # thus matching batch_shape [2, 3].
+  dist.prob(x)         # Shape [2, 3]
+  ```
+
+  Compute the gradients of samples w.r.t. the parameters:
+
+  ```python
+  alpha = tf.constant(1.0)
+  beta = tf.constant(2.0)
+  dist = tfd.Beta(alpha, beta)
+  samples = dist.sample(5)  # Shape [5]
+  loss = tf.reduce_mean(tf.square(samples))  # Arbitrary loss function
+  # Unbiased stochastic gradients of the loss function
+  grads = tf.gradients(loss, [alpha, beta])
+  ```
+
+  References:
+    Implicit Reparameterization Gradients:
+      [Figurnov et al., 2018]
+      (http://papers.nips.cc/paper/7326-implicit-reparameterization-gradients)
+      ([pdf]
+      (http://papers.nips.cc/paper/7326-implicit-reparameterization-gradients.pdf))
+  """
+
+  @deprecation.deprecated(
+      "2019-01-01",
+      "The TensorFlow Distributions library has moved to "
+      "TensorFlow Probability "
+      "(https://github.com/tensorflow/probability). You "
+      "should update all references to use `tfp.distributions` "
+      "instead of `tf.distributions`.",
+      warn_once=True)
+  def __init__(self,
+               concentration1=None,
+               concentration0=None,
+               validate_args=False,
+               allow_nan_stats=True,
+               name="Beta"):
+    """Initialize a batch of Beta distributions.
+
+    Args:
+      concentration1: Positive floating-point `Tensor` indicating mean
+        number of successes; aka "alpha". Implies `self.dtype` and
+        `self.batch_shape`, i.e.,
+        `concentration1.shape = [N1, N2, ..., Nm] = self.batch_shape`.
+      concentration0: Positive floating-point `Tensor` indicating mean
+        number of failures; aka "beta". Otherwise has same semantics as
+        `concentration1`.
+      validate_args: Python `bool`, default `False`. When `True` distribution
+        parameters are checked for validity despite possibly degrading runtime
+        performance. When `False` invalid inputs may silently render incorrect
+        outputs.
+      allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
+        (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
+        result is undefined. When `False`, an exception is raised if one or
+        more of the statistic's batch members are undefined.
+      name: Python `str` name prefixed to Ops created by this class.
+    """
+    parameters = dict(locals())
+    with ops.name_scope(name, values=[concentration1, concentration0]) as name:
+      self._concentration1 = self._maybe_assert_valid_concentration(
+          ops.convert_to_tensor(concentration1, name="concentration1"),
+          validate_args)
+      self._concentration0 = self._maybe_assert_valid_concentration(
+          ops.convert_to_tensor(concentration0, name="concentration0"),
+          validate_args)
+      check_ops.assert_same_float_dtype([
+          self._concentration1, self._concentration0])
+      self._total_concentration = self._concentration1 + self._concentration0
+    super(Beta, self).__init__(
+        dtype=self._total_concentration.dtype,
+        validate_args=validate_args,
+        allow_nan_stats=allow_nan_stats,
+        reparameterization_type=distribution.FULLY_REPARAMETERIZED,
+        parameters=parameters,
+        graph_parents=[self._concentration1,
+                       self._concentration0,
+                       self._total_concentration],
+        name=name)
+
+  @staticmethod
+  def _param_shapes(sample_shape):
+    return dict(zip(
+        ["concentration1", "concentration0"],
+        [ops.convert_to_tensor(sample_shape, dtype=dtypes.int32)] * 2))
+
+  @property
+  def concentration1(self):
+    """Concentration parameter associated with a `1` outcome."""
+    return self._concentration1
+
+  @property
+  def concentration0(self):
+    """Concentration parameter associated with a `0` outcome."""
+    return self._concentration0
+
+  @property
+  def total_concentration(self):
+    """Sum of concentration parameters."""
+    return self._total_concentration
+
+  def _batch_shape_tensor(self):
+    return array_ops.shape(self.total_concentration)
+
+  def _batch_shape(self):
+    return self.total_concentration.get_shape()
+
+  def _event_shape_tensor(self):
+    return constant_op.constant([], dtype=dtypes.int32)
+
+  def _event_shape(self):
+    return tensor_shape.TensorShape([])
+
+  def _sample_n(self, n, seed=None):
+    expanded_concentration1 = array_ops.ones_like(
+        self.total_concentration, dtype=self.dtype) * self.concentration1
+    expanded_concentration0 = array_ops.ones_like(
+        self.total_concentration, dtype=self.dtype) * self.concentration0
+    gamma1_sample = random_ops.random_gamma(
+        shape=[n],
+        alpha=expanded_concentration1,
+        dtype=self.dtype,
+        seed=seed)
+    gamma2_sample = random_ops.random_gamma(
+        shape=[n],
+        alpha=expanded_concentration0,
+        dtype=self.dtype,
+        seed=distribution_util.gen_new_seed(seed, "beta"))
+    beta_sample = gamma1_sample / (gamma1_sample + gamma2_sample)
+    return beta_sample
+
+  @distribution_util.AppendDocstring(_beta_sample_note)
+  def _log_prob(self, x):
+    return self._log_unnormalized_prob(x) - self._log_normalization()
+
+  @distribution_util.AppendDocstring(_beta_sample_note)
+  def _prob(self, x):
+    return math_ops.exp(self._log_prob(x))
+
+  @distribution_util.AppendDocstring(_beta_sample_note)
+  def _log_cdf(self, x):
+    return math_ops.log(self._cdf(x))
+
+  @distribution_util.AppendDocstring(_beta_sample_note)
+  def _cdf(self, x):
+    return math_ops.betainc(self.concentration1, self.concentration0, x)
+
+  def _log_unnormalized_prob(self, x):
+    x = self._maybe_assert_valid_sample(x)
+    return (math_ops.xlogy(self.concentration1 - 1., x) +
+            (self.concentration0 - 1.) * math_ops.log1p(-x))  # pylint: disable=invalid-unary-operand-type
+
+  def _log_normalization(self):
+    return (math_ops.lgamma(self.concentration1)
+            + math_ops.lgamma(self.concentration0)
+            - math_ops.lgamma(self.total_concentration))
+
+  def _entropy(self):
+    return (
+        self._log_normalization()
+        - (self.concentration1 - 1.) * math_ops.digamma(self.concentration1)
+        - (self.concentration0 - 1.) * math_ops.digamma(self.concentration0)
+        + ((self.total_concentration - 2.) *
+           math_ops.digamma(self.total_concentration)))
+
+  def _mean(self):
+    return self._concentration1 / self._total_concentration
+
+  def _variance(self):
+    return self._mean() * (1. - self._mean()) / (1. + self.total_concentration)
+
+  @distribution_util.AppendDocstring(
+      """Note: The mode is undefined when `concentration1 <= 1` or
+      `concentration0 <= 1`. If `self.allow_nan_stats` is `True`, `NaN`
+      is used for undefined modes. If `self.allow_nan_stats` is `False` an
+      exception is raised when one or more modes are undefined.""")
+  def _mode(self):
+    mode = (self.concentration1 - 1.) / (self.total_concentration - 2.)
+    if self.allow_nan_stats:
+      nan = array_ops.fill(
+          self.batch_shape_tensor(),
+          np.array(np.nan, dtype=self.dtype.as_numpy_dtype()),
+          name="nan")
+      is_defined = math_ops.logical_and(self.concentration1 > 1.,
+                                        self.concentration0 > 1.)
+      return array_ops.where_v2(is_defined, mode, nan)
+    return control_flow_ops.with_dependencies([
+        check_ops.assert_less(
+            array_ops.ones([], dtype=self.dtype),
+            self.concentration1,
+            message="Mode undefined for concentration1 <= 1."),
+        check_ops.assert_less(
+            array_ops.ones([], dtype=self.dtype),
+            self.concentration0,
+            message="Mode undefined for concentration0 <= 1.")
+    ], mode)
+
+  def _maybe_assert_valid_concentration(self, concentration, validate_args):
+    """Checks the validity of a concentration parameter."""
+    if not validate_args:
+      return concentration
+    return control_flow_ops.with_dependencies([
+        check_ops.assert_positive(
+            concentration,
+            message="Concentration parameter must be positive."),
+    ], concentration)
+
+  def _maybe_assert_valid_sample(self, x):
+    """Checks the validity of a sample."""
+    if not self.validate_args:
+      return x
+    return control_flow_ops.with_dependencies([
+        check_ops.assert_positive(x, message="sample must be positive"),
+        check_ops.assert_less(
+            x,
+            array_ops.ones([], self.dtype),
+            message="sample must be less than `1`."),
+    ], x)
+
+
+class BetaWithSoftplusConcentration(Beta):
+  """Beta with softplus transform of `concentration1` and `concentration0`."""
+
+  @deprecation.deprecated(
+      "2019-01-01",
+      "Use `tfd.Beta(tf.nn.softplus(concentration1), "
+      "tf.nn.softplus(concentration2))` instead.",
+      warn_once=True)
+  def __init__(self,
+               concentration1,
+               concentration0,
+               validate_args=False,
+               allow_nan_stats=True,
+               name="BetaWithSoftplusConcentration"):
+    parameters = dict(locals())
+    with ops.name_scope(name, values=[concentration1,
+                                      concentration0]) as name:
+      super(BetaWithSoftplusConcentration, self).__init__(
+          concentration1=nn.softplus(concentration1,
+                                     name="softplus_concentration1"),
+          concentration0=nn.softplus(concentration0,
+                                     name="softplus_concentration0"),
+          validate_args=validate_args,
+          allow_nan_stats=allow_nan_stats,
+          name=name)
+    self._parameters = parameters
+
+
+@kullback_leibler.RegisterKL(Beta, Beta)
+def _kl_beta_beta(d1, d2, name=None):
+  """Calculate the batchwise KL divergence KL(d1 || d2) with d1 and d2 Beta.
+
+  Args:
+    d1: instance of a Beta distribution object.
+    d2: instance of a Beta distribution object.
+    name: (optional) Name to use for created operations.
+      default is "kl_beta_beta".
+
+  Returns:
+    Batchwise KL(d1 || d2)
+  """
+  def delta(fn, is_property=True):
+    fn1 = getattr(d1, fn)
+    fn2 = getattr(d2, fn)
+    return (fn2 - fn1) if is_property else (fn2() - fn1())
+  with ops.name_scope(name, "kl_beta_beta", values=[
+      d1.concentration1,
+      d1.concentration0,
+      d1.total_concentration,
+      d2.concentration1,
+      d2.concentration0,
+      d2.total_concentration,
+  ]):
+    return (delta("_log_normalization", is_property=False)
+            - math_ops.digamma(d1.concentration1) * delta("concentration1")
+            - math_ops.digamma(d1.concentration0) * delta("concentration0")
+            + (math_ops.digamma(d1.total_concentration)
+               * delta("total_concentration")))

tensorflow/python/ops/distributions/bijector.py

@@ -0,0 +1,21 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Bijector base."""
+
+# go/tf-wildcard-import
+# pylint: disable=wildcard-import,unused-import
+from tensorflow.python.ops.distributions.bijector_impl import Bijector
+
+# pylint: enable=wildcard-import,unused-import

tensorflow/python/ops/distributions/bijector_test_util.py

@@ -0,0 +1,221 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Bijector unit-test utilities."""
+
+import numpy as np
+
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.distributions import uniform as uniform_lib
+
+
+def assert_finite(array):
+  if not np.isfinite(array).all():
+    raise AssertionError("array was not all finite. %s" % array[:15])
+
+
+def assert_strictly_increasing(array):
+  np.testing.assert_array_less(0., np.diff(array))
+
+
+def assert_strictly_decreasing(array):
+  np.testing.assert_array_less(np.diff(array), 0.)
+
+
+def assert_strictly_monotonic(array):
+  if array[0] < array[-1]:
+    assert_strictly_increasing(array)
+  else:
+    assert_strictly_decreasing(array)
+
+
+def assert_scalar_congruency(bijector,
+                             lower_x,
+                             upper_x,
+                             n=int(10e3),
+                             rtol=0.01,
+                             sess=None):
+  """Assert `bijector`'s forward/inverse/inverse_log_det_jacobian are congruent.
+
+  We draw samples `X ~ U(lower_x, upper_x)`, then feed these through the
+  `bijector` in order to check that:
+
+  1. the forward is strictly monotonic.
+  2. the forward/inverse methods are inverses of each other.
+  3. the jacobian is the correct change of measure.
+
+  This can only be used for a Bijector mapping open subsets of the real line
+  to themselves.  This is due to the fact that this test compares the `prob`
+  before/after transformation with the Lebesgue measure on the line.
+
+  Args:
+    bijector:  Instance of Bijector
+    lower_x:  Python scalar.
+    upper_x:  Python scalar.  Must have `lower_x < upper_x`, and both must be in
+      the domain of the `bijector`.  The `bijector` should probably not produce
+      huge variation in values in the interval `(lower_x, upper_x)`, or else
+      the variance based check of the Jacobian will require small `rtol` or
+      huge `n`.
+    n:  Number of samples to draw for the checks.
+    rtol:  Positive number.  Used for the Jacobian check.
+    sess:  `tf.compat.v1.Session`.  Defaults to the default session.
+
+  Raises:
+    AssertionError:  If tests fail.
+  """
+  # Checks and defaults.
+  if sess is None:
+    sess = ops.get_default_session()
+
+  # Should be monotonic over this interval
+  ten_x_pts = np.linspace(lower_x, upper_x, num=10).astype(np.float32)
+  if bijector.dtype is not None:
+    ten_x_pts = ten_x_pts.astype(bijector.dtype.as_numpy_dtype)
+  forward_on_10_pts = bijector.forward(ten_x_pts)
+
+  # Set the lower/upper limits in the range of the bijector.
+  lower_y, upper_y = sess.run(
+      [bijector.forward(lower_x), bijector.forward(upper_x)])
+  if upper_y < lower_y:  # If bijector.forward is a decreasing function.
+    lower_y, upper_y = upper_y, lower_y
+
+  # Uniform samples from the domain, range.
+  uniform_x_samps = uniform_lib.Uniform(
+      low=lower_x, high=upper_x).sample(n, seed=0)
+  uniform_y_samps = uniform_lib.Uniform(
+      low=lower_y, high=upper_y).sample(n, seed=1)
+
+  # These compositions should be the identity.
+  inverse_forward_x = bijector.inverse(bijector.forward(uniform_x_samps))
+  forward_inverse_y = bijector.forward(bijector.inverse(uniform_y_samps))
+
+  # For a < b, and transformation y = y(x),
+  # (b - a) = \int_a^b dx = \int_{y(a)}^{y(b)} |dx/dy| dy
+  # "change_measure_dy_dx" below is a Monte Carlo approximation to the right
+  # hand side, which should then be close to the left, which is (b - a).
+  # We assume event_ndims=0 because we assume scalar -> scalar. The log_det
+  # methods will handle whether they expect event_ndims > 0.
+  dy_dx = math_ops.exp(bijector.inverse_log_det_jacobian(
+      uniform_y_samps, event_ndims=0))
+  # E[|dx/dy|] under Uniform[lower_y, upper_y]
+  # = \int_{y(a)}^{y(b)} |dx/dy| dP(u), where dP(u) is the uniform measure
+  expectation_of_dy_dx_under_uniform = math_ops.reduce_mean(dy_dx)
+  # dy = dP(u) * (upper_y - lower_y)
+  change_measure_dy_dx = (
+      (upper_y - lower_y) * expectation_of_dy_dx_under_uniform)
+
+  # We'll also check that dy_dx = 1 / dx_dy.
+  dx_dy = math_ops.exp(
+      bijector.forward_log_det_jacobian(
+          bijector.inverse(uniform_y_samps), event_ndims=0))
+
+  [
+      forward_on_10_pts_v,
+      dy_dx_v,
+      dx_dy_v,
+      change_measure_dy_dx_v,
+      uniform_x_samps_v,
+      uniform_y_samps_v,
+      inverse_forward_x_v,
+      forward_inverse_y_v,
+  ] = sess.run([
+      forward_on_10_pts,
+      dy_dx,
+      dx_dy,
+      change_measure_dy_dx,
+      uniform_x_samps,
+      uniform_y_samps,
+      inverse_forward_x,
+      forward_inverse_y,
+  ])
+
+  assert_strictly_monotonic(forward_on_10_pts_v)
+  # Composition of forward/inverse should be the identity.
+  np.testing.assert_allclose(
+      inverse_forward_x_v, uniform_x_samps_v, atol=1e-5, rtol=1e-3)
+  np.testing.assert_allclose(
+      forward_inverse_y_v, uniform_y_samps_v, atol=1e-5, rtol=1e-3)
+  # Change of measure should be correct.
+  np.testing.assert_allclose(
+      upper_x - lower_x, change_measure_dy_dx_v, atol=0, rtol=rtol)
+  # Inverse Jacobian should be equivalent to the reciprocal of the forward
+  # Jacobian.
+  np.testing.assert_allclose(
+      dy_dx_v, np.divide(1., dx_dy_v), atol=1e-5, rtol=1e-3)
+
+
+def assert_bijective_and_finite(
+    bijector, x, y, event_ndims, atol=0, rtol=1e-5, sess=None):
+  """Assert that forward/inverse (along with jacobians) are inverses and finite.
+
+  It is recommended to use x and y values that are very very close to the edge
+  of the Bijector's domain.
+
+  Args:
+    bijector:  A Bijector instance.
+    x:  np.array of values in the domain of bijector.forward.
+    y:  np.array of values in the domain of bijector.inverse.
+    event_ndims: Integer describing the number of event dimensions this bijector
+      operates on.
+    atol:  Absolute tolerance.
+    rtol:  Relative tolerance.
+    sess:  TensorFlow session.  Defaults to the default session.
+
+  Raises:
+    AssertionError:  If tests fail.
+  """
+  sess = sess or ops.get_default_session()
+
+  # These are the incoming points, but people often create a crazy range of
+  # values for which these end up being bad, especially in 16bit.
+  assert_finite(x)
+  assert_finite(y)
+
+  f_x = bijector.forward(x)
+  g_y = bijector.inverse(y)
+
+  [
+      x_from_x,
+      y_from_y,
+      ildj_f_x,
+      fldj_x,
+      ildj_y,
+      fldj_g_y,
+      f_x_v,
+      g_y_v,
+  ] = sess.run([
+      bijector.inverse(f_x),
+      bijector.forward(g_y),
+      bijector.inverse_log_det_jacobian(f_x, event_ndims=event_ndims),
+      bijector.forward_log_det_jacobian(x, event_ndims=event_ndims),
+      bijector.inverse_log_det_jacobian(y, event_ndims=event_ndims),
+      bijector.forward_log_det_jacobian(g_y, event_ndims=event_ndims),
+      f_x,
+      g_y,
+  ])
+
+  assert_finite(x_from_x)
+  assert_finite(y_from_y)
+  assert_finite(ildj_f_x)
+  assert_finite(fldj_x)
+  assert_finite(ildj_y)
+  assert_finite(fldj_g_y)
+  assert_finite(f_x_v)
+  assert_finite(g_y_v)
+
+  np.testing.assert_allclose(x_from_x, x, atol=atol, rtol=rtol)
+  np.testing.assert_allclose(y_from_y, y, atol=atol, rtol=rtol)
+  np.testing.assert_allclose(-ildj_f_x, fldj_x, atol=atol, rtol=rtol)
+  np.testing.assert_allclose(-ildj_y, fldj_g_y, atol=atol, rtol=rtol)

tensorflow/python/ops/distributions/categorical.py

@@ -0,0 +1,345 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""The Categorical distribution class."""
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops.distributions import distribution
+from tensorflow.python.ops.distributions import kullback_leibler
+from tensorflow.python.ops.distributions import util as distribution_util
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+def _broadcast_cat_event_and_params(event, params, base_dtype):
+  """Broadcasts the event or distribution parameters."""
+  if event.dtype.is_integer:
+    pass
+  elif event.dtype.is_floating:
+    # When `validate_args=True` we've already ensured int/float casting
+    # is closed.
+    event = math_ops.cast(event, dtype=dtypes.int32)
+  else:
+    raise TypeError("`value` should have integer `dtype` or "
+                    "`self.dtype` ({})".format(base_dtype))
+  shape_known_statically = (
+      params.shape.ndims is not None and
+      params.shape[:-1].is_fully_defined() and
+      event.shape.is_fully_defined())
+  if not shape_known_statically or params.shape[:-1] != event.shape:
+    params *= array_ops.ones_like(event[..., array_ops.newaxis],
+                                  dtype=params.dtype)
+    params_shape = array_ops.shape(params)[:-1]
+    event *= array_ops.ones(params_shape, dtype=event.dtype)
+    if params.shape.ndims is not None:
+      event.set_shape(tensor_shape.TensorShape(params.shape[:-1]))
+
+  return event, params
+
+
+@tf_export(v1=["distributions.Categorical"])
+class Categorical(distribution.Distribution):
+  """Categorical distribution.
+
+  The Categorical distribution is parameterized by either probabilities or
+  log-probabilities of a set of `K` classes. It is defined over the integers
+  `{0, 1, ..., K}`.
+
+  The Categorical distribution is closely related to the `OneHotCategorical` and
+  `Multinomial` distributions.  The Categorical distribution can be intuited as
+  generating samples according to `argmax{ OneHotCategorical(probs) }` itself
+  being identical to `argmax{ Multinomial(probs, total_count=1) }`.
+
+  #### Mathematical Details
+
+  The probability mass function (pmf) is,
+
+  ```none
+  pmf(k; pi) = prod_j pi_j**[k == j]
+  ```
+
+  #### Pitfalls
+
+  The number of classes, `K`, must not exceed:
+  - the largest integer representable by `self.dtype`, i.e.,
+    `2**(mantissa_bits+1)` (IEEE 754),
+  - the maximum `Tensor` index, i.e., `2**31-1`.
+
+  In other words,
+
+  ```python
+  K <= min(2**31-1, {
+    tf.float16: 2**11,
+    tf.float32: 2**24,
+    tf.float64: 2**53 }[param.dtype])
+  ```
+
+  Note: This condition is validated only when `self.validate_args = True`.
+
+  #### Examples
+
+  Creates a 3-class distribution with the 2nd class being most likely.
+
+  ```python
+  dist = Categorical(probs=[0.1, 0.5, 0.4])
+  n = 1e4
+  empirical_prob = tf.cast(
+      tf.histogram_fixed_width(
+        dist.sample(int(n)),
+        [0., 2],
+        nbins=3),
+      dtype=tf.float32) / n
+  # ==> array([ 0.1005,  0.5037,  0.3958], dtype=float32)
+  ```
+
+  Creates a 3-class distribution with the 2nd class being most likely.
+  Parameterized by [logits](https://en.wikipedia.org/wiki/Logit) rather than
+  probabilities.
+
+  ```python
+  dist = Categorical(logits=np.log([0.1, 0.5, 0.4])
+  n = 1e4
+  empirical_prob = tf.cast(
+      tf.histogram_fixed_width(
+        dist.sample(int(n)),
+        [0., 2],
+        nbins=3),
+      dtype=tf.float32) / n
+  # ==> array([0.1045,  0.5047, 0.3908], dtype=float32)
+  ```
+
+  Creates a 3-class distribution with the 3rd class being most likely.
+  The distribution functions can be evaluated on counts.
+
+  ```python
+  # counts is a scalar.
+  p = [0.1, 0.4, 0.5]
+  dist = Categorical(probs=p)
+  dist.prob(0)  # Shape []
+
+  # p will be broadcast to [[0.1, 0.4, 0.5], [0.1, 0.4, 0.5]] to match counts.
+  counts = [1, 0]
+  dist.prob(counts)  # Shape [2]
+
+  # p will be broadcast to shape [3, 5, 7, 3] to match counts.
+  counts = [[...]] # Shape [5, 7, 3]
+  dist.prob(counts)  # Shape [5, 7, 3]
+  ```
+
+  """
+
+  @deprecation.deprecated(
+      "2019-01-01",
+      "The TensorFlow Distributions library has moved to "
+      "TensorFlow Probability "
+      "(https://github.com/tensorflow/probability). You "
+      "should update all references to use `tfp.distributions` "
+      "instead of `tf.distributions`.",
+      warn_once=True)
+  def __init__(
+      self,
+      logits=None,
+      probs=None,
+      dtype=dtypes.int32,
+      validate_args=False,
+      allow_nan_stats=True,
+      name="Categorical"):
+    """Initialize Categorical distributions using class log-probabilities.
+
+    Args:
+      logits: An N-D `Tensor`, `N >= 1`, representing the log probabilities
+        of a set of Categorical distributions. The first `N - 1` dimensions
+        index into a batch of independent distributions and the last dimension
+        represents a vector of logits for each class. Only one of `logits` or
+        `probs` should be passed in.
+      probs: An N-D `Tensor`, `N >= 1`, representing the probabilities
+        of a set of Categorical distributions. The first `N - 1` dimensions
+        index into a batch of independent distributions and the last dimension
+        represents a vector of probabilities for each class. Only one of
+        `logits` or `probs` should be passed in.
+      dtype: The type of the event samples (default: int32).
+      validate_args: Python `bool`, default `False`. When `True` distribution
+        parameters are checked for validity despite possibly degrading runtime
+        performance. When `False` invalid inputs may silently render incorrect
+        outputs.
+      allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
+        (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
+        result is undefined. When `False`, an exception is raised if one or
+        more of the statistic's batch members are undefined.
+      name: Python `str` name prefixed to Ops created by this class.
+    """
+    parameters = dict(locals())
+    with ops.name_scope(name, values=[logits, probs]) as name:
+      self._logits, self._probs = distribution_util.get_logits_and_probs(
+          logits=logits,
+          probs=probs,
+          validate_args=validate_args,
+          multidimensional=True,
+          name=name)
+
+      if validate_args:
+        self._logits = distribution_util.embed_check_categorical_event_shape(
+            self._logits)
+
+      logits_shape_static = self._logits.get_shape().with_rank_at_least(1)
+      if logits_shape_static.ndims is not None:
+        self._batch_rank = ops.convert_to_tensor(
+            logits_shape_static.ndims - 1,
+            dtype=dtypes.int32,
+            name="batch_rank")
+      else:
+        with ops.name_scope(name="batch_rank"):
+          self._batch_rank = array_ops.rank(self._logits) - 1
+
+      logits_shape = array_ops.shape(self._logits, name="logits_shape")
+      if tensor_shape.dimension_value(logits_shape_static[-1]) is not None:
+        self._event_size = ops.convert_to_tensor(
+            logits_shape_static.dims[-1].value,
+            dtype=dtypes.int32,
+            name="event_size")
+      else:
+        with ops.name_scope(name="event_size"):
+          self._event_size = logits_shape[self._batch_rank]
+
+      if logits_shape_static[:-1].is_fully_defined():
+        self._batch_shape_val = constant_op.constant(
+            logits_shape_static[:-1].as_list(),
+            dtype=dtypes.int32,
+            name="batch_shape")
+      else:
+        with ops.name_scope(name="batch_shape"):
+          self._batch_shape_val = logits_shape[:-1]
+    super(Categorical, self).__init__(
+        dtype=dtype,
+        reparameterization_type=distribution.NOT_REPARAMETERIZED,
+        validate_args=validate_args,
+        allow_nan_stats=allow_nan_stats,
+        parameters=parameters,
+        graph_parents=[self._logits,
+                       self._probs],
+        name=name)
+
+  @property
+  def event_size(self):
+    """Scalar `int32` tensor: the number of classes."""
+    return self._event_size
+
+  @property
+  def logits(self):
+    """Vector of coordinatewise logits."""
+    return self._logits
+
+  @property
+  def probs(self):
+    """Vector of coordinatewise probabilities."""
+    return self._probs
+
+  def _batch_shape_tensor(self):
+    return array_ops.identity(self._batch_shape_val)
+
+  def _batch_shape(self):
+    return self.logits.get_shape()[:-1]
+
+  def _event_shape_tensor(self):
+    return constant_op.constant([], dtype=dtypes.int32)
+
+  def _event_shape(self):
+    return tensor_shape.TensorShape([])
+
+  def _sample_n(self, n, seed=None):
+    if self.logits.get_shape().ndims == 2:
+      logits_2d = self.logits
+    else:
+      logits_2d = array_ops.reshape(self.logits, [-1, self.event_size])
+    sample_dtype = dtypes.int64 if self.dtype.size > 4 else dtypes.int32
+    draws = random_ops.multinomial(
+        logits_2d, n, seed=seed, output_dtype=sample_dtype)
+    draws = array_ops.reshape(
+        array_ops.transpose(draws),
+        array_ops.concat([[n], self.batch_shape_tensor()], 0))
+    return math_ops.cast(draws, self.dtype)
+
+  def _cdf(self, k):
+    k = ops.convert_to_tensor(k, name="k")
+    if self.validate_args:
+      k = distribution_util.embed_check_integer_casting_closed(
+          k, target_dtype=dtypes.int32)
+
+    k, probs = _broadcast_cat_event_and_params(
+        k, self.probs, base_dtype=self.dtype.base_dtype)
+
+    # batch-flatten everything in order to use `sequence_mask()`.
+    batch_flattened_probs = array_ops.reshape(probs,
+                                              (-1, self._event_size))
+    batch_flattened_k = array_ops.reshape(k, [-1])
+
+    to_sum_over = array_ops.where(
+        array_ops.sequence_mask(batch_flattened_k, self._event_size),
+        batch_flattened_probs,
+        array_ops.zeros_like(batch_flattened_probs))
+    batch_flattened_cdf = math_ops.reduce_sum(to_sum_over, axis=-1)
+    # Reshape back to the shape of the argument.
+    return array_ops.reshape(batch_flattened_cdf, array_ops.shape(k))
+
+  def _log_prob(self, k):
+    k = ops.convert_to_tensor(k, name="k")
+    if self.validate_args:
+      k = distribution_util.embed_check_integer_casting_closed(
+          k, target_dtype=dtypes.int32)
+    k, logits = _broadcast_cat_event_and_params(
+        k, self.logits, base_dtype=self.dtype.base_dtype)
+
+    # pylint: disable=invalid-unary-operand-type
+    return -nn_ops.sparse_softmax_cross_entropy_with_logits(
+        labels=k,
+        logits=logits)
+
+  def _entropy(self):
+    return -math_ops.reduce_sum(
+        nn_ops.log_softmax(self.logits) * self.probs, axis=-1)
+
+  def _mode(self):
+    ret = math_ops.argmax(self.logits, axis=self._batch_rank)
+    ret = math_ops.cast(ret, self.dtype)
+    ret.set_shape(self.batch_shape)
+    return ret
+
+
+@kullback_leibler.RegisterKL(Categorical, Categorical)
+def _kl_categorical_categorical(a, b, name=None):
+  """Calculate the batched KL divergence KL(a || b) with a and b Categorical.
+
+  Args:
+    a: instance of a Categorical distribution object.
+    b: instance of a Categorical distribution object.
+    name: (optional) Name to use for created operations.
+      default is "kl_categorical_categorical".
+
+  Returns:
+    Batchwise KL(a || b)
+  """
+  with ops.name_scope(name, "kl_categorical_categorical",
+                      values=[a.logits, b.logits]):
+    # sum(probs log(probs / (1 - probs)))
+    delta_log_probs1 = (nn_ops.log_softmax(a.logits) -
+                        nn_ops.log_softmax(b.logits))
+    return math_ops.reduce_sum(nn_ops.softmax(a.logits) * delta_log_probs1,
+                               axis=-1)

tensorflow/python/ops/distributions/dirichlet.py

@@ -0,0 +1,410 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""The Dirichlet distribution class."""
+
+import numpy as np
+
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import special_math_ops
+from tensorflow.python.ops.distributions import distribution
+from tensorflow.python.ops.distributions import kullback_leibler
+from tensorflow.python.ops.distributions import util as distribution_util
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+__all__ = [
+    "Dirichlet",
+]
+
+
+_dirichlet_sample_note = """Note: `value` must be a non-negative tensor with
+dtype `self.dtype` and be in the `(self.event_shape() - 1)`-simplex, i.e.,
+`tf.reduce_sum(value, -1) = 1`. It must have a shape compatible with
+`self.batch_shape() + self.event_shape()`."""
+
+
+@tf_export(v1=["distributions.Dirichlet"])
+class Dirichlet(distribution.Distribution):
+  """Dirichlet distribution.
+
+  The Dirichlet distribution is defined over the
+  [`(k-1)`-simplex](https://en.wikipedia.org/wiki/Simplex) using a positive,
+  length-`k` vector `concentration` (`k > 1`). The Dirichlet is identically the
+  Beta distribution when `k = 2`.
+
+  #### Mathematical Details
+
+  The Dirichlet is a distribution over the open `(k-1)`-simplex, i.e.,
+
+  ```none
+  S^{k-1} = { (x_0, ..., x_{k-1}) in R^k : sum_j x_j = 1 and all_j x_j > 0 }.
+  ```
+
+  The probability density function (pdf) is,
+
+  ```none
+  pdf(x; alpha) = prod_j x_j**(alpha_j - 1) / Z
+  Z = prod_j Gamma(alpha_j) / Gamma(sum_j alpha_j)
+  ```
+
+  where:
+
+  * `x in S^{k-1}`, i.e., the `(k-1)`-simplex,
+  * `concentration = alpha = [alpha_0, ..., alpha_{k-1}]`, `alpha_j > 0`,
+  * `Z` is the normalization constant aka the [multivariate beta function](
+    https://en.wikipedia.org/wiki/Beta_function#Multivariate_beta_function),
+    and,
+  * `Gamma` is the [gamma function](
+    https://en.wikipedia.org/wiki/Gamma_function).
+
+  The `concentration` represents mean total counts of class occurrence, i.e.,
+
+  ```none
+  concentration = alpha = mean * total_concentration
+  ```
+
+  where `mean` in `S^{k-1}` and `total_concentration` is a positive real number
+  representing a mean total count.
+
+  Distribution parameters are automatically broadcast in all functions; see
+  examples for details.
+
+  Warning: Some components of the samples can be zero due to finite precision.
+  This happens more often when some of the concentrations are very small.
+  Make sure to round the samples to `np.finfo(dtype).tiny` before computing the
+  density.
+
+  Samples of this distribution are reparameterized (pathwise differentiable).
+  The derivatives are computed using the approach described in
+  (Figurnov et al., 2018).
+
+  #### Examples
+
+  ```python
+  import tensorflow_probability as tfp
+  tfd = tfp.distributions
+
+  # Create a single trivariate Dirichlet, with the 3rd class being three times
+  # more frequent than the first. I.e., batch_shape=[], event_shape=[3].
+  alpha = [1., 2, 3]
+  dist = tfd.Dirichlet(alpha)
+
+  dist.sample([4, 5])  # shape: [4, 5, 3]
+
+  # x has one sample, one batch, three classes:
+  x = [.2, .3, .5]   # shape: [3]
+  dist.prob(x)       # shape: []
+
+  # x has two samples from one batch:
+  x = [[.1, .4, .5],
+       [.2, .3, .5]]
+  dist.prob(x)         # shape: [2]
+
+  # alpha will be broadcast to shape [5, 7, 3] to match x.
+  x = [[...]]   # shape: [5, 7, 3]
+  dist.prob(x)  # shape: [5, 7]
+  ```
+
+  ```python
+  # Create batch_shape=[2], event_shape=[3]:
+  alpha = [[1., 2, 3],
+           [4, 5, 6]]   # shape: [2, 3]
+  dist = tfd.Dirichlet(alpha)
+
+  dist.sample([4, 5])  # shape: [4, 5, 2, 3]
+
+  x = [.2, .3, .5]
+  # x will be broadcast as [[.2, .3, .5],
+  #                         [.2, .3, .5]],
+  # thus matching batch_shape [2, 3].
+  dist.prob(x)         # shape: [2]
+  ```
+
+  Compute the gradients of samples w.r.t. the parameters:
+
+  ```python
+  alpha = tf.constant([1.0, 2.0, 3.0])
+  dist = tfd.Dirichlet(alpha)
+  samples = dist.sample(5)  # Shape [5, 3]
+  loss = tf.reduce_mean(tf.square(samples))  # Arbitrary loss function
+  # Unbiased stochastic gradients of the loss function
+  grads = tf.gradients(loss, alpha)
+  ```
+
+  References:
+    Implicit Reparameterization Gradients:
+      [Figurnov et al., 2018]
+      (http://papers.nips.cc/paper/7326-implicit-reparameterization-gradients)
+      ([pdf]
+      (http://papers.nips.cc/paper/7326-implicit-reparameterization-gradients.pdf))
+  """
+
+  @deprecation.deprecated(
+      "2019-01-01",
+      "The TensorFlow Distributions library has moved to "
+      "TensorFlow Probability "
+      "(https://github.com/tensorflow/probability). You "
+      "should update all references to use `tfp.distributions` "
+      "instead of `tf.distributions`.",
+      warn_once=True)
+  def __init__(self,
+               concentration,
+               validate_args=False,
+               allow_nan_stats=True,
+               name="Dirichlet"):
+    """Initialize a batch of Dirichlet distributions.
+
+    Args:
+      concentration: Positive floating-point `Tensor` indicating mean number
+        of class occurrences; aka "alpha". Implies `self.dtype`, and
+        `self.batch_shape`, `self.event_shape`, i.e., if
+        `concentration.shape = [N1, N2, ..., Nm, k]` then
+        `batch_shape = [N1, N2, ..., Nm]` and
+        `event_shape = [k]`.
+      validate_args: Python `bool`, default `False`. When `True` distribution
+        parameters are checked for validity despite possibly degrading runtime
+        performance. When `False` invalid inputs may silently render incorrect
+        outputs.
+      allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
+        (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
+        result is undefined. When `False`, an exception is raised if one or
+        more of the statistic's batch members are undefined.
+      name: Python `str` name prefixed to Ops created by this class.
+    """
+    parameters = dict(locals())
+    with ops.name_scope(name, values=[concentration]) as name:
+      self._concentration = self._maybe_assert_valid_concentration(
+          ops.convert_to_tensor(concentration, name="concentration"),
+          validate_args)
+      self._total_concentration = math_ops.reduce_sum(self._concentration, -1)
+    super(Dirichlet, self).__init__(
+        dtype=self._concentration.dtype,
+        validate_args=validate_args,
+        allow_nan_stats=allow_nan_stats,
+        reparameterization_type=distribution.FULLY_REPARAMETERIZED,
+        parameters=parameters,
+        graph_parents=[self._concentration,
+                       self._total_concentration],
+        name=name)
+
+  @property
+  def concentration(self):
+    """Concentration parameter; expected counts for that coordinate."""
+    return self._concentration
+
+  @property
+  def total_concentration(self):
+    """Sum of last dim of concentration parameter."""
+    return self._total_concentration
+
+  def _batch_shape_tensor(self):
+    return array_ops.shape(self.total_concentration)
+
+  def _batch_shape(self):
+    return self.total_concentration.get_shape()
+
+  def _event_shape_tensor(self):
+    return array_ops.shape(self.concentration)[-1:]
+
+  def _event_shape(self):
+    return self.concentration.get_shape().with_rank_at_least(1)[-1:]
+
+  def _sample_n(self, n, seed=None):
+    gamma_sample = random_ops.random_gamma(
+        shape=[n],
+        alpha=self.concentration,
+        dtype=self.dtype,
+        seed=seed)
+    return gamma_sample / math_ops.reduce_sum(gamma_sample, -1, keepdims=True)
+
+  @distribution_util.AppendDocstring(_dirichlet_sample_note)
+  def _log_prob(self, x):
+    return self._log_unnormalized_prob(x) - self._log_normalization()
+
+  @distribution_util.AppendDocstring(_dirichlet_sample_note)
+  def _prob(self, x):
+    return math_ops.exp(self._log_prob(x))
+
+  def _log_unnormalized_prob(self, x):
+    x = self._maybe_assert_valid_sample(x)
+    return math_ops.reduce_sum(math_ops.xlogy(self.concentration - 1., x), -1)
+
+  def _log_normalization(self):
+    return special_math_ops.lbeta(self.concentration)
+
+  def _entropy(self):
+    k = math_ops.cast(self.event_shape_tensor()[0], self.dtype)
+    return (
+        self._log_normalization()
+        + ((self.total_concentration - k)
+           * math_ops.digamma(self.total_concentration))
+        - math_ops.reduce_sum(
+            (self.concentration - 1.) * math_ops.digamma(self.concentration),
+            axis=-1))
+
+  def _mean(self):
+    return self.concentration / self.total_concentration[..., array_ops.newaxis]
+
+  def _covariance(self):
+    x = self._variance_scale_term() * self._mean()
+    # pylint: disable=invalid-unary-operand-type
+    return array_ops.matrix_set_diag(
+        -math_ops.matmul(
+            x[..., array_ops.newaxis],
+            x[..., array_ops.newaxis, :]),  # outer prod
+        self._variance())
+
+  def _variance(self):
+    scale = self._variance_scale_term()
+    x = scale * self._mean()
+    return x * (scale - x)
+
+  def _variance_scale_term(self):
+    """Helper to `_covariance` and `_variance` which computes a shared scale."""
+    return math_ops.rsqrt(1. + self.total_concentration[..., array_ops.newaxis])
+
+  @distribution_util.AppendDocstring(
+      """Note: The mode is undefined when any `concentration <= 1`. If
+      `self.allow_nan_stats` is `True`, `NaN` is used for undefined modes. If
+      `self.allow_nan_stats` is `False` an exception is raised when one or more
+      modes are undefined.""")
+  def _mode(self):
+    k = math_ops.cast(self.event_shape_tensor()[0], self.dtype)
+    mode = (self.concentration - 1.) / (
+        self.total_concentration[..., array_ops.newaxis] - k)
+    if self.allow_nan_stats:
+      nan = array_ops.fill(
+          array_ops.shape(mode),
+          np.array(np.nan, dtype=self.dtype.as_numpy_dtype()),
+          name="nan")
+      return array_ops.where_v2(
+          math_ops.reduce_all(self.concentration > 1., axis=-1), mode, nan)
+    return control_flow_ops.with_dependencies([
+        check_ops.assert_less(
+            array_ops.ones([], self.dtype),
+            self.concentration,
+            message="Mode undefined when any concentration <= 1"),
+    ], mode)
+
+  def _maybe_assert_valid_concentration(self, concentration, validate_args):
+    """Checks the validity of the concentration parameter."""
+    if not validate_args:
+      return concentration
+    return control_flow_ops.with_dependencies([
+        check_ops.assert_positive(
+            concentration,
+            message="Concentration parameter must be positive."),
+        check_ops.assert_rank_at_least(
+            concentration, 1,
+            message="Concentration parameter must have >=1 dimensions."),
+        check_ops.assert_less(
+            1, array_ops.shape(concentration)[-1],
+            message="Concentration parameter must have event_size >= 2."),
+    ], concentration)
+
+  def _maybe_assert_valid_sample(self, x):
+    """Checks the validity of a sample."""
+    if not self.validate_args:
+      return x
+    return control_flow_ops.with_dependencies([
+        check_ops.assert_positive(x, message="samples must be positive"),
+        check_ops.assert_near(
+            array_ops.ones([], dtype=self.dtype),
+            math_ops.reduce_sum(x, -1),
+            message="sample last-dimension must sum to `1`"),
+    ], x)
+
+
+@kullback_leibler.RegisterKL(Dirichlet, Dirichlet)
+def _kl_dirichlet_dirichlet(d1, d2, name=None):
+  """Batchwise KL divergence KL(d1 || d2) with d1 and d2 Dirichlet.
+
+  Args:
+    d1: instance of a Dirichlet distribution object.
+    d2: instance of a Dirichlet distribution object.
+    name: (optional) Name to use for created operations.
+      default is "kl_dirichlet_dirichlet".
+
+  Returns:
+    Batchwise KL(d1 || d2)
+  """
+  with ops.name_scope(name, "kl_dirichlet_dirichlet", values=[
+      d1.concentration, d2.concentration]):
+    # The KL between Dirichlet distributions can be derived as follows. We have
+    #
+    #   Dir(x; a) = 1 / B(a) * prod_i[x[i]^(a[i] - 1)]
+    #
+    # where B(a) is the multivariate Beta function:
+    #
+    #   B(a) = Gamma(a[1]) * ... * Gamma(a[n]) / Gamma(a[1] + ... + a[n])
+    #
+    # The KL is
+    #
+    #   KL(Dir(x; a), Dir(x; b)) = E_Dir(x; a){log(Dir(x; a) / Dir(x; b))}
+    #
+    # so we'll need to know the log density of the Dirichlet. This is
+    #
+    #   log(Dir(x; a)) = sum_i[(a[i] - 1) log(x[i])] - log B(a)
+    #
+    # The only term that matters for the expectations is the log(x[i]). To
+    # compute the expectation of this term over the Dirichlet density, we can
+    # use the following facts about the Dirichlet in exponential family form:
+    #   1. log(x[i]) is a sufficient statistic
+    #   2. expected sufficient statistics (of any exp family distribution) are
+    #      equal to derivatives of the log normalizer with respect to
+    #      corresponding natural parameters: E{T[i](x)} = dA/d(eta[i])
+    #
+    # To proceed, we can rewrite the Dirichlet density in exponential family
+    # form as follows:
+    #
+    #   Dir(x; a) = exp{eta(a) . T(x) - A(a)}
+    #
+    # where '.' is the dot product of vectors eta and T, and A is a scalar:
+    #
+    #   eta[i](a) = a[i] - 1
+    #     T[i](x) = log(x[i])
+    #        A(a) = log B(a)
+    #
+    # Now, we can use fact (2) above to write
+    #
+    #   E_Dir(x; a)[log(x[i])]
+    #       = dA(a) / da[i]
+    #       = d/da[i] log B(a)
+    #       = d/da[i] (sum_j lgamma(a[j])) - lgamma(sum_j a[j])
+    #       = digamma(a[i])) - digamma(sum_j a[j])
+    #
+    # Putting it all together, we have
+    #
+    # KL[Dir(x; a) || Dir(x; b)]
+    #     = E_Dir(x; a){log(Dir(x; a) / Dir(x; b)}
+    #     = E_Dir(x; a){sum_i[(a[i] - b[i]) log(x[i])} - (lbeta(a) - lbeta(b))
+    #     = sum_i[(a[i] - b[i]) * E_Dir(x; a){log(x[i])}] - lbeta(a) + lbeta(b)
+    #     = sum_i[(a[i] - b[i]) * (digamma(a[i]) - digamma(sum_j a[j]))]
+    #          - lbeta(a) + lbeta(b))
+
+    digamma_sum_d1 = math_ops.digamma(
+        math_ops.reduce_sum(d1.concentration, axis=-1, keepdims=True))
+    digamma_diff = math_ops.digamma(d1.concentration) - digamma_sum_d1
+    concentration_diff = d1.concentration - d2.concentration
+
+    return (math_ops.reduce_sum(concentration_diff * digamma_diff, axis=-1) -
+            special_math_ops.lbeta(d1.concentration) +
+            special_math_ops.lbeta(d2.concentration))

tensorflow/python/ops/distributions/dirichlet_multinomial.py

@@ -0,0 +1,353 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""The DirichletMultinomial distribution class."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import special_math_ops
+from tensorflow.python.ops.distributions import distribution
+from tensorflow.python.ops.distributions import util as distribution_util
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+__all__ = [
+    "DirichletMultinomial",
+]
+
+
+_dirichlet_multinomial_sample_note = """For each batch of counts,
+`value = [n_0, ..., n_{K-1}]`, `P[value]` is the probability that after
+sampling `self.total_count` draws from this Dirichlet-Multinomial distribution,
+the number of draws falling in class `j` is `n_j`. Since this definition is
+[exchangeable](https://en.wikipedia.org/wiki/Exchangeable_random_variables);
+different sequences have the same counts so the probability includes a
+combinatorial coefficient.
+
+Note: `value` must be a non-negative tensor with dtype `self.dtype`, have no
+fractional components, and such that
+`tf.reduce_sum(value, -1) = self.total_count`. Its shape must be broadcastable
+with `self.concentration` and `self.total_count`."""
+
+
+@tf_export(v1=["distributions.DirichletMultinomial"])
+class DirichletMultinomial(distribution.Distribution):
+  """Dirichlet-Multinomial compound distribution.
+
+  The Dirichlet-Multinomial distribution is parameterized by a (batch of)
+  length-`K` `concentration` vectors (`K > 1`) and a `total_count` number of
+  trials, i.e., the number of trials per draw from the DirichletMultinomial. It
+  is defined over a (batch of) length-`K` vector `counts` such that
+  `tf.reduce_sum(counts, -1) = total_count`. The Dirichlet-Multinomial is
+  identically the Beta-Binomial distribution when `K = 2`.
+
+  #### Mathematical Details
+
+  The Dirichlet-Multinomial is a distribution over `K`-class counts, i.e., a
+  length-`K` vector of non-negative integer `counts = n = [n_0, ..., n_{K-1}]`.
+
+  The probability mass function (pmf) is,
+
+  ```none
+  pmf(n; alpha, N) = Beta(alpha + n) / (prod_j n_j!) / Z
+  Z = Beta(alpha) / N!
+  ```
+
+  where:
+
+  * `concentration = alpha = [alpha_0, ..., alpha_{K-1}]`, `alpha_j > 0`,
+  * `total_count = N`, `N` a positive integer,
+  * `N!` is `N` factorial, and,
+  * `Beta(x) = prod_j Gamma(x_j) / Gamma(sum_j x_j)` is the
+    [multivariate beta function](
+    https://en.wikipedia.org/wiki/Beta_function#Multivariate_beta_function),
+    and,
+  * `Gamma` is the [gamma function](
+    https://en.wikipedia.org/wiki/Gamma_function).
+
+  Dirichlet-Multinomial is a [compound distribution](
+  https://en.wikipedia.org/wiki/Compound_probability_distribution), i.e., its
+  samples are generated as follows.
+
+    1. Choose class probabilities:
+       `probs = [p_0,...,p_{K-1}] ~ Dir(concentration)`
+    2. Draw integers:
+       `counts = [n_0,...,n_{K-1}] ~ Multinomial(total_count, probs)`
+
+  The last `concentration` dimension parametrizes a single Dirichlet-Multinomial
+  distribution. When calling distribution functions (e.g., `dist.prob(counts)`),
+  `concentration`, `total_count` and `counts` are broadcast to the same shape.
+  The last dimension of `counts` corresponds single Dirichlet-Multinomial
+  distributions.
+
+  Distribution parameters are automatically broadcast in all functions; see
+  examples for details.
+
+  #### Pitfalls
+
+  The number of classes, `K`, must not exceed:
+  - the largest integer representable by `self.dtype`, i.e.,
+    `2**(mantissa_bits+1)` (IEE754),
+  - the maximum `Tensor` index, i.e., `2**31-1`.
+
+  In other words,
+
+  ```python
+  K <= min(2**31-1, {
+    tf.float16: 2**11,
+    tf.float32: 2**24,
+    tf.float64: 2**53 }[param.dtype])
+  ```
+
+  Note: This condition is validated only when `self.validate_args = True`.
+
+  #### Examples
+
+  ```python
+  alpha = [1., 2., 3.]
+  n = 2.
+  dist = DirichletMultinomial(n, alpha)
+  ```
+
+  Creates a 3-class distribution, with the 3rd class is most likely to be
+  drawn.
+  The distribution functions can be evaluated on counts.
+
+  ```python
+  # counts same shape as alpha.
+  counts = [0., 0., 2.]
+  dist.prob(counts)  # Shape []
+
+  # alpha will be broadcast to [[1., 2., 3.], [1., 2., 3.]] to match counts.
+  counts = [[1., 1., 0.], [1., 0., 1.]]
+  dist.prob(counts)  # Shape [2]
+
+  # alpha will be broadcast to shape [5, 7, 3] to match counts.
+  counts = [[...]]  # Shape [5, 7, 3]
+  dist.prob(counts)  # Shape [5, 7]
+  ```
+
+  Creates a 2-batch of 3-class distributions.
+
+  ```python
+  alpha = [[1., 2., 3.], [4., 5., 6.]]  # Shape [2, 3]
+  n = [3., 3.]
+  dist = DirichletMultinomial(n, alpha)
+
+  # counts will be broadcast to [[2., 1., 0.], [2., 1., 0.]] to match alpha.
+  counts = [2., 1., 0.]
+  dist.prob(counts)  # Shape [2]
+  ```
+
+  """
+
+  # TODO(b/27419586) Change docstring for dtype of concentration once int
+  # allowed.
+  @deprecation.deprecated(
+      "2019-01-01",
+      "The TensorFlow Distributions library has moved to "
+      "TensorFlow Probability "
+      "(https://github.com/tensorflow/probability). You "
+      "should update all references to use `tfp.distributions` "
+      "instead of `tf.distributions`.",
+      warn_once=True)
+  def __init__(self,
+               total_count,
+               concentration,
+               validate_args=False,
+               allow_nan_stats=True,
+               name="DirichletMultinomial"):
+    """Initialize a batch of DirichletMultinomial distributions.
+
+    Args:
+      total_count:  Non-negative floating point tensor, whose dtype is the same
+        as `concentration`. The shape is broadcastable to `[N1,..., Nm]` with
+        `m >= 0`. Defines this as a batch of `N1 x ... x Nm` different
+        Dirichlet multinomial distributions. Its components should be equal to
+        integer values.
+      concentration: Positive floating point tensor, whose dtype is the
+        same as `n` with shape broadcastable to `[N1,..., Nm, K]` `m >= 0`.
+        Defines this as a batch of `N1 x ... x Nm` different `K` class Dirichlet
+        multinomial distributions.
+      validate_args: Python `bool`, default `False`. When `True` distribution
+        parameters are checked for validity despite possibly degrading runtime
+        performance. When `False` invalid inputs may silently render incorrect
+        outputs.
+      allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
+        (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
+        result is undefined. When `False`, an exception is raised if one or
+        more of the statistic's batch members are undefined.
+      name: Python `str` name prefixed to Ops created by this class.
+    """
+    parameters = dict(locals())
+    with ops.name_scope(name, values=[total_count, concentration]) as name:
+      # Broadcasting works because:
+      # * The broadcasting convention is to prepend dimensions of size [1], and
+      #   we use the last dimension for the distribution, whereas
+      #   the batch dimensions are the leading dimensions, which forces the
+      #   distribution dimension to be defined explicitly (i.e. it cannot be
+      #   created automatically by prepending). This forces enough explicitness.
+      # * All calls involving `counts` eventually require a broadcast between
+      #  `counts` and concentration.
+      self._total_count = ops.convert_to_tensor(total_count, name="total_count")
+      if validate_args:
+        self._total_count = (
+            distribution_util.embed_check_nonnegative_integer_form(
+                self._total_count))
+      self._concentration = self._maybe_assert_valid_concentration(
+          ops.convert_to_tensor(concentration,
+                                name="concentration"),
+          validate_args)
+      self._total_concentration = math_ops.reduce_sum(self._concentration, -1)
+    super(DirichletMultinomial, self).__init__(
+        dtype=self._concentration.dtype,
+        validate_args=validate_args,
+        allow_nan_stats=allow_nan_stats,
+        reparameterization_type=distribution.NOT_REPARAMETERIZED,
+        parameters=parameters,
+        graph_parents=[self._total_count,
+                       self._concentration],
+        name=name)
+
+  @property
+  def total_count(self):
+    """Number of trials used to construct a sample."""
+    return self._total_count
+
+  @property
+  def concentration(self):
+    """Concentration parameter; expected prior counts for that coordinate."""
+    return self._concentration
+
+  @property
+  def total_concentration(self):
+    """Sum of last dim of concentration parameter."""
+    return self._total_concentration
+
+  def _batch_shape_tensor(self):
+    return array_ops.shape(self.total_concentration)
+
+  def _batch_shape(self):
+    return self.total_concentration.get_shape()
+
+  def _event_shape_tensor(self):
+    return array_ops.shape(self.concentration)[-1:]
+
+  def _event_shape(self):
+    # Event shape depends only on total_concentration, not "n".
+    return self.concentration.get_shape().with_rank_at_least(1)[-1:]
+
+  def _sample_n(self, n, seed=None):
+    n_draws = math_ops.cast(self.total_count, dtype=dtypes.int32)
+    k = self.event_shape_tensor()[0]
+    unnormalized_logits = array_ops.reshape(
+        math_ops.log(random_ops.random_gamma(
+            shape=[n],
+            alpha=self.concentration,
+            dtype=self.dtype,
+            seed=seed)),
+        shape=[-1, k])
+    draws = random_ops.multinomial(
+        logits=unnormalized_logits,
+        num_samples=n_draws,
+        seed=distribution_util.gen_new_seed(seed, salt="dirichlet_multinomial"))
+    x = math_ops.reduce_sum(array_ops.one_hot(draws, depth=k), -2)
+    final_shape = array_ops.concat([[n], self.batch_shape_tensor(), [k]], 0)
+    x = array_ops.reshape(x, final_shape)
+    return math_ops.cast(x, self.dtype)
+
+  @distribution_util.AppendDocstring(_dirichlet_multinomial_sample_note)
+  def _log_prob(self, counts):
+    counts = self._maybe_assert_valid_sample(counts)
+    ordered_prob = (
+        special_math_ops.lbeta(self.concentration + counts)
+        - special_math_ops.lbeta(self.concentration))
+    return ordered_prob + distribution_util.log_combinations(
+        self.total_count, counts)
+
+  @distribution_util.AppendDocstring(_dirichlet_multinomial_sample_note)
+  def _prob(self, counts):
+    return math_ops.exp(self._log_prob(counts))
+
+  def _mean(self):
+    return self.total_count * (self.concentration /
+                               self.total_concentration[..., array_ops.newaxis])
+
+  @distribution_util.AppendDocstring(
+      """The covariance for each batch member is defined as the following:
+
+      ```none
+      Var(X_j) = n * alpha_j / alpha_0 * (1 - alpha_j / alpha_0) *
+      (n + alpha_0) / (1 + alpha_0)
+      ```
+
+      where `concentration = alpha` and
+      `total_concentration = alpha_0 = sum_j alpha_j`.
+
+      The covariance between elements in a batch is defined as:
+
+      ```none
+      Cov(X_i, X_j) = -n * alpha_i * alpha_j / alpha_0 ** 2 *
+      (n + alpha_0) / (1 + alpha_0)
+      ```
+      """)
+  def _covariance(self):
+    x = self._variance_scale_term() * self._mean()
+    # pylint: disable=invalid-unary-operand-type
+    return array_ops.matrix_set_diag(
+        -math_ops.matmul(
+            x[..., array_ops.newaxis],
+            x[..., array_ops.newaxis, :]),  # outer prod
+        self._variance())
+
+  def _variance(self):
+    scale = self._variance_scale_term()
+    x = scale * self._mean()
+    return x * (self.total_count * scale - x)
+
+  def _variance_scale_term(self):
+    """Helper to `_covariance` and `_variance` which computes a shared scale."""
+    # We must take care to expand back the last dim whenever we use the
+    # total_concentration.
+    c0 = self.total_concentration[..., array_ops.newaxis]
+    return math_ops.sqrt((1. + c0 / self.total_count) / (1. + c0))
+
+  def _maybe_assert_valid_concentration(self, concentration, validate_args):
+    """Checks the validity of the concentration parameter."""
+    if not validate_args:
+      return concentration
+    concentration = distribution_util.embed_check_categorical_event_shape(
+        concentration)
+    return control_flow_ops.with_dependencies([
+        check_ops.assert_positive(
+            concentration,
+            message="Concentration parameter must be positive."),
+    ], concentration)
+
+  def _maybe_assert_valid_sample(self, counts):
+    """Check counts for proper shape, values, then return tensor version."""
+    if not self.validate_args:
+      return counts
+    counts = distribution_util.embed_check_nonnegative_integer_form(counts)
+    return control_flow_ops.with_dependencies([
+        check_ops.assert_equal(
+            self.total_count, math_ops.reduce_sum(counts, -1),
+            message="counts last-dimension must sum to `self.total_count`"),
+    ], counts)

tensorflow/python/ops/distributions/distributions.py

@@ -0,0 +1,36 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Core module for TensorFlow distribution objects and helpers."""
+from tensorflow.python.util import deprecation
+
+
+# pylint: disable=wildcard-import,unused-import,g-import-not-at-top
+with deprecation.silence():
+  from tensorflow.python.ops.distributions.bernoulli import Bernoulli
+  from tensorflow.python.ops.distributions.beta import Beta
+  from tensorflow.python.ops.distributions.categorical import Categorical
+  from tensorflow.python.ops.distributions.dirichlet import Dirichlet
+  from tensorflow.python.ops.distributions.dirichlet_multinomial import DirichletMultinomial
+  from tensorflow.python.ops.distributions.distribution import *
+  from tensorflow.python.ops.distributions.exponential import Exponential
+  from tensorflow.python.ops.distributions.gamma import Gamma
+  from tensorflow.python.ops.distributions.kullback_leibler import *
+  from tensorflow.python.ops.distributions.laplace import Laplace
+  from tensorflow.python.ops.distributions.multinomial import Multinomial
+  from tensorflow.python.ops.distributions.normal import Normal
+  from tensorflow.python.ops.distributions.student_t import StudentT
+  from tensorflow.python.ops.distributions.uniform import Uniform
+# pylint: enable=wildcard-import,unused-import
+del deprecation

tensorflow/python/ops/distributions/exponential.py

@@ -0,0 +1,162 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""The Exponential distribution class."""
+
+import numpy as np
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops.distributions import gamma
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+__all__ = [
+    "Exponential",
+    "ExponentialWithSoftplusRate",
+]
+
+
+@tf_export(v1=["distributions.Exponential"])
+class Exponential(gamma.Gamma):
+  """Exponential distribution.
+
+  The Exponential distribution is parameterized by an event `rate` parameter.
+
+  #### Mathematical Details
+
+  The probability density function (pdf) is,
+
+  ```none
+  pdf(x; lambda, x > 0) = exp(-lambda x) / Z
+  Z = 1 / lambda
+  ```
+
+  where `rate = lambda` and `Z` is the normalizaing constant.
+
+  The Exponential distribution is a special case of the Gamma distribution,
+  i.e.,
+
+  ```python
+  Exponential(rate) = Gamma(concentration=1., rate)
+  ```
+
+  The Exponential distribution uses a `rate` parameter, or "inverse scale",
+  which can be intuited as,
+
+  ```none
+  X ~ Exponential(rate=1)
+  Y = X / rate
+  ```
+
+  """
+
+  @deprecation.deprecated(
+      "2019-01-01",
+      "The TensorFlow Distributions library has moved to "
+      "TensorFlow Probability "
+      "(https://github.com/tensorflow/probability). You "
+      "should update all references to use `tfp.distributions` "
+      "instead of `tf.distributions`.",
+      warn_once=True)
+  def __init__(self,
+               rate,
+               validate_args=False,
+               allow_nan_stats=True,
+               name="Exponential"):
+    """Construct Exponential distribution with parameter `rate`.
+
+    Args:
+      rate: Floating point tensor, equivalent to `1 / mean`. Must contain only
+        positive values.
+      validate_args: Python `bool`, default `False`. When `True` distribution
+        parameters are checked for validity despite possibly degrading runtime
+        performance. When `False` invalid inputs may silently render incorrect
+        outputs.
+      allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
+        (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
+        result is undefined. When `False`, an exception is raised if one or
+        more of the statistic's batch members are undefined.
+      name: Python `str` name prefixed to Ops created by this class.
+    """
+    parameters = dict(locals())
+    # Even though all statistics of are defined for valid inputs, this is not
+    # true in the parent class "Gamma."  Therefore, passing
+    # allow_nan_stats=True
+    # through to the parent class results in unnecessary asserts.
+    with ops.name_scope(name, values=[rate]) as name:
+      self._rate = ops.convert_to_tensor(rate, name="rate")
+    super(Exponential, self).__init__(
+        concentration=array_ops.ones([], dtype=self._rate.dtype),
+        rate=self._rate,
+        allow_nan_stats=allow_nan_stats,
+        validate_args=validate_args,
+        name=name)
+    self._parameters = parameters
+    self._graph_parents += [self._rate]
+
+  @staticmethod
+  def _param_shapes(sample_shape):
+    return {"rate": ops.convert_to_tensor(sample_shape, dtype=dtypes.int32)}
+
+  @property
+  def rate(self):
+    return self._rate
+
+  def _log_survival_function(self, value):
+    return self._log_prob(value) - math_ops.log(self._rate)
+
+  def _sample_n(self, n, seed=None):
+    shape = array_ops.concat([[n], array_ops.shape(self._rate)], 0)
+    # Uniform variates must be sampled from the open-interval `(0, 1)` rather
+    # than `[0, 1)`. To do so, we use `np.finfo(self.dtype.as_numpy_dtype).tiny`
+    # because it is the smallest, positive, "normal" number. A "normal" number
+    # is such that the mantissa has an implicit leading 1. Normal, positive
+    # numbers x, y have the reasonable property that, `x + y >= max(x, y)`. In
+    # this case, a subnormal number (i.e., np.nextafter) can cause us to sample
+    # 0.
+    sampled = random_ops.random_uniform(
+        shape,
+        minval=np.finfo(self.dtype.as_numpy_dtype).tiny,
+        maxval=1.,
+        seed=seed,
+        dtype=self.dtype)
+    return -math_ops.log(sampled) / self._rate
+
+
+class ExponentialWithSoftplusRate(Exponential):
+  """Exponential with softplus transform on `rate`."""
+
+  @deprecation.deprecated(
+      "2019-01-01",
+      "Use `tfd.Exponential(tf.nn.softplus(rate)).",
+      warn_once=True)
+  def __init__(self,
+               rate,
+               validate_args=False,
+               allow_nan_stats=True,
+               name="ExponentialWithSoftplusRate"):
+    parameters = dict(locals())
+    with ops.name_scope(name, values=[rate]) as name:
+      super(ExponentialWithSoftplusRate, self).__init__(
+          rate=nn.softplus(rate, name="softplus_rate"),
+          validate_args=validate_args,
+          allow_nan_stats=allow_nan_stats,
+          name=name)
+    self._parameters = parameters

tensorflow/python/ops/distributions/gamma.py

@@ -0,0 +1,338 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""The Gamma distribution class."""
+
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops.distributions import distribution
+from tensorflow.python.ops.distributions import kullback_leibler
+from tensorflow.python.ops.distributions import util as distribution_util
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+__all__ = [
+    "Gamma",
+    "GammaWithSoftplusConcentrationRate",
+]
+
+
+@tf_export(v1=["distributions.Gamma"])
+class Gamma(distribution.Distribution):
+  """Gamma distribution.
+
+  The Gamma distribution is defined over positive real numbers using
+  parameters `concentration` (aka "alpha") and `rate` (aka "beta").
+
+  #### Mathematical Details
+
+  The probability density function (pdf) is,
+
+  ```none
+  pdf(x; alpha, beta, x > 0) = x**(alpha - 1) exp(-x beta) / Z
+  Z = Gamma(alpha) beta**(-alpha)
+  ```
+
+  where:
+
+  * `concentration = alpha`, `alpha > 0`,
+  * `rate = beta`, `beta > 0`,
+  * `Z` is the normalizing constant, and,
+  * `Gamma` is the [gamma function](
+    https://en.wikipedia.org/wiki/Gamma_function).
+
+  The cumulative density function (cdf) is,
+
+  ```none
+  cdf(x; alpha, beta, x > 0) = GammaInc(alpha, beta x) / Gamma(alpha)
+  ```
+
+  where `GammaInc` is the [lower incomplete Gamma function](
+  https://en.wikipedia.org/wiki/Incomplete_gamma_function).
+
+  The parameters can be intuited via their relationship to mean and stddev,
+
+  ```none
+  concentration = alpha = (mean / stddev)**2
+  rate = beta = mean / stddev**2 = concentration / mean
+  ```
+
+  Distribution parameters are automatically broadcast in all functions; see
+  examples for details.
+
+  Warning: The samples of this distribution are always non-negative. However,
+  the samples that are smaller than `np.finfo(dtype).tiny` are rounded
+  to this value, so it appears more often than it should.
+  This should only be noticeable when the `concentration` is very small, or the
+  `rate` is very large. See note in `tf.random.gamma` docstring.
+
+  Samples of this distribution are reparameterized (pathwise differentiable).
+  The derivatives are computed using the approach described in
+  (Figurnov et al., 2018).
+
+  #### Examples
+
+  ```python
+  import tensorflow_probability as tfp
+  tfd = tfp.distributions
+
+  dist = tfd.Gamma(concentration=3.0, rate=2.0)
+  dist2 = tfd.Gamma(concentration=[3.0, 4.0], rate=[2.0, 3.0])
+  ```
+
+  Compute the gradients of samples w.r.t. the parameters:
+
+  ```python
+  concentration = tf.constant(3.0)
+  rate = tf.constant(2.0)
+  dist = tfd.Gamma(concentration, rate)
+  samples = dist.sample(5)  # Shape [5]
+  loss = tf.reduce_mean(tf.square(samples))  # Arbitrary loss function
+  # Unbiased stochastic gradients of the loss function
+  grads = tf.gradients(loss, [concentration, rate])
+  ```
+
+  References:
+    Implicit Reparameterization Gradients:
+      [Figurnov et al., 2018]
+      (http://papers.nips.cc/paper/7326-implicit-reparameterization-gradients)
+      ([pdf](http://papers.nips.cc/paper/7326-implicit-reparameterization-gradients.pdf))
+  """
+
+  @deprecation.deprecated(
+      "2019-01-01",
+      "The TensorFlow Distributions library has moved to "
+      "TensorFlow Probability "
+      "(https://github.com/tensorflow/probability). You "
+      "should update all references to use `tfp.distributions` "
+      "instead of `tf.distributions`.",
+      warn_once=True)
+  def __init__(self,
+               concentration,
+               rate,
+               validate_args=False,
+               allow_nan_stats=True,
+               name="Gamma"):
+    """Construct Gamma with `concentration` and `rate` parameters.
+
+    The parameters `concentration` and `rate` must be shaped in a way that
+    supports broadcasting (e.g. `concentration + rate` is a valid operation).
+
+    Args:
+      concentration: Floating point tensor, the concentration params of the
+        distribution(s). Must contain only positive values.
+      rate: Floating point tensor, the inverse scale params of the
+        distribution(s). Must contain only positive values.
+      validate_args: Python `bool`, default `False`. When `True` distribution
+        parameters are checked for validity despite possibly degrading runtime
+        performance. When `False` invalid inputs may silently render incorrect
+        outputs.
+      allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
+        (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
+        result is undefined. When `False`, an exception is raised if one or
+        more of the statistic's batch members are undefined.
+      name: Python `str` name prefixed to Ops created by this class.
+
+    Raises:
+      TypeError: if `concentration` and `rate` are different dtypes.
+    """
+    parameters = dict(locals())
+    with ops.name_scope(name, values=[concentration, rate]) as name:
+      with ops.control_dependencies([
+          check_ops.assert_positive(concentration),
+          check_ops.assert_positive(rate),
+      ] if validate_args else []):
+        self._concentration = array_ops.identity(
+            concentration, name="concentration")
+        self._rate = array_ops.identity(rate, name="rate")
+        check_ops.assert_same_float_dtype(
+            [self._concentration, self._rate])
+    super(Gamma, self).__init__(
+        dtype=self._concentration.dtype,
+        validate_args=validate_args,
+        allow_nan_stats=allow_nan_stats,
+        reparameterization_type=distribution.FULLY_REPARAMETERIZED,
+        parameters=parameters,
+        graph_parents=[self._concentration,
+                       self._rate],
+        name=name)
+
+  @staticmethod
+  def _param_shapes(sample_shape):
+    return dict(
+        zip(("concentration", "rate"), ([ops.convert_to_tensor(
+            sample_shape, dtype=dtypes.int32)] * 2)))
+
+  @property
+  def concentration(self):
+    """Concentration parameter."""
+    return self._concentration
+
+  @property
+  def rate(self):
+    """Rate parameter."""
+    return self._rate
+
+  def _batch_shape_tensor(self):
+    return array_ops.broadcast_dynamic_shape(
+        array_ops.shape(self.concentration),
+        array_ops.shape(self.rate))
+
+  def _batch_shape(self):
+    return array_ops.broadcast_static_shape(
+        self.concentration.get_shape(),
+        self.rate.get_shape())
+
+  def _event_shape_tensor(self):
+    return constant_op.constant([], dtype=dtypes.int32)
+
+  def _event_shape(self):
+    return tensor_shape.TensorShape([])
+
+  @distribution_util.AppendDocstring(
+      """Note: See `tf.random.gamma` docstring for sampling details and
+      caveats.""")
+  def _sample_n(self, n, seed=None):
+    return random_ops.random_gamma(
+        shape=[n],
+        alpha=self.concentration,
+        beta=self.rate,
+        dtype=self.dtype,
+        seed=seed)
+
+  def _log_prob(self, x):
+    return self._log_unnormalized_prob(x) - self._log_normalization()
+
+  def _cdf(self, x):
+    x = self._maybe_assert_valid_sample(x)
+    # Note that igamma returns the regularized incomplete gamma function,
+    # which is what we want for the CDF.
+    return math_ops.igamma(self.concentration, self.rate * x)
+
+  def _log_unnormalized_prob(self, x):
+    x = self._maybe_assert_valid_sample(x)
+    return math_ops.xlogy(self.concentration - 1., x) - self.rate * x
+
+  def _log_normalization(self):
+    return (math_ops.lgamma(self.concentration)
+            - self.concentration * math_ops.log(self.rate))
+
+  def _entropy(self):
+    return (self.concentration
+            - math_ops.log(self.rate)
+            + math_ops.lgamma(self.concentration)
+            + ((1. - self.concentration) *
+               math_ops.digamma(self.concentration)))
+
+  def _mean(self):
+    return self.concentration / self.rate
+
+  def _variance(self):
+    return self.concentration / math_ops.square(self.rate)
+
+  def _stddev(self):
+    return math_ops.sqrt(self.concentration) / self.rate
+
+  @distribution_util.AppendDocstring(
+      """The mode of a gamma distribution is `(shape - 1) / rate` when
+      `shape > 1`, and `NaN` otherwise. If `self.allow_nan_stats` is `False`,
+      an exception will be raised rather than returning `NaN`.""")
+  def _mode(self):
+    mode = (self.concentration - 1.) / self.rate
+    if self.allow_nan_stats:
+      nan = array_ops.fill(
+          self.batch_shape_tensor(),
+          np.array(np.nan, dtype=self.dtype.as_numpy_dtype()),
+          name="nan")
+      return array_ops.where_v2(self.concentration > 1., mode, nan)
+    else:
+      return control_flow_ops.with_dependencies([
+          check_ops.assert_less(
+              array_ops.ones([], self.dtype),
+              self.concentration,
+              message="mode not defined when any concentration <= 1"),
+          ], mode)
+
+  def _maybe_assert_valid_sample(self, x):
+    check_ops.assert_same_float_dtype(tensors=[x], dtype=self.dtype)
+    if not self.validate_args:
+      return x
+    return control_flow_ops.with_dependencies([
+        check_ops.assert_positive(x),
+    ], x)
+
+
+class GammaWithSoftplusConcentrationRate(Gamma):
+  """`Gamma` with softplus of `concentration` and `rate`."""
+
+  @deprecation.deprecated(
+      "2019-01-01",
+      "Use `tfd.Gamma(tf.nn.softplus(concentration), "
+      "tf.nn.softplus(rate))` instead.",
+      warn_once=True)
+  def __init__(self,
+               concentration,
+               rate,
+               validate_args=False,
+               allow_nan_stats=True,
+               name="GammaWithSoftplusConcentrationRate"):
+    parameters = dict(locals())
+    with ops.name_scope(name, values=[concentration, rate]) as name:
+      super(GammaWithSoftplusConcentrationRate, self).__init__(
+          concentration=nn.softplus(concentration,
+                                    name="softplus_concentration"),
+          rate=nn.softplus(rate, name="softplus_rate"),
+          validate_args=validate_args,
+          allow_nan_stats=allow_nan_stats,
+          name=name)
+    self._parameters = parameters
+
+
+@kullback_leibler.RegisterKL(Gamma, Gamma)
+def _kl_gamma_gamma(g0, g1, name=None):
+  """Calculate the batched KL divergence KL(g0 || g1) with g0 and g1 Gamma.
+
+  Args:
+    g0: instance of a Gamma distribution object.
+    g1: instance of a Gamma distribution object.
+    name: (optional) Name to use for created operations.
+      Default is "kl_gamma_gamma".
+
+  Returns:
+    kl_gamma_gamma: `Tensor`. The batchwise KL(g0 || g1).
+  """
+  with ops.name_scope(name, "kl_gamma_gamma", values=[
+      g0.concentration, g0.rate, g1.concentration, g1.rate]):
+    # Result from:
+    #   http://www.fil.ion.ucl.ac.uk/~wpenny/publications/densities.ps
+    # For derivation see:
+    #   http://stats.stackexchange.com/questions/11646/kullback-leibler-divergence-between-two-gamma-distributions   pylint: disable=line-too-long
+    return (((g0.concentration - g1.concentration)
+             * math_ops.digamma(g0.concentration))
+            + math_ops.lgamma(g1.concentration)
+            - math_ops.lgamma(g0.concentration)
+            + g1.concentration * math_ops.log(g0.rate)
+            - g1.concentration * math_ops.log(g1.rate)
+            + g0.concentration * (g1.rate / g0.rate - 1.))

tensorflow/python/ops/distributions/identity_bijector.py

@@ -0,0 +1,68 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Identity bijector."""
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops.distributions import bijector
+from tensorflow.python.util import deprecation
+
+
+__all__ = [
+    "Identity",
+]
+
+
+class Identity(bijector.Bijector):
+  """Compute Y = g(X) = X.
+
+    Example Use:
+
+    ```python
+    # Create the Y=g(X)=X transform which is intended for Tensors with 1 batch
+    # ndim and 1 event ndim (i.e., vector of vectors).
+    identity = Identity()
+    x = [[1., 2],
+         [3, 4]]
+    x == identity.forward(x) == identity.inverse(x)
+    ```
+
+  """
+
+  @deprecation.deprecated(
+      "2019-01-01",
+      "The TensorFlow Distributions library has moved to "
+      "TensorFlow Probability "
+      "(https://github.com/tensorflow/probability). You "
+      "should update all references to use `tfp.distributions` "
+      "instead of `tf.distributions`.",
+      warn_once=True)
+  def __init__(self, validate_args=False, name="identity"):
+    super(Identity, self).__init__(
+        forward_min_event_ndims=0,
+        is_constant_jacobian=True,
+        validate_args=validate_args,
+        name=name)
+
+  def _forward(self, x):
+    return x
+
+  def _inverse(self, y):
+    return y
+
+  def _inverse_log_det_jacobian(self, y):
+    return constant_op.constant(0., dtype=y.dtype)
+
+  def _forward_log_det_jacobian(self, x):
+    return constant_op.constant(0., dtype=x.dtype)

tensorflow/python/ops/distributions/kullback_leibler.py

@@ -0,0 +1,210 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Registration and usage mechanisms for KL-divergences."""
+
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_assert
+from tensorflow.python.ops import math_ops
+from tensorflow.python.util import deprecation
+from tensorflow.python.util import tf_inspect
+from tensorflow.python.util.tf_export import tf_export
+
+
+_DIVERGENCES = {}
+
+
+__all__ = [
+    "RegisterKL",
+    "kl_divergence",
+]
+
+
+def _registered_kl(type_a, type_b):
+  """Get the KL function registered for classes a and b."""
+  hierarchy_a = tf_inspect.getmro(type_a)
+  hierarchy_b = tf_inspect.getmro(type_b)
+  dist_to_children = None
+  kl_fn = None
+  for mro_to_a, parent_a in enumerate(hierarchy_a):
+    for mro_to_b, parent_b in enumerate(hierarchy_b):
+      candidate_dist = mro_to_a + mro_to_b
+      candidate_kl_fn = _DIVERGENCES.get((parent_a, parent_b), None)
+      if not kl_fn or (candidate_kl_fn and candidate_dist < dist_to_children):
+        dist_to_children = candidate_dist
+        kl_fn = candidate_kl_fn
+  return kl_fn
+
+
+@deprecation.deprecated(
+    "2019-01-01",
+    "The TensorFlow Distributions library has moved to "
+    "TensorFlow Probability "
+    "(https://github.com/tensorflow/probability). You "
+    "should update all references to use `tfp.distributions` "
+    "instead of `tf.distributions`.",
+    warn_once=True)
+@tf_export(v1=["distributions.kl_divergence"])
+def kl_divergence(distribution_a, distribution_b,
+                  allow_nan_stats=True, name=None):
+  """Get the KL-divergence KL(distribution_a || distribution_b).
+
+  If there is no KL method registered specifically for `type(distribution_a)`
+  and `type(distribution_b)`, then the class hierarchies of these types are
+  searched.
+
+  If one KL method is registered between any pairs of classes in these two
+  parent hierarchies, it is used.
+
+  If more than one such registered method exists, the method whose registered
+  classes have the shortest sum MRO paths to the input types is used.
+
+  If more than one such shortest path exists, the first method
+  identified in the search is used (favoring a shorter MRO distance to
+  `type(distribution_a)`).
+
+  Args:
+    distribution_a: The first distribution.
+    distribution_b: The second distribution.
+    allow_nan_stats: Python `bool`, default `True`. When `True`,
+      statistics (e.g., mean, mode, variance) use the value "`NaN`" to
+      indicate the result is undefined. When `False`, an exception is raised
+      if one or more of the statistic's batch members are undefined.
+    name: Python `str` name prefixed to Ops created by this class.
+
+  Returns:
+    A Tensor with the batchwise KL-divergence between `distribution_a`
+    and `distribution_b`.
+
+  Raises:
+    NotImplementedError: If no KL method is defined for distribution types
+      of `distribution_a` and `distribution_b`.
+  """
+  kl_fn = _registered_kl(type(distribution_a), type(distribution_b))
+  if kl_fn is None:
+    raise NotImplementedError(
+        "No KL(distribution_a || distribution_b) registered for distribution_a "
+        "type %s and distribution_b type %s"
+        % (type(distribution_a).__name__, type(distribution_b).__name__))
+
+  with ops.name_scope("KullbackLeibler"):
+    kl_t = kl_fn(distribution_a, distribution_b, name=name)
+    if allow_nan_stats:
+      return kl_t
+
+    # Check KL for NaNs
+    kl_t = array_ops.identity(kl_t, name="kl")
+
+    with ops.control_dependencies([
+        control_flow_assert.Assert(
+            math_ops.logical_not(math_ops.reduce_any(math_ops.is_nan(kl_t))), [
+                "KL calculation between %s and %s returned NaN values "
+                "(and was called with allow_nan_stats=False). Values:" %
+                (distribution_a.name, distribution_b.name), kl_t
+            ])
+    ]):
+      return array_ops.identity(kl_t, name="checked_kl")
+
+
+@deprecation.deprecated(
+    "2019-01-01",
+    "The TensorFlow Distributions library has moved to "
+    "TensorFlow Probability "
+    "(https://github.com/tensorflow/probability). You "
+    "should update all references to use `tfp.distributions` "
+    "instead of `tf.distributions`.",
+    warn_once=True)
+def cross_entropy(ref, other,
+                  allow_nan_stats=True, name=None):
+  """Computes the (Shannon) cross entropy.
+
+  Denote two distributions by `P` (`ref`) and `Q` (`other`). Assuming `P, Q`
+  are absolutely continuous with respect to one another and permit densities
+  `p(x) dr(x)` and `q(x) dr(x)`, (Shanon) cross entropy is defined as:
+
+  ```none
+  H[P, Q] = E_p[-log q(X)] = -int_F p(x) log q(x) dr(x)
+  ```
+
+  where `F` denotes the support of the random variable `X ~ P`.
+
+  Args:
+    ref: `tfd.Distribution` instance.
+    other: `tfd.Distribution` instance.
+    allow_nan_stats: Python `bool`, default `True`. When `True`,
+      statistics (e.g., mean, mode, variance) use the value "`NaN`" to
+      indicate the result is undefined. When `False`, an exception is raised
+      if one or more of the statistic's batch members are undefined.
+    name: Python `str` prepended to names of ops created by this function.
+
+  Returns:
+    cross_entropy: `ref.dtype` `Tensor` with shape `[B1, ..., Bn]`
+      representing `n` different calculations of (Shanon) cross entropy.
+  """
+  with ops.name_scope(name, "cross_entropy"):
+    return ref.entropy() + kl_divergence(
+        ref, other, allow_nan_stats=allow_nan_stats)
+
+
+@tf_export(v1=["distributions.RegisterKL"])
+class RegisterKL:
+  """Decorator to register a KL divergence implementation function.
+
+  Usage:
+
+  @distributions.RegisterKL(distributions.Normal, distributions.Normal)
+  def _kl_normal_mvn(norm_a, norm_b):
+    # Return KL(norm_a || norm_b)
+  """
+
+  @deprecation.deprecated(
+      "2019-01-01",
+      "The TensorFlow Distributions library has moved to "
+      "TensorFlow Probability "
+      "(https://github.com/tensorflow/probability). You "
+      "should update all references to use `tfp.distributions` "
+      "instead of `tf.distributions`.",
+      warn_once=True)
+  def __init__(self, dist_cls_a, dist_cls_b):
+    """Initialize the KL registrar.
+
+    Args:
+      dist_cls_a: the class of the first argument of the KL divergence.
+      dist_cls_b: the class of the second argument of the KL divergence.
+    """
+    self._key = (dist_cls_a, dist_cls_b)
+
+  def __call__(self, kl_fn):
+    """Perform the KL registration.
+
+    Args:
+      kl_fn: The function to use for the KL divergence.
+
+    Returns:
+      kl_fn
+
+    Raises:
+      TypeError: if kl_fn is not a callable.
+      ValueError: if a KL divergence function has already been registered for
+        the given argument classes.
+    """
+    if not callable(kl_fn):
+      raise TypeError("kl_fn must be callable, received: %s" % kl_fn)
+    if self._key in _DIVERGENCES:
+      raise ValueError("KL(%s || %s) has already been registered to: %s"
+                       % (self._key[0].__name__, self._key[1].__name__,
+                          _DIVERGENCES[self._key]))
+    _DIVERGENCES[self._key] = kl_fn
+    return kl_fn

tensorflow/python/ops/distributions/laplace.py

@@ -0,0 +1,238 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""The Laplace distribution class."""
+
+import math
+
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops.distributions import distribution
+from tensorflow.python.ops.distributions import special_math
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+__all__ = [
+    "Laplace",
+    "LaplaceWithSoftplusScale",
+]
+
+
+@tf_export(v1=["distributions.Laplace"])
+class Laplace(distribution.Distribution):
+  """The Laplace distribution with location `loc` and `scale` parameters.
+
+  #### Mathematical details
+
+  The probability density function (pdf) of this distribution is,
+
+  ```none
+  pdf(x; mu, sigma) = exp(-|x - mu| / sigma) / Z
+  Z = 2 sigma
+  ```
+
+  where `loc = mu`, `scale = sigma`, and `Z` is the normalization constant.
+
+  Note that the Laplace distribution can be thought of two exponential
+  distributions spliced together "back-to-back."
+
+  The Lpalce distribution is a member of the [location-scale family](
+  https://en.wikipedia.org/wiki/Location-scale_family), i.e., it can be
+  constructed as,
+
+  ```none
+  X ~ Laplace(loc=0, scale=1)
+  Y = loc + scale * X
+  ```
+
+  """
+
+  @deprecation.deprecated(
+      "2019-01-01",
+      "The TensorFlow Distributions library has moved to "
+      "TensorFlow Probability "
+      "(https://github.com/tensorflow/probability). You "
+      "should update all references to use `tfp.distributions` "
+      "instead of `tf.distributions`.",
+      warn_once=True)
+  def __init__(self,
+               loc,
+               scale,
+               validate_args=False,
+               allow_nan_stats=True,
+               name="Laplace"):
+    """Construct Laplace distribution with parameters `loc` and `scale`.
+
+    The parameters `loc` and `scale` must be shaped in a way that supports
+    broadcasting (e.g., `loc / scale` is a valid operation).
+
+    Args:
+      loc: Floating point tensor which characterizes the location (center)
+        of the distribution.
+      scale: Positive floating point tensor which characterizes the spread of
+        the distribution.
+      validate_args: Python `bool`, default `False`. When `True` distribution
+        parameters are checked for validity despite possibly degrading runtime
+        performance. When `False` invalid inputs may silently render incorrect
+        outputs.
+      allow_nan_stats: Python `bool`, default `True`. When `True`,
+        statistics (e.g., mean, mode, variance) use the value "`NaN`" to
+        indicate the result is undefined. When `False`, an exception is raised
+        if one or more of the statistic's batch members are undefined.
+      name: Python `str` name prefixed to Ops created by this class.
+
+    Raises:
+      TypeError: if `loc` and `scale` are of different dtype.
+    """
+    parameters = dict(locals())
+    with ops.name_scope(name, values=[loc, scale]) as name:
+      with ops.control_dependencies([check_ops.assert_positive(scale)] if
+                                    validate_args else []):
+        self._loc = array_ops.identity(loc, name="loc")
+        self._scale = array_ops.identity(scale, name="scale")
+        check_ops.assert_same_float_dtype([self._loc, self._scale])
+      super(Laplace, self).__init__(
+          dtype=self._loc.dtype,
+          reparameterization_type=distribution.FULLY_REPARAMETERIZED,
+          validate_args=validate_args,
+          allow_nan_stats=allow_nan_stats,
+          parameters=parameters,
+          graph_parents=[self._loc, self._scale],
+          name=name)
+
+  @staticmethod
+  def _param_shapes(sample_shape):
+    return dict(
+        zip(("loc", "scale"), ([ops.convert_to_tensor(
+            sample_shape, dtype=dtypes.int32)] * 2)))
+
+  @property
+  def loc(self):
+    """Distribution parameter for the location."""
+    return self._loc
+
+  @property
+  def scale(self):
+    """Distribution parameter for scale."""
+    return self._scale
+
+  def _batch_shape_tensor(self):
+    return array_ops.broadcast_dynamic_shape(
+        array_ops.shape(self.loc), array_ops.shape(self.scale))
+
+  def _batch_shape(self):
+    return array_ops.broadcast_static_shape(
+        self.loc.get_shape(), self.scale.get_shape())
+
+  def _event_shape_tensor(self):
+    return constant_op.constant([], dtype=dtypes.int32)
+
+  def _event_shape(self):
+    return tensor_shape.TensorShape([])
+
+  def _sample_n(self, n, seed=None):
+    shape = array_ops.concat([[n], self.batch_shape_tensor()], 0)
+    # Uniform variates must be sampled from the open-interval `(-1, 1)` rather
+    # than `[-1, 1)`. In the case of `(0, 1)` we'd use
+    # `np.finfo(self.dtype.as_numpy_dtype).tiny` because it is the smallest,
+    # positive, "normal" number. However, the concept of subnormality exists
+    # only at zero; here we need the smallest usable number larger than -1,
+    # i.e., `-1 + eps/2`.
+    uniform_samples = random_ops.random_uniform(
+        shape=shape,
+        minval=np.nextafter(self.dtype.as_numpy_dtype(-1.),
+                            self.dtype.as_numpy_dtype(0.)),
+        maxval=1.,
+        dtype=self.dtype,
+        seed=seed)
+    return (self.loc - self.scale * math_ops.sign(uniform_samples) *
+            math_ops.log1p(-math_ops.abs(uniform_samples)))
+
+  def _log_prob(self, x):
+    return self._log_unnormalized_prob(x) - self._log_normalization()
+
+  def _prob(self, x):
+    return math_ops.exp(self._log_prob(x))
+
+  def _log_cdf(self, x):
+    return special_math.log_cdf_laplace(self._z(x))
+
+  def _log_survival_function(self, x):
+    return special_math.log_cdf_laplace(-self._z(x))
+
+  def _cdf(self, x):
+    z = self._z(x)
+    return (0.5 + 0.5 * math_ops.sign(z) *
+            (1. - math_ops.exp(-math_ops.abs(z))))
+
+  def _log_unnormalized_prob(self, x):
+    return -math_ops.abs(self._z(x))
+
+  def _log_normalization(self):
+    return math.log(2.) + math_ops.log(self.scale)
+
+  def _entropy(self):
+    # Use broadcasting rules to calculate the full broadcast scale.
+    scale = self.scale + array_ops.zeros_like(self.loc)
+    return math.log(2.) + 1. + math_ops.log(scale)
+
+  def _mean(self):
+    return self.loc + array_ops.zeros_like(self.scale)
+
+  def _stddev(self):
+    return math.sqrt(2.) * self.scale + array_ops.zeros_like(self.loc)
+
+  def _median(self):
+    return self._mean()
+
+  def _mode(self):
+    return self._mean()
+
+  def _z(self, x):
+    return (x - self.loc) / self.scale
+
+
+class LaplaceWithSoftplusScale(Laplace):
+  """Laplace with softplus applied to `scale`."""
+
+  @deprecation.deprecated(
+      "2019-01-01",
+      "Use `tfd.Laplace(loc, tf.nn.softplus(scale)) "
+      "instead.",
+      warn_once=True)
+  def __init__(self,
+               loc,
+               scale,
+               validate_args=False,
+               allow_nan_stats=True,
+               name="LaplaceWithSoftplusScale"):
+    parameters = dict(locals())
+    with ops.name_scope(name, values=[loc, scale]) as name:
+      super(LaplaceWithSoftplusScale, self).__init__(
+          loc=loc,
+          scale=nn.softplus(scale, name="softplus_scale"),
+          validate_args=validate_args,
+          allow_nan_stats=allow_nan_stats,
+          name=name)
+    self._parameters = parameters

tensorflow/python/ops/distributions/multinomial.py

@@ -0,0 +1,314 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""The Multinomial distribution class."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import map_fn
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops.distributions import distribution
+from tensorflow.python.ops.distributions import util as distribution_util
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+__all__ = [
+    "Multinomial",
+]
+
+
+_multinomial_sample_note = """For each batch of counts, `value = [n_0, ...
+,n_{k-1}]`, `P[value]` is the probability that after sampling `self.total_count`
+draws from this Multinomial distribution, the number of draws falling in class
+`j` is `n_j`. Since this definition is [exchangeable](
+https://en.wikipedia.org/wiki/Exchangeable_random_variables); different
+sequences have the same counts so the probability includes a combinatorial
+coefficient.
+
+Note: `value` must be a non-negative tensor with dtype `self.dtype`, have no
+fractional components, and such that
+`tf.reduce_sum(value, -1) = self.total_count`. Its shape must be broadcastable
+with `self.probs` and `self.total_count`."""
+
+
+@tf_export(v1=["distributions.Multinomial"])
+class Multinomial(distribution.Distribution):
+  """Multinomial distribution.
+
+  This Multinomial distribution is parameterized by `probs`, a (batch of)
+  length-`K` `prob` (probability) vectors (`K > 1`) such that
+  `tf.reduce_sum(probs, -1) = 1`, and a `total_count` number of trials, i.e.,
+  the number of trials per draw from the Multinomial. It is defined over a
+  (batch of) length-`K` vector `counts` such that
+  `tf.reduce_sum(counts, -1) = total_count`. The Multinomial is identically the
+  Binomial distribution when `K = 2`.
+
+  #### Mathematical Details
+
+  The Multinomial is a distribution over `K`-class counts, i.e., a length-`K`
+  vector of non-negative integer `counts = n = [n_0, ..., n_{K-1}]`.
+
+  The probability mass function (pmf) is,
+
+  ```none
+  pmf(n; pi, N) = prod_j (pi_j)**n_j / Z
+  Z = (prod_j n_j!) / N!
+  ```
+
+  where:
+  * `probs = pi = [pi_0, ..., pi_{K-1}]`, `pi_j > 0`, `sum_j pi_j = 1`,
+  * `total_count = N`, `N` a positive integer,
+  * `Z` is the normalization constant, and,
+  * `N!` denotes `N` factorial.
+
+  Distribution parameters are automatically broadcast in all functions; see
+  examples for details.
+
+  #### Pitfalls
+
+  The number of classes, `K`, must not exceed:
+  - the largest integer representable by `self.dtype`, i.e.,
+    `2**(mantissa_bits+1)` (IEE754),
+  - the maximum `Tensor` index, i.e., `2**31-1`.
+
+  In other words,
+
+  ```python
+  K <= min(2**31-1, {
+    tf.float16: 2**11,
+    tf.float32: 2**24,
+    tf.float64: 2**53 }[param.dtype])
+  ```
+
+  Note: This condition is validated only when `self.validate_args = True`.
+
+  #### Examples
+
+  Create a 3-class distribution, with the 3rd class is most likely to be drawn,
+  using logits.
+
+  ```python
+  logits = [-50., -43, 0]
+  dist = Multinomial(total_count=4., logits=logits)
+  ```
+
+  Create a 3-class distribution, with the 3rd class is most likely to be drawn.
+
+  ```python
+  p = [.2, .3, .5]
+  dist = Multinomial(total_count=4., probs=p)
+  ```
+
+  The distribution functions can be evaluated on counts.
+
+  ```python
+  # counts same shape as p.
+  counts = [1., 0, 3]
+  dist.prob(counts)  # Shape []
+
+  # p will be broadcast to [[.2, .3, .5], [.2, .3, .5]] to match counts.
+  counts = [[1., 2, 1], [2, 2, 0]]
+  dist.prob(counts)  # Shape [2]
+
+  # p will be broadcast to shape [5, 7, 3] to match counts.
+  counts = [[...]]  # Shape [5, 7, 3]
+  dist.prob(counts)  # Shape [5, 7]
+  ```
+
+  Create a 2-batch of 3-class distributions.
+
+  ```python
+  p = [[.1, .2, .7], [.3, .3, .4]]  # Shape [2, 3]
+  dist = Multinomial(total_count=[4., 5], probs=p)
+
+  counts = [[2., 1, 1], [3, 1, 1]]
+  dist.prob(counts)  # Shape [2]
+
+  dist.sample(5) # Shape [5, 2, 3]
+  ```
+  """
+
+  @deprecation.deprecated(
+      "2019-01-01",
+      "The TensorFlow Distributions library has moved to "
+      "TensorFlow Probability "
+      "(https://github.com/tensorflow/probability). You "
+      "should update all references to use `tfp.distributions` "
+      "instead of `tf.distributions`.",
+      warn_once=True)
+  def __init__(self,
+               total_count,
+               logits=None,
+               probs=None,
+               validate_args=False,
+               allow_nan_stats=True,
+               name="Multinomial"):
+    """Initialize a batch of Multinomial distributions.
+
+    Args:
+      total_count: Non-negative floating point tensor with shape broadcastable
+        to `[N1,..., Nm]` with `m >= 0`. Defines this as a batch of
+        `N1 x ... x Nm` different Multinomial distributions. Its components
+        should be equal to integer values.
+      logits: Floating point tensor representing unnormalized log-probabilities
+        of a positive event with shape broadcastable to
+        `[N1,..., Nm, K]` `m >= 0`, and the same dtype as `total_count`. Defines
+        this as a batch of `N1 x ... x Nm` different `K` class Multinomial
+        distributions. Only one of `logits` or `probs` should be passed in.
+      probs: Positive floating point tensor with shape broadcastable to
+        `[N1,..., Nm, K]` `m >= 0` and same dtype as `total_count`. Defines
+        this as a batch of `N1 x ... x Nm` different `K` class Multinomial
+        distributions. `probs`'s components in the last portion of its shape
+        should sum to `1`. Only one of `logits` or `probs` should be passed in.
+      validate_args: Python `bool`, default `False`. When `True` distribution
+        parameters are checked for validity despite possibly degrading runtime
+        performance. When `False` invalid inputs may silently render incorrect
+        outputs.
+      allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
+        (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
+        result is undefined. When `False`, an exception is raised if one or
+        more of the statistic's batch members are undefined.
+      name: Python `str` name prefixed to Ops created by this class.
+    """
+    parameters = dict(locals())
+    with ops.name_scope(name, values=[total_count, logits, probs]) as name:
+      self._total_count = ops.convert_to_tensor(total_count, name="total_count")
+      if validate_args:
+        self._total_count = (
+            distribution_util.embed_check_nonnegative_integer_form(
+                self._total_count))
+      self._logits, self._probs = distribution_util.get_logits_and_probs(
+          logits=logits,
+          probs=probs,
+          multidimensional=True,
+          validate_args=validate_args,
+          name=name)
+      self._mean_val = self._total_count[..., array_ops.newaxis] * self._probs
+    super(Multinomial, self).__init__(
+        dtype=self._probs.dtype,
+        reparameterization_type=distribution.NOT_REPARAMETERIZED,
+        validate_args=validate_args,
+        allow_nan_stats=allow_nan_stats,
+        parameters=parameters,
+        graph_parents=[self._total_count,
+                       self._logits,
+                       self._probs],
+        name=name)
+
+  @property
+  def total_count(self):
+    """Number of trials used to construct a sample."""
+    return self._total_count
+
+  @property
+  def logits(self):
+    """Vector of coordinatewise logits."""
+    return self._logits
+
+  @property
+  def probs(self):
+    """Probability of drawing a `1` in that coordinate."""
+    return self._probs
+
+  def _batch_shape_tensor(self):
+    return array_ops.shape(self._mean_val)[:-1]
+
+  def _batch_shape(self):
+    return self._mean_val.get_shape().with_rank_at_least(1)[:-1]
+
+  def _event_shape_tensor(self):
+    return array_ops.shape(self._mean_val)[-1:]
+
+  def _event_shape(self):
+    return self._mean_val.get_shape().with_rank_at_least(1)[-1:]
+
+  def _sample_n(self, n, seed=None):
+    n_draws = math_ops.cast(self.total_count, dtype=dtypes.int32)
+    k = self.event_shape_tensor()[0]
+
+    # broadcast the total_count and logits to same shape
+    n_draws = array_ops.ones_like(
+        self.logits[..., 0], dtype=n_draws.dtype) * n_draws
+    logits = array_ops.ones_like(
+        n_draws[..., array_ops.newaxis], dtype=self.logits.dtype) * self.logits
+
+    # flatten the total_count and logits
+    flat_logits = array_ops.reshape(logits, [-1, k])  # [B1B2...Bm, k]
+    flat_ndraws = n * array_ops.reshape(n_draws, [-1])  # [B1B2...Bm]
+
+    # computes each total_count and logits situation by map_fn
+    def _sample_single(args):
+      logits, n_draw = args[0], args[1]  # [K], []
+      x = random_ops.multinomial(logits[array_ops.newaxis, ...], n_draw,
+                                 seed)  # [1, n*n_draw]
+      x = array_ops.reshape(x, shape=[n, -1])  # [n, n_draw]
+      x = math_ops.reduce_sum(array_ops.one_hot(x, depth=k), axis=-2)  # [n, k]
+      return x
+
+    x = map_fn.map_fn(
+        _sample_single, [flat_logits, flat_ndraws],
+        dtype=self.dtype)  # [B1B2...Bm, n, k]
+
+    # reshape the results to proper shape
+    x = array_ops.transpose(x, perm=[1, 0, 2])
+    final_shape = array_ops.concat([[n], self.batch_shape_tensor(), [k]], 0)
+    x = array_ops.reshape(x, final_shape)  # [n, B1, B2,..., Bm, k]
+    return x
+
+  @distribution_util.AppendDocstring(_multinomial_sample_note)
+  def _log_prob(self, counts):
+    return self._log_unnormalized_prob(counts) - self._log_normalization(counts)
+
+  def _log_unnormalized_prob(self, counts):
+    counts = self._maybe_assert_valid_sample(counts)
+    return math_ops.reduce_sum(counts * nn_ops.log_softmax(self.logits), -1)
+
+  def _log_normalization(self, counts):
+    counts = self._maybe_assert_valid_sample(counts)
+    return -distribution_util.log_combinations(self.total_count, counts)
+
+  def _mean(self):
+    return array_ops.identity(self._mean_val)
+
+  def _covariance(self):
+    p = self.probs * array_ops.ones_like(
+        self.total_count)[..., array_ops.newaxis]
+    # pylint: disable=invalid-unary-operand-type
+    return array_ops.matrix_set_diag(
+        -math_ops.matmul(
+            self._mean_val[..., array_ops.newaxis],
+            p[..., array_ops.newaxis, :]),  # outer product
+        self._variance())
+
+  def _variance(self):
+    p = self.probs * array_ops.ones_like(
+        self.total_count)[..., array_ops.newaxis]
+    return self._mean_val - self._mean_val * p
+
+  def _maybe_assert_valid_sample(self, counts):
+    """Check counts for proper shape, values, then return tensor version."""
+    if not self.validate_args:
+      return counts
+    counts = distribution_util.embed_check_nonnegative_integer_form(counts)
+    return control_flow_ops.with_dependencies([
+        check_ops.assert_equal(
+            self.total_count, math_ops.reduce_sum(counts, -1),
+            message="counts must sum to `self.total_count`"),
+    ], counts)

tensorflow/python/ops/distributions/normal.py

@@ -0,0 +1,291 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""The Normal (Gaussian) distribution class."""
+
+import math
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops.distributions import distribution
+from tensorflow.python.ops.distributions import kullback_leibler
+from tensorflow.python.ops.distributions import special_math
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+__all__ = [
+    "Normal",
+    "NormalWithSoftplusScale",
+]
+
+
+@tf_export(v1=["distributions.Normal"])
+class Normal(distribution.Distribution):
+  """The Normal distribution with location `loc` and `scale` parameters.
+
+  #### Mathematical details
+
+  The probability density function (pdf) is,
+
+  ```none
+  pdf(x; mu, sigma) = exp(-0.5 (x - mu)**2 / sigma**2) / Z
+  Z = (2 pi sigma**2)**0.5
+  ```
+
+  where `loc = mu` is the mean, `scale = sigma` is the std. deviation, and, `Z`
+  is the normalization constant.
+
+  The Normal distribution is a member of the [location-scale family](
+  https://en.wikipedia.org/wiki/Location-scale_family), i.e., it can be
+  constructed as,
+
+  ```none
+  X ~ Normal(loc=0, scale=1)
+  Y = loc + scale * X
+  ```
+
+  #### Examples
+
+  Examples of initialization of one or a batch of distributions.
+
+  ```python
+  import tensorflow_probability as tfp
+  tfd = tfp.distributions
+
+  # Define a single scalar Normal distribution.
+  dist = tfd.Normal(loc=0., scale=3.)
+
+  # Evaluate the cdf at 1, returning a scalar.
+  dist.cdf(1.)
+
+  # Define a batch of two scalar valued Normals.
+  # The first has mean 1 and standard deviation 11, the second 2 and 22.
+  dist = tfd.Normal(loc=[1, 2.], scale=[11, 22.])
+
+  # Evaluate the pdf of the first distribution on 0, and the second on 1.5,
+  # returning a length two tensor.
+  dist.prob([0, 1.5])
+
+  # Get 3 samples, returning a 3 x 2 tensor.
+  dist.sample([3])
+  ```
+
+  Arguments are broadcast when possible.
+
+  ```python
+  # Define a batch of two scalar valued Normals.
+  # Both have mean 1, but different standard deviations.
+  dist = tfd.Normal(loc=1., scale=[11, 22.])
+
+  # Evaluate the pdf of both distributions on the same point, 3.0,
+  # returning a length 2 tensor.
+  dist.prob(3.0)
+  ```
+
+  """
+
+  @deprecation.deprecated(
+      "2019-01-01",
+      "The TensorFlow Distributions library has moved to "
+      "TensorFlow Probability "
+      "(https://github.com/tensorflow/probability). You "
+      "should update all references to use `tfp.distributions` "
+      "instead of `tf.distributions`.",
+      warn_once=True)
+  def __init__(self,
+               loc,
+               scale,
+               validate_args=False,
+               allow_nan_stats=True,
+               name="Normal"):
+    """Construct Normal distributions with mean and stddev `loc` and `scale`.
+
+    The parameters `loc` and `scale` must be shaped in a way that supports
+    broadcasting (e.g. `loc + scale` is a valid operation).
+
+    Args:
+      loc: Floating point tensor; the means of the distribution(s).
+      scale: Floating point tensor; the stddevs of the distribution(s).
+        Must contain only positive values.
+      validate_args: Python `bool`, default `False`. When `True` distribution
+        parameters are checked for validity despite possibly degrading runtime
+        performance. When `False` invalid inputs may silently render incorrect
+        outputs.
+      allow_nan_stats: Python `bool`, default `True`. When `True`,
+        statistics (e.g., mean, mode, variance) use the value "`NaN`" to
+        indicate the result is undefined. When `False`, an exception is raised
+        if one or more of the statistic's batch members are undefined.
+      name: Python `str` name prefixed to Ops created by this class.
+
+    Raises:
+      TypeError: if `loc` and `scale` have different `dtype`.
+    """
+    parameters = dict(locals())
+    with ops.name_scope(name, values=[loc, scale]) as name:
+      with ops.control_dependencies([check_ops.assert_positive(scale)] if
+                                    validate_args else []):
+        self._loc = array_ops.identity(loc, name="loc")
+        self._scale = array_ops.identity(scale, name="scale")
+        check_ops.assert_same_float_dtype([self._loc, self._scale])
+    super(Normal, self).__init__(
+        dtype=self._scale.dtype,
+        reparameterization_type=distribution.FULLY_REPARAMETERIZED,
+        validate_args=validate_args,
+        allow_nan_stats=allow_nan_stats,
+        parameters=parameters,
+        graph_parents=[self._loc, self._scale],
+        name=name)
+
+  @staticmethod
+  def _param_shapes(sample_shape):
+    return dict(
+        zip(("loc", "scale"), ([ops.convert_to_tensor(
+            sample_shape, dtype=dtypes.int32)] * 2)))
+
+  @property
+  def loc(self):
+    """Distribution parameter for the mean."""
+    return self._loc
+
+  @property
+  def scale(self):
+    """Distribution parameter for standard deviation."""
+    return self._scale
+
+  def _batch_shape_tensor(self):
+    return array_ops.broadcast_dynamic_shape(
+        array_ops.shape(self.loc),
+        array_ops.shape(self.scale))
+
+  def _batch_shape(self):
+    return array_ops.broadcast_static_shape(
+        self.loc.get_shape(),
+        self.scale.get_shape())
+
+  def _event_shape_tensor(self):
+    return constant_op.constant([], dtype=dtypes.int32)
+
+  def _event_shape(self):
+    return tensor_shape.TensorShape([])
+
+  def _sample_n(self, n, seed=None):
+    shape = array_ops.concat([[n], self.batch_shape_tensor()], 0)
+    sampled = random_ops.random_normal(
+        shape=shape, mean=0., stddev=1., dtype=self.loc.dtype, seed=seed)
+    return sampled * self.scale + self.loc
+
+  def _log_prob(self, x):
+    return self._log_unnormalized_prob(x) - self._log_normalization()
+
+  def _log_cdf(self, x):
+    return special_math.log_ndtr(self._z(x))
+
+  def _cdf(self, x):
+    return special_math.ndtr(self._z(x))
+
+  def _log_survival_function(self, x):
+    return special_math.log_ndtr(-self._z(x))
+
+  def _survival_function(self, x):
+    return special_math.ndtr(-self._z(x))
+
+  def _log_unnormalized_prob(self, x):
+    return -0.5 * math_ops.square(self._z(x))
+
+  def _log_normalization(self):
+    return 0.5 * math.log(2. * math.pi) + math_ops.log(self.scale)
+
+  def _entropy(self):
+    # Use broadcasting rules to calculate the full broadcast scale.
+    scale = self.scale * array_ops.ones_like(self.loc)
+    return 0.5 * math.log(2. * math.pi * math.e) + math_ops.log(scale)
+
+  def _mean(self):
+    return self.loc * array_ops.ones_like(self.scale)
+
+  def _quantile(self, p):
+    return self._inv_z(special_math.ndtri(p))
+
+  def _stddev(self):
+    return self.scale * array_ops.ones_like(self.loc)
+
+  def _mode(self):
+    return self._mean()
+
+  def _z(self, x):
+    """Standardize input `x` to a unit normal."""
+    with ops.name_scope("standardize", values=[x]):
+      return (x - self.loc) / self.scale
+
+  def _inv_z(self, z):
+    """Reconstruct input `x` from a its normalized version."""
+    with ops.name_scope("reconstruct", values=[z]):
+      return z * self.scale + self.loc
+
+
+class NormalWithSoftplusScale(Normal):
+  """Normal with softplus applied to `scale`."""
+
+  @deprecation.deprecated(
+      "2019-01-01",
+      "Use `tfd.Normal(loc, tf.nn.softplus(scale)) "
+      "instead.",
+      warn_once=True)
+  def __init__(self,
+               loc,
+               scale,
+               validate_args=False,
+               allow_nan_stats=True,
+               name="NormalWithSoftplusScale"):
+    parameters = dict(locals())
+    with ops.name_scope(name, values=[scale]) as name:
+      super(NormalWithSoftplusScale, self).__init__(
+          loc=loc,
+          scale=nn.softplus(scale, name="softplus_scale"),
+          validate_args=validate_args,
+          allow_nan_stats=allow_nan_stats,
+          name=name)
+    self._parameters = parameters
+
+
+@kullback_leibler.RegisterKL(Normal, Normal)
+def _kl_normal_normal(n_a, n_b, name=None):
+  """Calculate the batched KL divergence KL(n_a || n_b) with n_a and n_b Normal.
+
+  Args:
+    n_a: instance of a Normal distribution object.
+    n_b: instance of a Normal distribution object.
+    name: (optional) Name to use for created operations.
+      default is "kl_normal_normal".
+
+  Returns:
+    Batchwise KL(n_a || n_b)
+  """
+  with ops.name_scope(name, "kl_normal_normal", [n_a.loc, n_b.loc]):
+    one = constant_op.constant(1, dtype=n_a.dtype)
+    two = constant_op.constant(2, dtype=n_a.dtype)
+    half = constant_op.constant(0.5, dtype=n_a.dtype)
+    s_a_squared = math_ops.square(n_a.scale)
+    s_b_squared = math_ops.square(n_b.scale)
+    ratio = s_a_squared / s_b_squared
+    return (math_ops.squared_difference(n_a.loc, n_b.loc) / (two * s_b_squared)
+            + half * (ratio - one - math_ops.log(ratio)))

tensorflow/python/ops/distributions/special_math.py

@@ -0,0 +1,470 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# Functions "ndtr" and "ndtri" are derived from calculations made in:
+# https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
+# In the following email exchange, the author gives his consent to redistribute
+# derived works under an Apache 2.0 license.
+#
+# From: Stephen Moshier <steve@moshier.net>
+# Date: Sat, Jun 9, 2018 at 2:36 PM
+# Subject: Re: Licensing cephes under Apache (BSD-like) license.
+# To: rif <rif@google.com>
+#
+#
+#
+# Hello Rif,
+#
+# Yes, Google may distribute Cephes files under the Apache 2 license.
+#
+# If clarification is needed, I do not favor BSD over other free licenses.
+# I would agree that Apache 2 seems to cover the concern you mentioned
+# about sublicensees.
+#
+# Best wishes for good luck with your projects!
+# Steve Moshier
+#
+#
+#
+# On Thu, 31 May 2018, rif wrote:
+#
+# > Hello Steve.
+# > My name is Rif. I work on machine learning software at Google.
+# >
+# > Your cephes software continues to be incredibly useful and widely used. I
+# > was wondering whether it would be permissible for us to use the Cephes code
+# > under the Apache 2.0 license, which is extremely similar in permissions to
+# > the BSD license (Wikipedia comparisons). This would be quite helpful to us
+# > in terms of avoiding multiple licenses on software.
+# >
+# > I'm sorry to bother you with this (I can imagine you're sick of hearing
+# > about this by now), but I want to be absolutely clear we're on the level and
+# > not misusing your important software. In former conversation with Eugene
+# > Brevdo (ebrevdo@google.com), you wrote "If your licensing is similar to BSD,
+# > the formal way that has been handled is simply to add a statement to the
+# > effect that you are incorporating the Cephes software by permission of the
+# > author." I wanted to confirm that (a) we could use the Apache license, (b)
+# > that we don't need to (and probably you don't want to) keep getting
+# > contacted about individual uses, because your intent is generally to allow
+# > this software to be reused under "BSD-like" license, and (c) you're OK
+# > letting incorporators decide whether a license is sufficiently BSD-like?
+# >
+# > Best,
+# >
+# > rif
+# >
+# >
+# >
+
+"""Special Math Ops."""
+
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+
+__all__ = [
+    "erfinv",
+    "ndtr",
+    "ndtri",
+    "log_ndtr",
+    "log_cdf_laplace",
+]
+
+
+# log_ndtr uses different functions over the ranges
+# (-infty, lower](lower, upper](upper, infty)
+# Lower bound values were chosen by examining where the support of ndtr
+# appears to be zero, relative to scipy's (which is always 64bit). They were
+# then made more conservative just to be safe. (Conservative means use the
+# expansion more than we probably need to.) See `NdtrTest` in
+# special_math_test.py.
+LOGNDTR_FLOAT64_LOWER = np.array(-20, np.float64)
+LOGNDTR_FLOAT32_LOWER = np.array(-10, np.float32)
+
+# Upper bound values were chosen by examining for which values of 'x'
+# Log[cdf(x)] is 0, after which point we need to use the approximation
+# Log[cdf(x)] = Log[1 - cdf(-x)] approx -cdf(-x). We chose a value slightly
+# conservative, meaning we use the approximation earlier than needed.
+LOGNDTR_FLOAT64_UPPER = np.array(8, np.float64)
+LOGNDTR_FLOAT32_UPPER = np.array(5, np.float32)
+
+
+def ndtr(x, name="ndtr"):
+  """Normal distribution function.
+
+  Returns the area under the Gaussian probability density function, integrated
+  from minus infinity to x:
+
+  ```
+                    1       / x
+     ndtr(x)  = ----------  |    exp(-0.5 t**2) dt
+                sqrt(2 pi)  /-inf
+
+              = 0.5 (1 + erf(x / sqrt(2)))
+              = 0.5 erfc(x / sqrt(2))
+  ```
+
+  Args:
+    x: `Tensor` of type `float32`, `float64`.
+    name: Python string. A name for the operation (default="ndtr").
+
+  Returns:
+    ndtr: `Tensor` with `dtype=x.dtype`.
+
+  Raises:
+    TypeError: if `x` is not floating-type.
+  """
+
+  with ops.name_scope(name, values=[x]):
+    x = ops.convert_to_tensor(x, name="x")
+    if x.dtype.as_numpy_dtype not in [np.float32, np.float64]:
+      raise TypeError(
+          "x.dtype=%s is not handled, see docstring for supported types."
+          % x.dtype)
+    return _ndtr(x)
+
+
+def _ndtr(x):
+  """Implements ndtr core logic."""
+  half_sqrt_2 = constant_op.constant(
+      0.5 * np.sqrt(2.), dtype=x.dtype, name="half_sqrt_2")
+  w = x * half_sqrt_2
+  z = math_ops.abs(w)
+  y = array_ops.where_v2(
+      math_ops.less(z, half_sqrt_2), 1. + math_ops.erf(w),
+      array_ops.where_v2(
+          math_ops.greater(w, 0.), 2. - math_ops.erfc(z), math_ops.erfc(z)))
+  return 0.5 * y
+
+
+def ndtri(p, name="ndtri"):
+  """The inverse of the CDF of the Normal distribution function.
+
+  Returns x such that the area under the pdf from minus infinity to x is equal
+  to p.
+
+  A piece-wise rational approximation is done for the function.
+  This is a port of the implementation in netlib.
+
+  Args:
+    p: `Tensor` of type `float32`, `float64`.
+    name: Python string. A name for the operation (default="ndtri").
+
+  Returns:
+    x: `Tensor` with `dtype=p.dtype`.
+
+  Raises:
+    TypeError: if `p` is not floating-type.
+  """
+
+  with ops.name_scope(name, values=[p]):
+    p = ops.convert_to_tensor(p, name="p")
+    if p.dtype.as_numpy_dtype not in [np.float32, np.float64]:
+      raise TypeError(
+          "p.dtype=%s is not handled, see docstring for supported types."
+          % p.dtype)
+    return _ndtri(p)
+
+
+def _ndtri(p):
+  """Implements ndtri core logic."""
+
+  # Constants used in piece-wise rational approximations. Taken from the cephes
+  # library:
+  # https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
+
+  p0 = [
+      -1.23916583867381258016E0, 1.39312609387279679503E1,
+      -5.66762857469070293439E1, 9.80010754185999661536E1,
+      -5.99633501014107895267E1
+  ]
+  q0 = [
+      -1.18331621121330003142E0, 1.59056225126211695515E1,
+      -8.20372256168333339912E1, 2.00260212380060660359E2,
+      -2.25462687854119370527E2, 8.63602421390890590575E1,
+      4.67627912898881538453E0, 1.95448858338141759834E0, 1.0
+  ]
+  p1 = [
+      -8.57456785154685413611E-4, -3.50424626827848203418E-2,
+      -1.40256079171354495875E-1, 2.18663306850790267539E0,
+      1.46849561928858024014E1, 4.40805073893200834700E1,
+      5.71628192246421288162E1, 3.15251094599893866154E1,
+      4.05544892305962419923E0
+  ]
+  q1 = [
+      -9.33259480895457427372E-4, -3.80806407691578277194E-2,
+      -1.42182922854787788574E-1, 2.50464946208309415979E0,
+      1.50425385692907503408E1, 4.13172038254672030440E1,
+      4.53907635128879210584E1, 1.57799883256466749731E1, 1.0
+  ]
+  p2 = [
+      6.23974539184983293730E-9, 2.65806974686737550832E-6,
+      3.01581553508235416007E-4, 1.23716634817820021358E-2,
+      2.01485389549179081538E-1, 1.33303460815807542389E0,
+      3.93881025292474443415E0, 6.91522889068984211695E0,
+      3.23774891776946035970E0
+  ]
+  q2 = [
+      6.79019408009981274425E-9, 2.89247864745380683936E-6,
+      3.28014464682127739104E-4, 1.34204006088543189037E-2,
+      2.16236993594496635890E-1, 1.37702099489081330271E0,
+      3.67983563856160859403E0, 6.02427039364742014255E0, 1.0
+  ]
+
+  def _create_polynomial(var, coeffs):
+    """Compute n_th order polynomial via Horner's method."""
+    coeffs = np.array(coeffs, var.dtype.as_numpy_dtype)
+    if not coeffs.size:
+      return array_ops.zeros_like(var)
+    return coeffs[0] + _create_polynomial(var, coeffs[1:]) * var
+
+  maybe_complement_p = array_ops.where_v2(p > -np.expm1(-2.), 1. - p, p)
+  # Write in an arbitrary value in place of 0 for p since 0 will cause NaNs
+  # later on. The result from the computation when p == 0 is not used so any
+  # number that doesn't result in NaNs is fine.
+  sanitized_mcp = array_ops.where_v2(
+      maybe_complement_p <= 0.,
+      array_ops.fill(array_ops.shape(p), np.array(0.5, p.dtype.as_numpy_dtype)),
+      maybe_complement_p)
+
+  # Compute x for p > exp(-2): x/sqrt(2pi) = w + w**3 P0(w**2)/Q0(w**2).
+  w = sanitized_mcp - 0.5
+  ww = w ** 2
+  x_for_big_p = w + w * ww * (_create_polynomial(ww, p0)
+                              / _create_polynomial(ww, q0))
+  x_for_big_p *= -np.sqrt(2. * np.pi)
+
+  # Compute x for p <= exp(-2): x = z - log(z)/z - (1/z) P(1/z) / Q(1/z),
+  # where z = sqrt(-2. * log(p)), and P/Q are chosen between two different
+  # arrays based on whether p < exp(-32).
+  z = math_ops.sqrt(-2. * math_ops.log(sanitized_mcp))
+  first_term = z - math_ops.log(z) / z
+  second_term_small_p = (
+      _create_polynomial(1. / z, p2) /
+      _create_polynomial(1. / z, q2) / z)
+  second_term_otherwise = (
+      _create_polynomial(1. / z, p1) /
+      _create_polynomial(1. / z, q1) / z)
+  x_for_small_p = first_term - second_term_small_p
+  x_otherwise = first_term - second_term_otherwise
+
+  x = array_ops.where_v2(
+      sanitized_mcp > np.exp(-2.), x_for_big_p,
+      array_ops.where_v2(z >= 8.0, x_for_small_p, x_otherwise))
+
+  x = array_ops.where_v2(p > 1. - np.exp(-2.), x, -x)
+  infinity_scalar = constant_op.constant(np.inf, dtype=p.dtype)
+  infinity = array_ops.fill(array_ops.shape(p), infinity_scalar)
+  x_nan_replaced = array_ops.where_v2(p <= 0.0, -infinity,
+                                      array_ops.where_v2(p >= 1.0, infinity, x))
+  return x_nan_replaced
+
+
+def log_ndtr(x, series_order=3, name="log_ndtr"):
+  """Log Normal distribution function.
+
+  For details of the Normal distribution function see `ndtr`.
+
+  This function calculates `(log o ndtr)(x)` by either calling `log(ndtr(x))` or
+  using an asymptotic series. Specifically:
+  - For `x > upper_segment`, use the approximation `-ndtr(-x)` based on
+    `log(1-x) ~= -x, x << 1`.
+  - For `lower_segment < x <= upper_segment`, use the existing `ndtr` technique
+    and take a log.
+  - For `x <= lower_segment`, we use the series approximation of erf to compute
+    the log CDF directly.
+
+  The `lower_segment` is set based on the precision of the input:
+
+  ```
+  lower_segment = { -20,  x.dtype=float64
+                  { -10,  x.dtype=float32
+  upper_segment = {   8,  x.dtype=float64
+                  {   5,  x.dtype=float32
+  ```
+
+  When `x < lower_segment`, the `ndtr` asymptotic series approximation is:
+
+  ```
+     ndtr(x) = scale * (1 + sum) + R_N
+     scale   = exp(-0.5 x**2) / (-x sqrt(2 pi))
+     sum     = Sum{(-1)^n (2n-1)!! / (x**2)^n, n=1:N}
+     R_N     = O(exp(-0.5 x**2) (2N+1)!! / |x|^{2N+3})
+  ```
+
+  where `(2n-1)!! = (2n-1) (2n-3) (2n-5) ...  (3) (1)` is a
+  [double-factorial](https://en.wikipedia.org/wiki/Double_factorial).
+
+
+  Args:
+    x: `Tensor` of type `float32`, `float64`.
+    series_order: Positive Python `integer`. Maximum depth to
+      evaluate the asymptotic expansion. This is the `N` above.
+    name: Python string. A name for the operation (default="log_ndtr").
+
+  Returns:
+    log_ndtr: `Tensor` with `dtype=x.dtype`.
+
+  Raises:
+    TypeError: if `x.dtype` is not handled.
+    TypeError: if `series_order` is a not Python `integer.`
+    ValueError:  if `series_order` is not in `[0, 30]`.
+  """
+  if not isinstance(series_order, int):
+    raise TypeError("series_order must be a Python integer.")
+  if series_order < 0:
+    raise ValueError("series_order must be non-negative.")
+  if series_order > 30:
+    raise ValueError("series_order must be <= 30.")
+
+  with ops.name_scope(name, values=[x]):
+    x = ops.convert_to_tensor(x, name="x")
+
+    if x.dtype.as_numpy_dtype == np.float64:
+      lower_segment = LOGNDTR_FLOAT64_LOWER
+      upper_segment = LOGNDTR_FLOAT64_UPPER
+    elif x.dtype.as_numpy_dtype == np.float32:
+      lower_segment = LOGNDTR_FLOAT32_LOWER
+      upper_segment = LOGNDTR_FLOAT32_UPPER
+    else:
+      raise TypeError("x.dtype=%s is not supported." % x.dtype)
+
+    # The basic idea here was ported from:
+    #   https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
+    # We copy the main idea, with a few changes
+    # * For x >> 1, and X ~ Normal(0, 1),
+    #     Log[P[X < x]] = Log[1 - P[X < -x]] approx -P[X < -x],
+    #     which extends the range of validity of this function.
+    # * We use one fixed series_order for all of 'x', rather than adaptive.
+    # * Our docstring properly reflects that this is an asymptotic series, not a
+    #   Taylor series. We also provided a correct bound on the remainder.
+    # * We need to use the max/min in the _log_ndtr_lower arg to avoid nan when
+    #   x=0. This happens even though the branch is unchosen because when x=0
+    #   the gradient of a select involves the calculation 1*dy+0*(-inf)=nan
+    #   regardless of whether dy is finite. Note that the minimum is a NOP if
+    #   the branch is chosen.
+    return array_ops.where_v2(
+        math_ops.greater(x, upper_segment),
+        -_ndtr(-x),  # log(1-x) ~= -x, x << 1  # pylint: disable=invalid-unary-operand-type
+        array_ops.where_v2(
+            math_ops.greater(x, lower_segment),
+            math_ops.log(_ndtr(math_ops.maximum(x, lower_segment))),
+            _log_ndtr_lower(math_ops.minimum(x, lower_segment), series_order)))
+
+
+def _log_ndtr_lower(x, series_order):
+  """Asymptotic expansion version of `Log[cdf(x)]`, appropriate for `x<<-1`."""
+  x_2 = math_ops.square(x)
+  # Log of the term multiplying (1 + sum)
+  log_scale = -0.5 * x_2 - math_ops.log(-x) - 0.5 * np.log(2. * np.pi)
+  return log_scale + math_ops.log(_log_ndtr_asymptotic_series(x, series_order))
+
+
+def _log_ndtr_asymptotic_series(x, series_order):
+  """Calculates the asymptotic series used in log_ndtr."""
+  dtype = x.dtype.as_numpy_dtype
+  if series_order <= 0:
+    return np.array(1, dtype)
+  x_2 = math_ops.square(x)
+  even_sum = array_ops.zeros_like(x)
+  odd_sum = array_ops.zeros_like(x)
+  x_2n = x_2  # Start with x^{2*1} = x^{2*n} with n = 1.
+  for n in range(1, series_order + 1):
+    y = np.array(_double_factorial(2 * n - 1), dtype) / x_2n
+    if n % 2:
+      odd_sum += y
+    else:
+      even_sum += y
+    x_2n *= x_2
+  return 1. + even_sum - odd_sum
+
+
+def erfinv(x, name="erfinv"):
+  """The inverse function for erf, the error function.
+
+  Args:
+    x: `Tensor` of type `float32`, `float64`.
+    name: Python string. A name for the operation (default="erfinv").
+
+  Returns:
+    x: `Tensor` with `dtype=x.dtype`.
+
+  Raises:
+    TypeError: if `x` is not floating-type.
+  """
+
+  with ops.name_scope(name, values=[x]):
+    x = ops.convert_to_tensor(x, name="x")
+    if x.dtype.as_numpy_dtype not in [np.float32, np.float64]:
+      raise TypeError(
+          "x.dtype=%s is not handled, see docstring for supported types."
+          % x.dtype)
+    return ndtri((x + 1.0) / 2.0) / np.sqrt(2)
+
+
+def _double_factorial(n):
+  """The double factorial function for small Python integer `n`."""
+  return np.prod(np.arange(n, 1, -2))
+
+
+def log_cdf_laplace(x, name="log_cdf_laplace"):
+  """Log Laplace distribution function.
+
+  This function calculates `Log[L(x)]`, where `L(x)` is the cumulative
+  distribution function of the Laplace distribution, i.e.
+
+  ```L(x) := 0.5 * int_{-infty}^x e^{-|t|} dt```
+
+  For numerical accuracy, `L(x)` is computed in different ways depending on `x`,
+
+  ```
+  x <= 0:
+    Log[L(x)] = Log[0.5] + x, which is exact
+
+  0 < x:
+    Log[L(x)] = Log[1 - 0.5 * e^{-x}], which is exact
+  ```
+
+  Args:
+    x: `Tensor` of type `float32`, `float64`.
+    name: Python string. A name for the operation (default="log_ndtr").
+
+  Returns:
+    `Tensor` with `dtype=x.dtype`.
+
+  Raises:
+    TypeError: if `x.dtype` is not handled.
+  """
+
+  with ops.name_scope(name, values=[x]):
+    x = ops.convert_to_tensor(x, name="x")
+
+    # For x < 0, L(x) = 0.5 * exp{x} exactly, so Log[L(x)] = log(0.5) + x.
+    lower_solution = -np.log(2.) + x
+
+    # safe_exp_neg_x = exp{-x} for x > 0, but is
+    # bounded above by 1, which avoids
+    #   log[1 - 1] = -inf for x = log(1/2), AND
+    #   exp{-x} --> inf, for x << -1
+    safe_exp_neg_x = math_ops.exp(-math_ops.abs(x))
+
+    # log1p(z) = log(1 + z) approx z for |z| << 1. This approximation is used
+    # internally by log1p, rather than being done explicitly here.
+    upper_solution = math_ops.log1p(-0.5 * safe_exp_neg_x)
+
+    return array_ops.where_v2(x < 0., lower_solution, upper_solution)

tensorflow/python/ops/distributions/student_t.py

@@ -0,0 +1,391 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Student's t distribution class."""
+
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import special_math_ops
+from tensorflow.python.ops.distributions import distribution
+from tensorflow.python.ops.distributions import util as distribution_util
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+__all__ = [
+    "StudentT",
+    "StudentTWithAbsDfSoftplusScale",
+]
+
+
+@tf_export(v1=["distributions.StudentT"])
+class StudentT(distribution.Distribution):
+  """Student's t-distribution.
+
+  This distribution has parameters: degree of freedom `df`, location `loc`,
+  and `scale`.
+
+  #### Mathematical details
+
+  The probability density function (pdf) is,
+
+  ```none
+  pdf(x; df, mu, sigma) = (1 + y**2 / df)**(-0.5 (df + 1)) / Z
+  where,
+  y = (x - mu) / sigma
+  Z = abs(sigma) sqrt(df pi) Gamma(0.5 df) / Gamma(0.5 (df + 1))
+  ```
+
+  where:
+  * `loc = mu`,
+  * `scale = sigma`, and,
+  * `Z` is the normalization constant, and,
+  * `Gamma` is the [gamma function](
+    https://en.wikipedia.org/wiki/Gamma_function).
+
+  The StudentT distribution is a member of the [location-scale family](
+  https://en.wikipedia.org/wiki/Location-scale_family), i.e., it can be
+  constructed as,
+
+  ```none
+  X ~ StudentT(df, loc=0, scale=1)
+  Y = loc + scale * X
+  ```
+
+  Notice that `scale` has semantics more similar to standard deviation than
+  variance. However it is not actually the std. deviation; the Student's
+  t-distribution std. dev. is `scale sqrt(df / (df - 2))` when `df > 2`.
+
+  Samples of this distribution are reparameterized (pathwise differentiable).
+  The derivatives are computed using the approach described in
+  (Figurnov et al., 2018).
+
+  #### Examples
+
+  Examples of initialization of one or a batch of distributions.
+
+  ```python
+  import tensorflow_probability as tfp
+  tfd = tfp.distributions
+
+  # Define a single scalar Student t distribution.
+  single_dist = tfd.StudentT(df=3)
+
+  # Evaluate the pdf at 1, returning a scalar Tensor.
+  single_dist.prob(1.)
+
+  # Define a batch of two scalar valued Student t's.
+  # The first has degrees of freedom 2, mean 1, and scale 11.
+  # The second 3, 2 and 22.
+  multi_dist = tfd.StudentT(df=[2, 3], loc=[1, 2.], scale=[11, 22.])
+
+  # Evaluate the pdf of the first distribution on 0, and the second on 1.5,
+  # returning a length two tensor.
+  multi_dist.prob([0, 1.5])
+
+  # Get 3 samples, returning a 3 x 2 tensor.
+  multi_dist.sample(3)
+  ```
+
+  Arguments are broadcast when possible.
+
+  ```python
+  # Define a batch of two Student's t distributions.
+  # Both have df 2 and mean 1, but different scales.
+  dist = tfd.StudentT(df=2, loc=1, scale=[11, 22.])
+
+  # Evaluate the pdf of both distributions on the same point, 3.0,
+  # returning a length 2 tensor.
+  dist.prob(3.0)
+  ```
+
+  Compute the gradients of samples w.r.t. the parameters:
+
+  ```python
+  df = tf.constant(2.0)
+  loc = tf.constant(2.0)
+  scale = tf.constant(11.0)
+  dist = tfd.StudentT(df=df, loc=loc, scale=scale)
+  samples = dist.sample(5)  # Shape [5]
+  loss = tf.reduce_mean(tf.square(samples))  # Arbitrary loss function
+  # Unbiased stochastic gradients of the loss function
+  grads = tf.gradients(loss, [df, loc, scale])
+  ```
+
+  References:
+    Implicit Reparameterization Gradients:
+      [Figurnov et al., 2018]
+      (http://papers.nips.cc/paper/7326-implicit-reparameterization-gradients)
+      ([pdf](http://papers.nips.cc/paper/7326-implicit-reparameterization-gradients.pdf))
+  """
+
+  @deprecation.deprecated(
+      "2019-01-01",
+      "The TensorFlow Distributions library has moved to "
+      "TensorFlow Probability "
+      "(https://github.com/tensorflow/probability). You "
+      "should update all references to use `tfp.distributions` "
+      "instead of `tf.distributions`.",
+      warn_once=True)
+  def __init__(self,
+               df,
+               loc,
+               scale,
+               validate_args=False,
+               allow_nan_stats=True,
+               name="StudentT"):
+    """Construct Student's t distributions.
+
+    The distributions have degree of freedom `df`, mean `loc`, and scale
+    `scale`.
+
+    The parameters `df`, `loc`, and `scale` must be shaped in a way that
+    supports broadcasting (e.g. `df + loc + scale` is a valid operation).
+
+    Args:
+      df: Floating-point `Tensor`. The degrees of freedom of the
+        distribution(s). `df` must contain only positive values.
+      loc: Floating-point `Tensor`. The mean(s) of the distribution(s).
+      scale: Floating-point `Tensor`. The scaling factor(s) for the
+        distribution(s). Note that `scale` is not technically the standard
+        deviation of this distribution but has semantics more similar to
+        standard deviation than variance.
+      validate_args: Python `bool`, default `False`. When `True` distribution
+        parameters are checked for validity despite possibly degrading runtime
+        performance. When `False` invalid inputs may silently render incorrect
+        outputs.
+      allow_nan_stats: Python `bool`, default `True`. When `True`,
+        statistics (e.g., mean, mode, variance) use the value "`NaN`" to
+        indicate the result is undefined. When `False`, an exception is raised
+        if one or more of the statistic's batch members are undefined.
+      name: Python `str` name prefixed to Ops created by this class.
+
+    Raises:
+      TypeError: if loc and scale are different dtypes.
+    """
+    parameters = dict(locals())
+    with ops.name_scope(name, values=[df, loc, scale]) as name:
+      with ops.control_dependencies([check_ops.assert_positive(df)]
+                                    if validate_args else []):
+        self._df = array_ops.identity(df, name="df")
+        self._loc = array_ops.identity(loc, name="loc")
+        self._scale = array_ops.identity(scale, name="scale")
+        check_ops.assert_same_float_dtype(
+            (self._df, self._loc, self._scale))
+    super(StudentT, self).__init__(
+        dtype=self._scale.dtype,
+        reparameterization_type=distribution.FULLY_REPARAMETERIZED,
+        validate_args=validate_args,
+        allow_nan_stats=allow_nan_stats,
+        parameters=parameters,
+        graph_parents=[self._df, self._loc, self._scale],
+        name=name)
+
+  @staticmethod
+  def _param_shapes(sample_shape):
+    return dict(
+        zip(("df", "loc", "scale"), (
+            [ops.convert_to_tensor(
+                sample_shape, dtype=dtypes.int32)] * 3)))
+
+  @property
+  def df(self):
+    """Degrees of freedom in these Student's t distribution(s)."""
+    return self._df
+
+  @property
+  def loc(self):
+    """Locations of these Student's t distribution(s)."""
+    return self._loc
+
+  @property
+  def scale(self):
+    """Scaling factors of these Student's t distribution(s)."""
+    return self._scale
+
+  def _batch_shape_tensor(self):
+    return array_ops.broadcast_dynamic_shape(
+        array_ops.shape(self.df),
+        array_ops.broadcast_dynamic_shape(
+            array_ops.shape(self.loc), array_ops.shape(self.scale)))
+
+  def _batch_shape(self):
+    return array_ops.broadcast_static_shape(
+        array_ops.broadcast_static_shape(self.df.get_shape(),
+                                         self.loc.get_shape()),
+        self.scale.get_shape())
+
+  def _event_shape_tensor(self):
+    return constant_op.constant([], dtype=math_ops.int32)
+
+  def _event_shape(self):
+    return tensor_shape.TensorShape([])
+
+  def _sample_n(self, n, seed=None):
+    # The sampling method comes from the fact that if:
+    #   X ~ Normal(0, 1)
+    #   Z ~ Chi2(df)
+    #   Y = X / sqrt(Z / df)
+    # then:
+    #   Y ~ StudentT(df).
+    shape = array_ops.concat([[n], self.batch_shape_tensor()], 0)
+    normal_sample = random_ops.random_normal(shape, dtype=self.dtype, seed=seed)
+    df = self.df * array_ops.ones(self.batch_shape_tensor(), dtype=self.dtype)
+    gamma_sample = random_ops.random_gamma(
+        [n],
+        0.5 * df,
+        beta=0.5,
+        dtype=self.dtype,
+        seed=distribution_util.gen_new_seed(seed, salt="student_t"))
+    samples = normal_sample * math_ops.rsqrt(gamma_sample / df)
+    return samples * self.scale + self.loc  # Abs(scale) not wanted.
+
+  def _log_prob(self, x):
+    return self._log_unnormalized_prob(x) - self._log_normalization()
+
+  def _log_unnormalized_prob(self, x):
+    y = (x - self.loc) / self.scale  # Abs(scale) superfluous.
+    return -0.5 * (self.df + 1.) * math_ops.log1p(y**2. / self.df)
+
+  def _log_normalization(self):
+    return (math_ops.log(math_ops.abs(self.scale)) +
+            0.5 * math_ops.log(self.df) +
+            0.5 * np.log(np.pi) +
+            math_ops.lgamma(0.5 * self.df) -
+            math_ops.lgamma(0.5 * (self.df + 1.)))
+
+  def _cdf(self, x):
+    # Take Abs(scale) to make subsequent where work correctly.
+    y = (x - self.loc) / math_ops.abs(self.scale)
+    x_t = self.df / (y**2. + self.df)
+    neg_cdf = 0.5 * math_ops.betainc(0.5 * self.df, 0.5, x_t)
+    return array_ops.where_v2(math_ops.less(y, 0.), neg_cdf, 1. - neg_cdf)
+
+  def _entropy(self):
+    v = array_ops.ones(self.batch_shape_tensor(),
+                       dtype=self.dtype)[..., array_ops.newaxis]
+    u = v * self.df[..., array_ops.newaxis]
+    beta_arg = array_ops.concat([u, v], -1) / 2.
+    return (math_ops.log(math_ops.abs(self.scale)) +
+            0.5 * math_ops.log(self.df) +
+            special_math_ops.lbeta(beta_arg) +
+            0.5 * (self.df + 1.) *
+            (math_ops.digamma(0.5 * (self.df + 1.)) -
+             math_ops.digamma(0.5 * self.df)))
+
+  @distribution_util.AppendDocstring(
+      """The mean of Student's T equals `loc` if `df > 1`, otherwise it is
+      `NaN`. If `self.allow_nan_stats=True`, then an exception will be raised
+      rather than returning `NaN`.""")
+  def _mean(self):
+    mean = self.loc * array_ops.ones(self.batch_shape_tensor(),
+                                     dtype=self.dtype)
+    if self.allow_nan_stats:
+      nan = np.array(np.nan, dtype=self.dtype.as_numpy_dtype())
+      return array_ops.where_v2(
+          math_ops.greater(
+              self.df,
+              array_ops.ones(self.batch_shape_tensor(), dtype=self.dtype)),
+          mean, array_ops.fill(self.batch_shape_tensor(), nan, name="nan"))
+    else:
+      return control_flow_ops.with_dependencies(
+          [
+              check_ops.assert_less(
+                  array_ops.ones([], dtype=self.dtype),
+                  self.df,
+                  message="mean not defined for components of df <= 1"),
+          ],
+          mean)
+
+  @distribution_util.AppendDocstring("""
+      The variance for Student's T equals
+
+      ```
+      df / (df - 2), when df > 2
+      infinity, when 1 < df <= 2
+      NaN, when df <= 1
+      ```
+      """)
+  def _variance(self):
+    # We need to put the tf.where inside the outer tf.where to ensure we never
+    # hit a NaN in the gradient.
+    denom = array_ops.where_v2(
+        math_ops.greater(self.df, 2.), self.df - 2.,
+        array_ops.ones_like(self.df))
+    # Abs(scale) superfluous.
+    var = (array_ops.ones(self.batch_shape_tensor(), dtype=self.dtype) *
+           math_ops.square(self.scale) * self.df / denom)
+    # When 1 < df <= 2, variance is infinite.
+    inf = np.array(np.inf, dtype=self.dtype.as_numpy_dtype())
+    result_where_defined = array_ops.where_v2(
+        self.df > array_ops.fill(self.batch_shape_tensor(), 2.), var,
+        array_ops.fill(self.batch_shape_tensor(), inf, name="inf"))
+
+    if self.allow_nan_stats:
+      nan = np.array(np.nan, dtype=self.dtype.as_numpy_dtype())
+      return array_ops.where_v2(
+          math_ops.greater(
+              self.df,
+              array_ops.ones(self.batch_shape_tensor(), dtype=self.dtype)),
+          result_where_defined,
+          array_ops.fill(self.batch_shape_tensor(), nan, name="nan"))
+    else:
+      return control_flow_ops.with_dependencies(
+          [
+              check_ops.assert_less(
+                  array_ops.ones([], dtype=self.dtype),
+                  self.df,
+                  message="variance not defined for components of df <= 1"),
+          ],
+          result_where_defined)
+
+  def _mode(self):
+    return array_ops.identity(self.loc)
+
+
+class StudentTWithAbsDfSoftplusScale(StudentT):
+  """StudentT with `df = floor(abs(df))` and `scale = softplus(scale)`."""
+
+  @deprecation.deprecated(
+      "2019-01-01",
+      "Use `tfd.StudentT(tf.floor(tf.abs(df)), loc, "
+      "tf.nn.softplus(scale)) instead.",
+      warn_once=True)
+  def __init__(self,
+               df,
+               loc,
+               scale,
+               validate_args=False,
+               allow_nan_stats=True,
+               name="StudentTWithAbsDfSoftplusScale"):
+    parameters = dict(locals())
+    with ops.name_scope(name, values=[df, scale]) as name:
+      super(StudentTWithAbsDfSoftplusScale, self).__init__(
+          df=math_ops.floor(math_ops.abs(df)),
+          loc=loc,
+          scale=nn.softplus(scale, name="softplus_scale"),
+          validate_args=validate_args,
+          allow_nan_stats=allow_nan_stats,
+          name=name)
+    self._parameters = parameters

tensorflow/python/ops/distributions/transformed_distribution.py

@@ -0,0 +1,642 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""A Transformed Distribution class."""
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.distributions import distribution as distribution_lib
+from tensorflow.python.ops.distributions import identity_bijector
+from tensorflow.python.ops.distributions import util as distribution_util
+from tensorflow.python.util import deprecation
+
+__all__ = [
+    "TransformedDistribution",
+]
+
+
+# The following helper functions attempt to statically perform a TF operation.
+# These functions make debugging easier since we can do more validation during
+# graph construction.
+
+
+def _static_value(x):
+  """Returns the static value of a `Tensor` or `None`."""
+  return tensor_util.constant_value(ops.convert_to_tensor(x))
+
+
+def _logical_and(*args):
+  """Convenience function which attempts to statically `reduce_all`."""
+  args_ = [_static_value(x) for x in args]
+  if any(x is not None and not bool(x) for x in args_):
+    return constant_op.constant(False)
+  if all(x is not None and bool(x) for x in args_):
+    return constant_op.constant(True)
+  if len(args) == 2:
+    return math_ops.logical_and(*args)
+  return math_ops.reduce_all(args)
+
+
+def _logical_equal(x, y):
+  """Convenience function which attempts to statically compute `x == y`."""
+  x_ = _static_value(x)
+  y_ = _static_value(y)
+  if x_ is None or y_ is None:
+    return math_ops.equal(x, y)
+  return constant_op.constant(np.array_equal(x_, y_))
+
+
+def _logical_not(x):
+  """Convenience function which attempts to statically apply `logical_not`."""
+  x_ = _static_value(x)
+  if x_ is None:
+    return math_ops.logical_not(x)
+  return constant_op.constant(np.logical_not(x_))
+
+
+def _concat_vectors(*args):
+  """Convenience function which concatenates input vectors."""
+  args_ = [_static_value(x) for x in args]
+  if any(x_ is None for x_ in args_):
+    return array_ops.concat(args, 0)
+  return constant_op.constant([x_ for vec_ in args_ for x_ in vec_])
+
+
+def _pick_scalar_condition(pred, cond_true, cond_false):
+  """Convenience function which chooses the condition based on the predicate."""
+  # Note: This function is only valid if all of pred, cond_true, and cond_false
+  # are scalars. This means its semantics are arguably more like tf.cond than
+  # tf.select even though we use tf.select to implement it.
+  pred_ = _static_value(pred)
+  if pred_ is None:
+    return array_ops.where_v2(pred, cond_true, cond_false)
+  return cond_true if pred_ else cond_false
+
+
+def _ones_like(x):
+  """Convenience function attempts to statically construct `ones_like`."""
+  # Should only be used for small vectors.
+  if x.get_shape().is_fully_defined():
+    return array_ops.ones(x.get_shape().as_list(), dtype=x.dtype)
+  return array_ops.ones_like(x)
+
+
+def _ndims_from_shape(shape):
+  """Returns `Tensor`'s `rank` implied by a `Tensor` shape."""
+  if shape.get_shape().ndims not in (None, 1):
+    raise ValueError("input is not a valid shape: not 1D")
+  if not shape.dtype.is_integer:
+    raise TypeError("input is not a valid shape: wrong dtype")
+  if shape.get_shape().is_fully_defined():
+    return constant_op.constant(shape.get_shape().as_list()[0])
+  return array_ops.shape(shape)[0]
+
+
+def _is_scalar_from_shape(shape):
+  """Returns `True` `Tensor` if `Tensor` shape implies a scalar."""
+  return _logical_equal(_ndims_from_shape(shape), 0)
+
+
+class TransformedDistribution(distribution_lib.Distribution):
+  """A Transformed Distribution.
+
+  A `TransformedDistribution` models `p(y)` given a base distribution `p(x)`,
+  and a deterministic, invertible, differentiable transform, `Y = g(X)`. The
+  transform is typically an instance of the `Bijector` class and the base
+  distribution is typically an instance of the `Distribution` class.
+
+  A `Bijector` is expected to implement the following functions:
+  - `forward`,
+  - `inverse`,
+  - `inverse_log_det_jacobian`.
+  The semantics of these functions are outlined in the `Bijector` documentation.
+
+  We now describe how a `TransformedDistribution` alters the input/outputs of a
+  `Distribution` associated with a random variable (rv) `X`.
+
+  Write `cdf(Y=y)` for an absolutely continuous cumulative distribution function
+  of random variable `Y`; write the probability density function `pdf(Y=y) :=
+  d^k / (dy_1,...,dy_k) cdf(Y=y)` for its derivative wrt to `Y` evaluated at
+  `y`. Assume that `Y = g(X)` where `g` is a deterministic diffeomorphism,
+  i.e., a non-random, continuous, differentiable, and invertible function.
+  Write the inverse of `g` as `X = g^{-1}(Y)` and `(J o g)(x)` for the Jacobian
+  of `g` evaluated at `x`.
+
+  A `TransformedDistribution` implements the following operations:
+
+    * `sample`
+      Mathematically:   `Y = g(X)`
+      Programmatically: `bijector.forward(distribution.sample(...))`
+
+    * `log_prob`
+      Mathematically:   `(log o pdf)(Y=y) = (log o pdf o g^{-1})(y)
+                         + (log o abs o det o J o g^{-1})(y)`
+      Programmatically: `(distribution.log_prob(bijector.inverse(y))
+                         + bijector.inverse_log_det_jacobian(y))`
+
+    * `log_cdf`
+      Mathematically:   `(log o cdf)(Y=y) = (log o cdf o g^{-1})(y)`
+      Programmatically: `distribution.log_cdf(bijector.inverse(x))`
+
+    * and similarly for: `cdf`, `prob`, `log_survival_function`,
+     `survival_function`.
+
+  A simple example constructing a Log-Normal distribution from a Normal
+  distribution:
+
+  ```python
+  ds = tfp.distributions
+  log_normal = ds.TransformedDistribution(
+    distribution=ds.Normal(loc=0., scale=1.),
+    bijector=ds.bijectors.Exp(),
+    name="LogNormalTransformedDistribution")
+  ```
+
+  A `LogNormal` made from callables:
+
+  ```python
+  ds = tfp.distributions
+  log_normal = ds.TransformedDistribution(
+    distribution=ds.Normal(loc=0., scale=1.),
+    bijector=ds.bijectors.Inline(
+      forward_fn=tf.exp,
+      inverse_fn=tf.math.log,
+      inverse_log_det_jacobian_fn=(
+        lambda y: -tf.reduce_sum(tf.math.log(y), axis=-1)),
+    name="LogNormalTransformedDistribution")
+  ```
+
+  Another example constructing a Normal from a StandardNormal:
+
+  ```python
+  ds = tfp.distributions
+  normal = ds.TransformedDistribution(
+    distribution=ds.Normal(loc=0., scale=1.),
+    bijector=ds.bijectors.Affine(
+      shift=-1.,
+      scale_identity_multiplier=2.)
+    name="NormalTransformedDistribution")
+  ```
+
+  A `TransformedDistribution`'s batch- and event-shape are implied by the base
+  distribution unless explicitly overridden by `batch_shape` or `event_shape`
+  arguments. Specifying an overriding `batch_shape` (`event_shape`) is
+  permitted only if the base distribution has scalar batch-shape (event-shape).
+  The bijector is applied to the distribution as if the distribution possessed
+  the overridden shape(s). The following example demonstrates how to construct a
+  multivariate Normal as a `TransformedDistribution`.
+
+  ```python
+  ds = tfp.distributions
+  # We will create two MVNs with batch_shape = event_shape = 2.
+  mean = [[-1., 0],      # batch:0
+          [0., 1]]       # batch:1
+  chol_cov = [[[1., 0],
+               [0, 1]],  # batch:0
+              [[1, 0],
+               [2, 2]]]  # batch:1
+  mvn1 = ds.TransformedDistribution(
+      distribution=ds.Normal(loc=0., scale=1.),
+      bijector=ds.bijectors.Affine(shift=mean, scale_tril=chol_cov),
+      batch_shape=[2],  # Valid because base_distribution.batch_shape == [].
+      event_shape=[2])  # Valid because base_distribution.event_shape == [].
+  mvn2 = ds.MultivariateNormalTriL(loc=mean, scale_tril=chol_cov)
+  # mvn1.log_prob(x) == mvn2.log_prob(x)
+  ```
+
+  """
+
+  @deprecation.deprecated(
+      "2019-01-01",
+      "The TensorFlow Distributions library has moved to "
+      "TensorFlow Probability "
+      "(https://github.com/tensorflow/probability). You "
+      "should update all references to use `tfp.distributions` "
+      "instead of `tf.distributions`.",
+      warn_once=True)
+  def __init__(self,
+               distribution,
+               bijector=None,
+               batch_shape=None,
+               event_shape=None,
+               validate_args=False,
+               name=None):
+    """Construct a Transformed Distribution.
+
+    Args:
+      distribution: The base distribution instance to transform. Typically an
+        instance of `Distribution`.
+      bijector: The object responsible for calculating the transformation.
+        Typically an instance of `Bijector`. `None` means `Identity()`.
+      batch_shape: `integer` vector `Tensor` which overrides `distribution`
+        `batch_shape`; valid only if `distribution.is_scalar_batch()`.
+      event_shape: `integer` vector `Tensor` which overrides `distribution`
+        `event_shape`; valid only if `distribution.is_scalar_event()`.
+      validate_args: Python `bool`, default `False`. When `True` distribution
+        parameters are checked for validity despite possibly degrading runtime
+        performance. When `False` invalid inputs may silently render incorrect
+        outputs.
+      name: Python `str` name prefixed to Ops created by this class. Default:
+        `bijector.name + distribution.name`.
+    """
+    parameters = dict(locals())
+    name = name or (("" if bijector is None else bijector.name) +
+                    distribution.name)
+    with ops.name_scope(name, values=[event_shape, batch_shape]) as name:
+      # For convenience we define some handy constants.
+      self._zero = constant_op.constant(0, dtype=dtypes.int32, name="zero")
+      self._empty = constant_op.constant([], dtype=dtypes.int32, name="empty")
+
+      if bijector is None:
+        bijector = identity_bijector.Identity(validate_args=validate_args)
+
+      # We will keep track of a static and dynamic version of
+      # self._is_{batch,event}_override. This way we can do more prior to graph
+      # execution, including possibly raising Python exceptions.
+
+      self._override_batch_shape = self._maybe_validate_shape_override(
+          batch_shape, distribution.is_scalar_batch(), validate_args,
+          "batch_shape")
+      self._is_batch_override = _logical_not(_logical_equal(
+          _ndims_from_shape(self._override_batch_shape), self._zero))
+      self._is_maybe_batch_override = bool(
+          tensor_util.constant_value(self._override_batch_shape) is None or
+          tensor_util.constant_value(self._override_batch_shape).size != 0)
+
+      self._override_event_shape = self._maybe_validate_shape_override(
+          event_shape, distribution.is_scalar_event(), validate_args,
+          "event_shape")
+      self._is_event_override = _logical_not(_logical_equal(
+          _ndims_from_shape(self._override_event_shape), self._zero))
+      self._is_maybe_event_override = bool(
+          tensor_util.constant_value(self._override_event_shape) is None or
+          tensor_util.constant_value(self._override_event_shape).size != 0)
+
+      # To convert a scalar distribution into a multivariate distribution we
+      # will draw dims from the sample dims, which are otherwise iid. This is
+      # easy to do except in the case that the base distribution has batch dims
+      # and we're overriding event shape. When that case happens the event dims
+      # will incorrectly be to the left of the batch dims. In this case we'll
+      # cyclically permute left the new dims.
+      self._needs_rotation = _logical_and(
+          self._is_event_override,
+          _logical_not(self._is_batch_override),
+          _logical_not(distribution.is_scalar_batch()))
+      override_event_ndims = _ndims_from_shape(self._override_event_shape)
+      self._rotate_ndims = _pick_scalar_condition(
+          self._needs_rotation, override_event_ndims, 0)
+      # We'll be reducing the head dims (if at all), i.e., this will be []
+      # if we don't need to reduce.
+      self._reduce_event_indices = math_ops.range(
+          self._rotate_ndims - override_event_ndims, self._rotate_ndims)
+
+    self._distribution = distribution
+    self._bijector = bijector
+    super(TransformedDistribution, self).__init__(
+        dtype=self._distribution.dtype,
+        reparameterization_type=self._distribution.reparameterization_type,
+        validate_args=validate_args,
+        allow_nan_stats=self._distribution.allow_nan_stats,
+        parameters=parameters,
+        # We let TransformedDistribution access _graph_parents since this class
+        # is more like a baseclass than derived.
+        graph_parents=(distribution._graph_parents +  # pylint: disable=protected-access
+                       bijector.graph_parents),
+        name=name)
+
+  @property
+  def distribution(self):
+    """Base distribution, p(x)."""
+    return self._distribution
+
+  @property
+  def bijector(self):
+    """Function transforming x => y."""
+    return self._bijector
+
+  def _event_shape_tensor(self):
+    return self.bijector.forward_event_shape_tensor(
+        distribution_util.pick_vector(
+            self._is_event_override,
+            self._override_event_shape,
+            self.distribution.event_shape_tensor()))
+
+  def _event_shape(self):
+    # If there's a chance that the event_shape has been overridden, we return
+    # what we statically know about the `event_shape_override`. This works
+    # because: `_is_maybe_event_override` means `static_override` is `None` or a
+    # non-empty list, i.e., we don't statically know the `event_shape` or we do.
+    #
+    # Since the `bijector` may change the `event_shape`, we then forward what we
+    # know to the bijector. This allows the `bijector` to have final say in the
+    # `event_shape`.
+    static_override = tensor_util.constant_value_as_shape(
+        self._override_event_shape)
+    return self.bijector.forward_event_shape(
+        static_override
+        if self._is_maybe_event_override
+        else self.distribution.event_shape)
+
+  def _batch_shape_tensor(self):
+    return distribution_util.pick_vector(
+        self._is_batch_override,
+        self._override_batch_shape,
+        self.distribution.batch_shape_tensor())
+
+  def _batch_shape(self):
+    # If there's a chance that the batch_shape has been overridden, we return
+    # what we statically know about the `batch_shape_override`. This works
+    # because: `_is_maybe_batch_override` means `static_override` is `None` or a
+    # non-empty list, i.e., we don't statically know the `batch_shape` or we do.
+    #
+    # Notice that this implementation parallels the `_event_shape` except that
+    # the `bijector` doesn't get to alter the `batch_shape`. Recall that
+    # `batch_shape` is a property of a distribution while `event_shape` is
+    # shared between both the `distribution` instance and the `bijector`.
+    static_override = tensor_util.constant_value_as_shape(
+        self._override_batch_shape)
+    return (static_override
+            if self._is_maybe_batch_override
+            else self.distribution.batch_shape)
+
+  def _sample_n(self, n, seed=None):
+    sample_shape = _concat_vectors(
+        distribution_util.pick_vector(self._needs_rotation, self._empty, [n]),
+        self._override_batch_shape,
+        self._override_event_shape,
+        distribution_util.pick_vector(self._needs_rotation, [n], self._empty))
+    x = self.distribution.sample(sample_shape=sample_shape, seed=seed)
+    x = self._maybe_rotate_dims(x)
+    # We'll apply the bijector in the `_call_sample_n` function.
+    return x
+
+  def _call_sample_n(self, sample_shape, seed, name, **kwargs):
+    # We override `_call_sample_n` rather than `_sample_n` so we can ensure that
+    # the result of `self.bijector.forward` is not modified (and thus caching
+    # works).
+    with self._name_scope(name, values=[sample_shape]):
+      sample_shape = ops.convert_to_tensor(
+          sample_shape, dtype=dtypes.int32, name="sample_shape")
+      sample_shape, n = self._expand_sample_shape_to_vector(
+          sample_shape, "sample_shape")
+
+      # First, generate samples. We will possibly generate extra samples in the
+      # event that we need to reinterpret the samples as part of the
+      # event_shape.
+      x = self._sample_n(n, seed, **kwargs)
+
+      # Next, we reshape `x` into its final form. We do this prior to the call
+      # to the bijector to ensure that the bijector caching works.
+      batch_event_shape = array_ops.shape(x)[1:]
+      final_shape = array_ops.concat([sample_shape, batch_event_shape], 0)
+      x = array_ops.reshape(x, final_shape)
+
+      # Finally, we apply the bijector's forward transformation. For caching to
+      # work, it is imperative that this is the last modification to the
+      # returned result.
+      y = self.bijector.forward(x, **kwargs)
+      y = self._set_sample_static_shape(y, sample_shape)
+
+      return y
+
+  def _log_prob(self, y):
+    # For caching to work, it is imperative that the bijector is the first to
+    # modify the input.
+    x = self.bijector.inverse(y)
+    event_ndims = self._maybe_get_static_event_ndims()
+
+    ildj = self.bijector.inverse_log_det_jacobian(y, event_ndims=event_ndims)
+    if self.bijector._is_injective:  # pylint: disable=protected-access
+      return self._finish_log_prob_for_one_fiber(y, x, ildj, event_ndims)
+
+    lp_on_fibers = [
+        self._finish_log_prob_for_one_fiber(y, x_i, ildj_i, event_ndims)
+        for x_i, ildj_i in zip(x, ildj)]
+    return math_ops.reduce_logsumexp(
+        array_ops_stack.stack(lp_on_fibers), axis=0)
+
+  def _finish_log_prob_for_one_fiber(self, y, x, ildj, event_ndims):
+    """Finish computation of log_prob on one element of the inverse image."""
+    x = self._maybe_rotate_dims(x, rotate_right=True)
+    log_prob = self.distribution.log_prob(x)
+    if self._is_maybe_event_override:
+      log_prob = math_ops.reduce_sum(log_prob, self._reduce_event_indices)
+    log_prob += math_ops.cast(ildj, log_prob.dtype)
+    if self._is_maybe_event_override and isinstance(event_ndims, int):
+      log_prob.set_shape(
+          array_ops.broadcast_static_shape(
+              y.get_shape().with_rank_at_least(1)[:-event_ndims],
+              self.batch_shape))
+    return log_prob
+
+  def _prob(self, y):
+    x = self.bijector.inverse(y)
+    event_ndims = self._maybe_get_static_event_ndims()
+    ildj = self.bijector.inverse_log_det_jacobian(y, event_ndims=event_ndims)
+    if self.bijector._is_injective:  # pylint: disable=protected-access
+      return self._finish_prob_for_one_fiber(y, x, ildj, event_ndims)
+
+    prob_on_fibers = [
+        self._finish_prob_for_one_fiber(y, x_i, ildj_i, event_ndims)
+        for x_i, ildj_i in zip(x, ildj)]
+    return sum(prob_on_fibers)
+
+  def _finish_prob_for_one_fiber(self, y, x, ildj, event_ndims):
+    """Finish computation of prob on one element of the inverse image."""
+    x = self._maybe_rotate_dims(x, rotate_right=True)
+    prob = self.distribution.prob(x)
+    if self._is_maybe_event_override:
+      prob = math_ops.reduce_prod(prob, self._reduce_event_indices)
+    prob *= math_ops.exp(math_ops.cast(ildj, prob.dtype))
+    if self._is_maybe_event_override and isinstance(event_ndims, int):
+      prob.set_shape(
+          array_ops.broadcast_static_shape(
+              y.get_shape().with_rank_at_least(1)[:-event_ndims],
+              self.batch_shape))
+    return prob
+
+  def _log_cdf(self, y):
+    if self._is_maybe_event_override:
+      raise NotImplementedError("log_cdf is not implemented when overriding "
+                                "event_shape")
+    if not self.bijector._is_injective:  # pylint: disable=protected-access
+      raise NotImplementedError("log_cdf is not implemented when "
+                                "bijector is not injective.")
+    x = self.bijector.inverse(y)
+    return self.distribution.log_cdf(x)
+
+  def _cdf(self, y):
+    if self._is_maybe_event_override:
+      raise NotImplementedError("cdf is not implemented when overriding "
+                                "event_shape")
+    if not self.bijector._is_injective:  # pylint: disable=protected-access
+      raise NotImplementedError("cdf is not implemented when "
+                                "bijector is not injective.")
+    x = self.bijector.inverse(y)
+    return self.distribution.cdf(x)
+
+  def _log_survival_function(self, y):
+    if self._is_maybe_event_override:
+      raise NotImplementedError("log_survival_function is not implemented when "
+                                "overriding event_shape")
+    if not self.bijector._is_injective:  # pylint: disable=protected-access
+      raise NotImplementedError("log_survival_function is not implemented when "
+                                "bijector is not injective.")
+    x = self.bijector.inverse(y)
+    return self.distribution.log_survival_function(x)
+
+  def _survival_function(self, y):
+    if self._is_maybe_event_override:
+      raise NotImplementedError("survival_function is not implemented when "
+                                "overriding event_shape")
+    if not self.bijector._is_injective:  # pylint: disable=protected-access
+      raise NotImplementedError("survival_function is not implemented when "
+                                "bijector is not injective.")
+    x = self.bijector.inverse(y)
+    return self.distribution.survival_function(x)
+
+  def _quantile(self, value):
+    if self._is_maybe_event_override:
+      raise NotImplementedError("quantile is not implemented when overriding "
+                                "event_shape")
+    if not self.bijector._is_injective:  # pylint: disable=protected-access
+      raise NotImplementedError("quantile is not implemented when "
+                                "bijector is not injective.")
+    # x_q is the "qth quantile" of X iff q = P[X <= x_q].  Now, since X =
+    # g^{-1}(Y), q = P[X <= x_q] = P[g^{-1}(Y) <= x_q] = P[Y <= g(x_q)],
+    # implies the qth quantile of Y is g(x_q).
+    inv_cdf = self.distribution.quantile(value)
+    return self.bijector.forward(inv_cdf)
+
+  def _entropy(self):
+    if not self.bijector.is_constant_jacobian:
+      raise NotImplementedError("entropy is not implemented")
+    if not self.bijector._is_injective:  # pylint: disable=protected-access
+      raise NotImplementedError("entropy is not implemented when "
+                                "bijector is not injective.")
+    # Suppose Y = g(X) where g is a diffeomorphism and X is a continuous rv. It
+    # can be shown that:
+    #   H[Y] = H[X] + E_X[(log o abs o det o J o g)(X)].
+    # If is_constant_jacobian then:
+    #   E_X[(log o abs o det o J o g)(X)] = (log o abs o det o J o g)(c)
+    # where c can by anything.
+    entropy = self.distribution.entropy()
+    if self._is_maybe_event_override:
+      # H[X] = sum_i H[X_i] if X_i are mutually independent.
+      # This means that a reduce_sum is a simple rescaling.
+      entropy *= math_ops.cast(math_ops.reduce_prod(self._override_event_shape),
+                               dtype=entropy.dtype.base_dtype)
+    if self._is_maybe_batch_override:
+      new_shape = array_ops.concat([
+          _ones_like(self._override_batch_shape),
+          self.distribution.batch_shape_tensor()
+      ], 0)
+      entropy = array_ops.reshape(entropy, new_shape)
+      multiples = array_ops.concat([
+          self._override_batch_shape,
+          _ones_like(self.distribution.batch_shape_tensor())
+      ], 0)
+      entropy = array_ops.tile(entropy, multiples)
+    dummy = array_ops.zeros(
+        shape=array_ops.concat(
+            [self.batch_shape_tensor(), self.event_shape_tensor()],
+            0),
+        dtype=self.dtype)
+    event_ndims = (self.event_shape.ndims if self.event_shape.ndims is not None
+                   else array_ops.size(self.event_shape_tensor()))
+    ildj = self.bijector.inverse_log_det_jacobian(
+        dummy, event_ndims=event_ndims)
+
+    entropy -= math_ops.cast(ildj, entropy.dtype)
+    entropy.set_shape(self.batch_shape)
+    return entropy
+
+  def _maybe_validate_shape_override(self, override_shape, base_is_scalar,
+                                     validate_args, name):
+    """Helper to __init__ which ensures override batch/event_shape are valid."""
+    if override_shape is None:
+      override_shape = []
+
+    override_shape = ops.convert_to_tensor(override_shape, dtype=dtypes.int32,
+                                           name=name)
+
+    if not override_shape.dtype.is_integer:
+      raise TypeError("shape override must be an integer")
+
+    override_is_scalar = _is_scalar_from_shape(override_shape)
+    if tensor_util.constant_value(override_is_scalar):
+      return self._empty
+
+    dynamic_assertions = []
+
+    if override_shape.get_shape().ndims is not None:
+      if override_shape.get_shape().ndims != 1:
+        raise ValueError("shape override must be a vector")
+    elif validate_args:
+      dynamic_assertions += [check_ops.assert_rank(
+          override_shape, 1,
+          message="shape override must be a vector")]
+
+    if tensor_util.constant_value(override_shape) is not None:
+      if any(s <= 0 for s in tensor_util.constant_value(override_shape)):
+        raise ValueError("shape override must have positive elements")
+    elif validate_args:
+      dynamic_assertions += [check_ops.assert_positive(
+          override_shape,
+          message="shape override must have positive elements")]
+
+    is_both_nonscalar = _logical_and(_logical_not(base_is_scalar),
+                                     _logical_not(override_is_scalar))
+    if tensor_util.constant_value(is_both_nonscalar) is not None:
+      if tensor_util.constant_value(is_both_nonscalar):
+        raise ValueError("base distribution not scalar")
+    elif validate_args:
+      dynamic_assertions += [check_ops.assert_equal(
+          is_both_nonscalar, False,
+          message="base distribution not scalar")]
+
+    if not dynamic_assertions:
+      return override_shape
+    return control_flow_ops.with_dependencies(
+        dynamic_assertions, override_shape)
+
+  def _maybe_rotate_dims(self, x, rotate_right=False):
+    """Helper which rolls left event_dims left or right event_dims right."""
+    needs_rotation_const = tensor_util.constant_value(self._needs_rotation)
+    if needs_rotation_const is not None and not needs_rotation_const:
+      return x
+    ndims = array_ops.rank(x)
+    n = (ndims - self._rotate_ndims) if rotate_right else self._rotate_ndims
+    return array_ops.transpose(
+        x, _concat_vectors(math_ops.range(n, ndims), math_ops.range(0, n)))
+
+  def _maybe_get_static_event_ndims(self):
+    if self.event_shape.ndims is not None:
+      return self.event_shape.ndims
+
+    event_ndims = array_ops.size(self.event_shape_tensor())
+    event_ndims_ = distribution_util.maybe_get_static_value(event_ndims)
+
+    if event_ndims_ is not None:
+      return event_ndims_
+
+    return event_ndims

tensorflow/python/ops/distributions/uniform.py

@@ -0,0 +1,204 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""The Uniform distribution class."""
+
+import math
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops.distributions import distribution
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["distributions.Uniform"])
+class Uniform(distribution.Distribution):
+  """Uniform distribution with `low` and `high` parameters.
+
+  #### Mathematical Details
+
+  The probability density function (pdf) is,
+
+  ```none
+  pdf(x; a, b) = I[a <= x < b] / Z
+  Z = b - a
+  ```
+
+  where
+
+  - `low = a`,
+  - `high = b`,
+  - `Z` is the normalizing constant, and
+  - `I[predicate]` is the [indicator function](
+    https://en.wikipedia.org/wiki/Indicator_function) for `predicate`.
+
+  The parameters `low` and `high` must be shaped in a way that supports
+  broadcasting (e.g., `high - low` is a valid operation).
+
+  #### Examples
+
+  ```python
+  # Without broadcasting:
+  u1 = Uniform(low=3.0, high=4.0)  # a single uniform distribution [3, 4]
+  u2 = Uniform(low=[1.0, 2.0],
+               high=[3.0, 4.0])  # 2 distributions [1, 3], [2, 4]
+  u3 = Uniform(low=[[1.0, 2.0],
+                    [3.0, 4.0]],
+               high=[[1.5, 2.5],
+                     [3.5, 4.5]])  # 4 distributions
+  ```
+
+  ```python
+  # With broadcasting:
+  u1 = Uniform(low=3.0, high=[5.0, 6.0, 7.0])  # 3 distributions
+  ```
+
+  """
+
+  @deprecation.deprecated(
+      "2019-01-01",
+      "The TensorFlow Distributions library has moved to "
+      "TensorFlow Probability "
+      "(https://github.com/tensorflow/probability). You "
+      "should update all references to use `tfp.distributions` "
+      "instead of `tf.distributions`.",
+      warn_once=True)
+  def __init__(self,
+               low=0.,
+               high=1.,
+               validate_args=False,
+               allow_nan_stats=True,
+               name="Uniform"):
+    """Initialize a batch of Uniform distributions.
+
+    Args:
+      low: Floating point tensor, lower boundary of the output interval. Must
+        have `low < high`.
+      high: Floating point tensor, upper boundary of the output interval. Must
+        have `low < high`.
+      validate_args: Python `bool`, default `False`. When `True` distribution
+        parameters are checked for validity despite possibly degrading runtime
+        performance. When `False` invalid inputs may silently render incorrect
+        outputs.
+      allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
+        (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
+        result is undefined. When `False`, an exception is raised if one or
+        more of the statistic's batch members are undefined.
+      name: Python `str` name prefixed to Ops created by this class.
+
+    Raises:
+      InvalidArgumentError: if `low >= high` and `validate_args=False`.
+    """
+    parameters = dict(locals())
+    with ops.name_scope(name, values=[low, high]) as name:
+      with ops.control_dependencies([
+          check_ops.assert_less(
+              low, high, message="uniform not defined when low >= high.")
+      ] if validate_args else []):
+        self._low = array_ops.identity(low, name="low")
+        self._high = array_ops.identity(high, name="high")
+        check_ops.assert_same_float_dtype([self._low, self._high])
+    super(Uniform, self).__init__(
+        dtype=self._low.dtype,
+        reparameterization_type=distribution.FULLY_REPARAMETERIZED,
+        validate_args=validate_args,
+        allow_nan_stats=allow_nan_stats,
+        parameters=parameters,
+        graph_parents=[self._low,
+                       self._high],
+        name=name)
+
+  @staticmethod
+  def _param_shapes(sample_shape):
+    return dict(
+        zip(("low", "high"),
+            ([ops.convert_to_tensor(sample_shape, dtype=dtypes.int32)] * 2)))
+
+  @property
+  def low(self):
+    """Lower boundary of the output interval."""
+    return self._low
+
+  @property
+  def high(self):
+    """Upper boundary of the output interval."""
+    return self._high
+
+  def range(self, name="range"):
+    """`high - low`."""
+    with self._name_scope(name):
+      return self.high - self.low
+
+  def _batch_shape_tensor(self):
+    return array_ops.broadcast_dynamic_shape(
+        array_ops.shape(self.low),
+        array_ops.shape(self.high))
+
+  def _batch_shape(self):
+    return array_ops.broadcast_static_shape(
+        self.low.get_shape(),
+        self.high.get_shape())
+
+  def _event_shape_tensor(self):
+    return constant_op.constant([], dtype=dtypes.int32)
+
+  def _event_shape(self):
+    return tensor_shape.TensorShape([])
+
+  def _sample_n(self, n, seed=None):
+    shape = array_ops.concat([[n], self.batch_shape_tensor()], 0)
+    samples = random_ops.random_uniform(shape=shape,
+                                        dtype=self.dtype,
+                                        seed=seed)
+    return self.low + self.range() * samples
+
+  def _prob(self, x):
+    broadcasted_x = x * array_ops.ones(
+        self.batch_shape_tensor(), dtype=x.dtype)
+    return array_ops.where_v2(
+        math_ops.is_nan(broadcasted_x), broadcasted_x,
+        array_ops.where_v2(
+            math_ops.logical_or(broadcasted_x < self.low,
+                                broadcasted_x >= self.high),
+            array_ops.zeros_like(broadcasted_x),
+            array_ops.ones_like(broadcasted_x) / self.range()))
+
+  def _cdf(self, x):
+    broadcast_shape = array_ops.broadcast_dynamic_shape(
+        array_ops.shape(x), self.batch_shape_tensor())
+    zeros = array_ops.zeros(broadcast_shape, dtype=self.dtype)
+    ones = array_ops.ones(broadcast_shape, dtype=self.dtype)
+    broadcasted_x = x * ones
+    result_if_not_big = array_ops.where_v2(
+        x < self.low, zeros, (broadcasted_x - self.low) / self.range())
+    return array_ops.where_v2(x >= self.high, ones, result_if_not_big)
+
+  def _entropy(self):
+    return math_ops.log(self.range())
+
+  def _mean(self):
+    return (self.low + self.high) / 2.
+
+  def _variance(self):
+    return math_ops.square(self.range()) / 12.
+
+  def _stddev(self):
+    return self.range() / math.sqrt(12.)

tensorflow/python/ops/embedding_ops_test.py

@@ -0,0 +1,99 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for tensorflow.ops.embedding_ops."""
+
+import numpy as np
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import embedding_ops
+from tensorflow.python.ops import gradients
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class EmbeddingLookupTest(test_util.TensorFlowTestCase):
+
+  def testEmbeddingLookupOnUninitializedVariableDoesSparseRead(self):
+    x = resource_variable_ops.UninitializedVariable(
+        trainable=True, shape=[3, 3], dtype=dtypes.float32)
+
+    @def_function.function(input_signature=[])
+    def _init():
+      return x.assign(np.zeros([3, 3]))
+
+    @def_function.function(input_signature=[])
+    def _call():
+      return embedding_ops.embedding_lookup_v2(x, [0])
+
+    self.assertAllClose(self.evaluate(_init()), np.zeros([3, 3]))
+
+    concrete_call = _call.get_concrete_function()
+    self.assertAllClose(self.evaluate(concrete_call()), [[0., 0., 0.]])
+
+    resource_gather_node = []
+    read_var_node = []
+    graph = concrete_call.graph.as_graph_def()
+    for n in graph.node:
+      if n.op == "ResourceGather":
+        resource_gather_node.append(n)
+      if n.op == "ReadVariableOp":
+        read_var_node.append(n)
+
+    for f in graph.library.function:
+      for n in f.node_def:
+        if n.op == "ResourceGather":
+          resource_gather_node.append(n)
+        if n.op == "ReadVariableOp":
+          read_var_node.append(n)
+    # There should be a single ResourceGather, but no ReadVariableOp
+    # (dense read).
+    self.assertLen(resource_gather_node, 1)
+    self.assertLen(read_var_node, 0)
+
+  def testEmbeddingLookupGradientsHaveKnownShape(self):
+    x = resource_variable_ops.ResourceVariable(
+        initial_value=np.zeros([3, 3]),
+        trainable=True,
+        shape=[3, 3],
+        dtype=dtypes.float32)
+
+    @def_function.function(input_signature=[])
+    def _init():
+      return x.assign(np.zeros([3, 3]))
+
+    @def_function.function(input_signature=[])
+    def _call():
+      with gradients.GradientTape() as tape:
+        y = embedding_ops.embedding_lookup_v2(x, [0])
+        loss = math_ops.reduce_sum(y)
+      grads = tape.gradient(loss, x)
+      self.assertAllEqual(grads.shape, [3, 3])
+      return ops.convert_to_tensor(grads)
+
+    self.assertAllClose(self.evaluate(_init()), np.zeros([3, 3]))
+
+    concrete_call = _call.get_concrete_function()
+    self.assertAllClose(
+        self.evaluate(concrete_call()),
+        [[1., 1., 1.], [0., 0., 0.], [0., 0., 0.]])
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/python/ops/factory_ops_test.py

@@ -0,0 +1,78 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests that sparse tensors work with GPU, such as placement of int and string.
+
+Test using sparse tensors with distributed dataset. Since GPU does
+not support strings, sparse tensors containing string should always be placed
+on CPU.
+"""
+
+from absl.testing import parameterized
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import mirrored_strategy
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import sparse_ops
+from tensorflow.python.platform import test
+
+
+def sparse_int64():
+  return sparse_tensor.SparseTensor(
+      indices=[[0, 0], [1, 1], [2, 2], [3, 3], [4, 0], [5, 1], [6, 2], [7, 3]],
+      values=constant_op.constant([1, 2, 3, 4, 5, 6, 7, 8], dtype=dtypes.int64),
+      dense_shape=[8, 4])
+
+
+def sparse_str():
+  return sparse_tensor.SparseTensor(
+      indices=[[0, 0], [1, 1], [2, 2], [3, 3], [4, 0], [5, 1], [6, 2], [7, 3]],
+      values=constant_op.constant(['1', '2', '3', '4', '5', '6', '7', '8']),
+      dense_shape=[8, 4])
+
+
+class FactoryOpsTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  @parameterized.parameters(
+      (sparse_int64,),
+      (sparse_str,),
+  )
+  @test_util.run_gpu_only
+  def testSparseWithDistributedDataset(self, sparse_factory):
+
+    @def_function.function
+    def distributed_dataset_producer(t):
+      strategy = mirrored_strategy.MirroredStrategy(['GPU:0', 'GPU:1'])
+      sparse_ds = dataset_ops.Dataset.from_tensor_slices(t).batch(2)
+      dist_dataset = strategy.experimental_distribute_dataset(sparse_ds)
+      ds = iter(dist_dataset)
+      result = strategy.experimental_local_results(next(ds))[0]
+      # Reach the end of the iterator
+      for ignore in ds:  # pylint: disable=unused-variable
+        pass
+      return result
+
+    t = sparse_factory()
+
+    result = distributed_dataset_producer(t)
+    self.assertAllEqual(
+        self.evaluate(sparse_ops.sparse_tensor_to_dense(t)[0]),
+        self.evaluate(sparse_ops.sparse_tensor_to_dense(result)[0]))
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/ops/filesystem_ops.py

@@ -0,0 +1,38 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Filesystem related operations."""
+
+from tensorflow.python.ops import gen_filesystem_ops as _gen_filesystem_ops
+
+
+# pylint: disable=protected-access
+def filesystem_set_configuration(scheme, key, value, name=None):
+  """Set configuration of the file system.
+
+  Args:
+    scheme: File system scheme.
+    key: The name of the configuration option.
+    value: The value of the configuration option.
+    name: A name for the operation (optional).
+
+  Returns:
+    None.
+  """
+
+  return _gen_filesystem_ops.file_system_set_configuration(
+      scheme, key=key, value=value, name=name)
+
+
+# pylint: enable=protected-access

tensorflow/python/ops/functional_ops_test.py

@@ -0,0 +1,82 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for functional operations."""
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import function
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import functional_ops
+from tensorflow.python.platform import test
+
+
+class FunctionalOpsTest(test.TestCase):
+
+  def testIfWithDefun(self):
+    # Defun should only be used in graph mode
+    with ops.Graph().as_default():
+      @function.Defun(dtypes.float32)
+      def Then(x):
+        return x + 1
+
+      @function.Defun(dtypes.float32)
+      def Else(x):
+        return x - 1
+
+      inputs = [10.]
+      result = self.evaluate(functional_ops.If(False, inputs, Then, Else))
+      self.assertEqual([9.0], result)
+
+  def testIfWithFunction(self):
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec((), dtypes.float32)])
+    def Then(x):
+      return x + 1
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec((), dtypes.float32)])
+    def Else(x):
+      return x - 1
+
+    inputs = [10.]
+    then_cf = Then.get_concrete_function()
+    else_cf = Else.get_concrete_function()
+    result = self.evaluate(functional_ops.If(False, inputs, then_cf, else_cf))
+    self.assertEqual([9.0], result)
+
+  def testIfWithFunctionComposite(self):
+
+    signature = [tensor_spec.TensorSpec([], dtypes.float32)]
+    @def_function.function(input_signature=signature)
+    def Then(x):
+      return sparse_tensor.SparseTensor([[0]], [x + 1], [1])
+
+    @def_function.function(input_signature=signature)
+    def Else(x):
+      return sparse_tensor.SparseTensor([[0]], [x - 1], [1])
+
+    inputs = [10.]
+    then_cf = Then.get_concrete_function()
+    else_cf = Else.get_concrete_function()
+    result = functional_ops.If(False, inputs, then_cf, else_cf)
+    self.assertIsInstance(result, sparse_tensor.SparseTensor)
+    self.assertAllEqual([9.0], result.values)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/ops/gradient_checker.py

@@ -0,0 +1,394 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Gradient checker for any ops, graphs.
+
+The gradient checker verifies numerically that an op/graph properly
+computes the gradients
+"""
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gradients
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util import deprecation
+from tensorflow.python.util import numpy_compat
+from tensorflow.python.util.tf_export import tf_export
+
+
+def _product(t):
+  if isinstance(t, int):
+    return t
+  else:
+    y = 1
+    for x in t:
+      y *= x
+    return y
+
+
+def _extra_feeds(extra_feed_dict, new_feeds):
+  if not extra_feed_dict:
+    return new_feeds
+  r = {}
+  r.update(extra_feed_dict)
+  r.update(new_feeds)
+  return r
+
+
+def _compute_theoretical_jacobian(x, x_shape, x_data, dy, dy_shape, dx,
+                                  extra_feed_dict):
+  """Computes the theoretical Jacobian for dy/dx.
+
+  Computes the theoretical Jacobian using the ops generated by
+  compute_gradient().
+
+  Args:
+    x: the tensor "x".
+    x_shape: the dimensions of x as a tuple or an array of ints.
+    x_data: a numpy parray as the input data for x
+    dy: the tensor "dy".
+    dy_shape: the dimensions of dy as a tuple or an array of ints.
+    dx: Tensor or IndexedSlices representing dx
+    extra_feed_dict: dict that allows fixing specified tensor values
+      during the jacobian calculation.
+
+  Returns:
+    A 2-d numpy array representing the Jacobian for dy/dx. It has "x_size" rows
+    and "dy_size" columns where "x_size" is the number of elements in x and
+    "dy_size" is the number of elements in dy.
+
+  Raises:
+    ValueError: If `dy` is empty but the gradient is nonzero.
+  """
+  # Complex vectors are treated as vectors of twice as many reals.
+  if x.dtype.is_complex:
+    x_shape = tuple(x_shape) + (2,)
+  dy_factor = 2 if dy.dtype.is_complex else 1
+
+  # To compute the jacobian, we treat x and y as one-dimensional vectors.
+  x_size = _product(x_shape)
+  x_val_size = _product(x_shape[1:])  # This is used for sparse gradients
+  dy_size = _product(dy_shape) * dy_factor
+
+  # Allocate 2-D Jacobian, with x dimensions smashed into the first
+  # dimension and y dimensions smashed into the second.
+  jacobian = np.zeros((x_size, dy_size),
+                      dtype=x.dtype.real_dtype.as_numpy_dtype)
+
+  # For each of the entry of dy, we set this to be 1 and
+  # everything else to be 0 and compute the backprop -- this will give us one
+  # one column of the Jacobian matrix.
+  dy_data = np.zeros(dy_shape, dtype=dy.dtype.as_numpy_dtype)
+  dy_data_flat = dy_data.ravel().view(dy.dtype.real_dtype.as_numpy_dtype)
+  sess = ops.get_default_session()
+  for col in range(dy_size):
+    dy_data_flat[col] = 1
+    if isinstance(dx, indexed_slices.IndexedSlices):
+      backprop_indices, backprop_values = sess.run(
+          [dx.indices, dx.values],
+          feed_dict=_extra_feeds(extra_feed_dict, {x: x_data, dy: dy_data}))
+      for i, v in zip(backprop_indices, backprop_values):
+        r_begin = i * x_val_size
+        r_end = r_begin + x_val_size
+        jacobian[r_begin:r_end, col] += v.flat
+    else:
+      assert isinstance(dx, tensor.Tensor), "dx = " + str(dx)
+      backprop = sess.run(
+          dx, feed_dict=_extra_feeds(extra_feed_dict, {x: x_data, dy: dy_data}))
+      jacobian[:, col] = backprop.ravel().view(jacobian.dtype)
+    dy_data_flat[col] = 0
+
+  # If the output is empty, run the gradients at least once and make sure
+  # they produce zeros.
+  if not dy_size:
+    backprop = sess.run(
+        dx, feed_dict=_extra_feeds(extra_feed_dict, {x: x_data, dy: dy_data}))
+    if backprop.shape != x_data.shape:
+      raise ValueError("Empty gradient has wrong shape: expected %s, got %s" %
+                       (x_data.shape, backprop.shape))
+    if np.any(backprop):
+      raise ValueError("Empty tensor with nonzero gradients")
+
+  logging.vlog(1, "Theoretical Jacobian =\n%s", jacobian)
+  return jacobian
+
+
+def _compute_numeric_jacobian(x, x_shape, x_data, y, y_shape, delta,
+                              extra_feed_dict):
+  """Computes the numeric Jacobian for dy/dx.
+
+  Computes the numeric Jacobian by slightly perturbing the inputs and
+  measuring the differences on the output.
+
+  Args:
+    x: the tensor "x".
+    x_shape: the dimensions of x as a tuple or an array of ints.
+    x_data: a numpy array as the input data for x
+    y: the tensor "y".
+    y_shape: the dimensions of y as a tuple or an array of ints.
+    delta: the amount of perturbation we give to the input
+    extra_feed_dict: dict that allows fixing specified tensor values
+      during the jacobian calculation.
+
+  Returns:
+    A 2-d numpy array representing the Jacobian for dy/dx. It has "x_size" rows
+    and "y_size" columns where "x_size" is the number of elements in x and
+    "y_size" is the number of elements in y.
+  """
+  # bfloat16 doesn't have enough bits to represent high precision numbers such
+  # as delta. Convert to float32 here. Since numeric_jacobian is expected to
+  # be the groundtruth to compare against, it shouldn't lose any information.
+  if x.dtype == dtypes.bfloat16:
+    x = math_ops.cast(x, dtypes.float32)  # TODO(wangpeng): Now that the new x
+            # is an output of the old x, isn't feeding to the new x a mistake?
+  if y.dtype == dtypes.bfloat16:
+    y = math_ops.cast(y, dtypes.float32)
+  if x_data.dtype == dtypes.bfloat16.as_numpy_dtype:
+    x_data = x_data.astype(np.float32)
+
+  # To compute the jacobian, we treat x and y as one-dimensional vectors
+  x_size = _product(x_shape) * (2 if x.dtype.is_complex else 1)
+  y_size = _product(y_shape) * (2 if y.dtype.is_complex else 1)
+  x_dtype = x.dtype.real_dtype.as_numpy_dtype
+  y_dtype = y.dtype.real_dtype.as_numpy_dtype
+
+  # Make sure we have the right types
+  x_data = numpy_compat.np_asarray(x_data, dtype=x.dtype.as_numpy_dtype)
+  scale = numpy_compat.np_asarray(2 * delta, dtype=y_dtype)[()]
+
+  jacobian = np.zeros((x_size, y_size), dtype=x_dtype)
+  # For each of the entry of x, we slightly perturbs this by adding and
+  # subtracting a delta and then compute difference between the outputs. This
+  # will give us one row of the Jacobian matrix.
+  for row in range(x_size):
+    x_pos = x_data.copy()
+    x_neg = x_data.copy()
+    x_pos.ravel().view(x_dtype)[row] += delta
+    y_pos = y.eval(feed_dict=_extra_feeds(extra_feed_dict, {x: x_pos}))
+    x_neg.ravel().view(x_dtype)[row] -= delta
+    y_neg = y.eval(feed_dict=_extra_feeds(extra_feed_dict, {x: x_neg}))
+    diff = (y_pos - y_neg) / scale
+    jacobian[row, :] = diff.ravel().view(y_dtype)
+
+  logging.vlog(1, "Numeric Jacobian =\n%s", jacobian)
+  return jacobian
+
+
+def _compute_dx_and_dy(x, y, y_shape):
+  """Returns a node to compute gradient of y wrt x."""
+  # We make up a dy so that we can compute the gradients. We don't really use
+  # the value of dy -- we will always feed it. We need to add an identity node
+  # so that we can always feed it properly. Otherwise, for the Add operation,
+  # dx is the same as dy and we cannot fetch the tensor that we are feeding.
+  with x.graph.as_default():
+    dy_orig = constant_op.constant(1.0, shape=y_shape, dtype=y.dtype)
+    dy = array_ops.identity(dy_orig)
+  # We compute the gradients for y wrt. x
+  grads = gradients.gradients(y, x, dy)
+  assert len(grads) == 1
+  return grads[0], dy_orig
+
+
+def _compute_gradient(x,
+                      x_shape,
+                      dx,
+                      y,
+                      y_shape,
+                      dy,
+                      x_init_value=None,
+                      delta=1e-3,
+                      extra_feed_dict=None):
+  """Computes the theoretical and numerical jacobian."""
+  t = dtypes.as_dtype(x.dtype)
+  allowed_types = [dtypes.float16, dtypes.bfloat16, dtypes.float32,
+                   dtypes.float64, dtypes.complex64, dtypes.complex128]
+  assert t.base_dtype in allowed_types, "Don't support type %s for x" % t.name
+  t2 = dtypes.as_dtype(y.dtype)
+  assert t2.base_dtype in allowed_types, "Don't support type %s for y" % t2.name
+
+  if x_init_value is not None:
+    i_shape = list(x_init_value.shape)
+    assert(list(x_shape) == i_shape), "x_shape = %s, init_data shape = %s" % (
+        x_shape, i_shape)
+    x_data = x_init_value
+  else:
+    x_data = np.random.random_sample(x_shape).astype(t.as_numpy_dtype)
+    if t.is_complex:
+      x_data.imag = np.random.random_sample(x_shape)
+
+  jacob_t = _compute_theoretical_jacobian(
+      x, x_shape, x_data, dy, y_shape, dx, extra_feed_dict=extra_feed_dict)
+  jacob_n = _compute_numeric_jacobian(
+      x, x_shape, x_data, y, y_shape, delta, extra_feed_dict=extra_feed_dict)
+  return jacob_t, jacob_n
+
+
+def _compute_gradient_list(x,
+                           x_shape,
+                           y,
+                           y_shape,
+                           x_init_value=None,
+                           delta=1e-3,
+                           init_targets=None,
+                           extra_feed_dict=None):
+  """Compute gradients for a list of x values."""
+  assert isinstance(x, list)
+  dx, dy = zip(*[_compute_dx_and_dy(xi, y, y_shape) for xi in x])
+
+  if init_targets is not None:
+    assert isinstance(init_targets, (list, tuple))
+    for init in init_targets:
+      init.run()
+  if x_init_value is None:
+    x_init_value = [None] * len(x)
+  # pylint: disable=g-complex-comprehension
+  ret = [_compute_gradient(xi, x_shapei, dxi, y, y_shape, dyi, x_init_valuei,
+                           delta, extra_feed_dict=extra_feed_dict)
+         for xi, x_shapei, dxi, dyi, x_init_valuei in zip(x, x_shape, dx, dy,
+                                                          x_init_value)]
+  return ret
+
+
+@tf_export(v1=["test.compute_gradient"])
+@deprecation.deprecated(
+    date=None,
+    instructions="Use tf.test.compute_gradient in 2.0, which has better "
+    "support for functions. Note that the two versions have different usage, "
+    "so code change is needed.")
+def compute_gradient(x,
+                     x_shape,
+                     y,
+                     y_shape,
+                     x_init_value=None,
+                     delta=1e-3,
+                     init_targets=None,
+                     extra_feed_dict=None):
+  """Computes and returns the theoretical and numerical Jacobian.
+
+  If `x` or `y` is complex, the Jacobian will still be real but the
+  corresponding Jacobian dimension(s) will be twice as large.  This is required
+  even if both input and output is complex since TensorFlow graphs are not
+  necessarily holomorphic, and may have gradients not expressible as complex
+  numbers.  For example, if `x` is complex with shape `[m]` and `y` is complex
+  with shape `[n]`, each Jacobian `J` will have shape `[m * 2, n * 2]` with
+
+      J[:m, :n] = d(Re y)/d(Re x)
+      J[:m, n:] = d(Im y)/d(Re x)
+      J[m:, :n] = d(Re y)/d(Im x)
+      J[m:, n:] = d(Im y)/d(Im x)
+
+  Args:
+    x: a tensor or list of tensors
+    x_shape: the dimensions of x as a tuple or an array of ints. If x is a list,
+    then this is the list of shapes.
+    y: a tensor
+    y_shape: the dimensions of y as a tuple or an array of ints.
+    x_init_value: (optional) a numpy array of the same shape as "x"
+      representing the initial value of x. If x is a list, this should be a list
+      of numpy arrays.  If this is none, the function will pick a random tensor
+      as the initial value.
+    delta: (optional) the amount of perturbation.
+    init_targets: list of targets to run to initialize model params.
+    extra_feed_dict: dict that allows fixing specified tensor values
+      during the Jacobian calculation.
+
+  Returns:
+    Two 2-d numpy arrays representing the theoretical and numerical
+    Jacobian for dy/dx. Each has "x_size" rows and "y_size" columns
+    where "x_size" is the number of elements in x and "y_size" is the
+    number of elements in y. If x is a list, returns a list of two numpy arrays.
+  """
+  # TODO(mrry): remove argument `init_targets`
+  if extra_feed_dict is None:
+    extra_feed_dict = {}
+
+  if isinstance(x, list):
+    return _compute_gradient_list(x, x_shape, y, y_shape, x_init_value, delta,
+                                  init_targets, extra_feed_dict=extra_feed_dict)
+  else:
+    if init_targets is not None:
+      assert isinstance(init_targets, (list, tuple))
+      for init in init_targets:
+        init.run()
+    dx, dy = _compute_dx_and_dy(x, y, y_shape)
+    ret = _compute_gradient(x, x_shape, dx, y, y_shape, dy, x_init_value, delta,
+                            extra_feed_dict=extra_feed_dict)
+    return ret
+
+
+def _compute_error(grad):
+  if isinstance(grad, tuple):
+    grad = [grad]
+  error = 0
+  for j_t, j_n in grad:
+    if j_t.size or j_n.size:  # Handle zero size tensors correctly
+      error = np.maximum(error, np.fabs(j_t - j_n).max())
+  return error
+
+
+@tf_export(v1=["test.compute_gradient_error"])
+@deprecation.deprecated(
+    date=None,
+    instructions="Use tf.test.compute_gradient in 2.0, which has better "
+    "support for functions. Note that the two versions have different usage, "
+    "so code change is needed.")
+def compute_gradient_error(x,
+                           x_shape,
+                           y,
+                           y_shape,
+                           x_init_value=None,
+                           delta=1e-3,
+                           init_targets=None,
+                           extra_feed_dict=None):
+  """Computes the gradient error.
+
+  Computes the maximum error for dy/dx between the computed Jacobian and the
+  numerically estimated Jacobian.
+
+  This function will modify the tensors passed in as it adds more operations
+  and hence changing the consumers of the operations of the input tensors.
+
+  This function adds operations to the current session. To compute the error
+  using a particular device, such as a GPU, use the standard methods for
+  setting a device (e.g. using with sess.graph.device() or setting a device
+  function in the session constructor).
+
+  Args:
+    x: a tensor or list of tensors
+    x_shape: the dimensions of x as a tuple or an array of ints. If x is a list,
+    then this is the list of shapes.
+    y: a tensor
+    y_shape: the dimensions of y as a tuple or an array of ints.
+    x_init_value: (optional) a numpy array of the same shape as "x"
+      representing the initial value of x. If x is a list, this should be a list
+      of numpy arrays.  If this is none, the function will pick a random tensor
+      as the initial value.
+    delta: (optional) the amount of perturbation.
+    init_targets: list of targets to run to initialize model params.
+    extra_feed_dict: dict that allows fixing specified tensor values
+      during the Jacobian calculation.
+
+  Returns:
+    The maximum error in between the two Jacobians.
+  """
+  grad = compute_gradient(x, x_shape, y, y_shape, x_init_value, delta,
+                          init_targets, extra_feed_dict=extra_feed_dict)
+  return _compute_error(grad)

tensorflow/python/ops/gradient_checker_v2.py

@@ -0,0 +1,365 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Gradient checker for functions.
+
+The gradient checker verifies numerically that an function properly
+computes the gradients
+"""
+import numpy as np
+
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gradients_impl  # pylint: disable=unused-import
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util import numpy_compat
+from tensorflow.python.util.tf_export import tf_export
+
+
+def _product(t):
+  if isinstance(t, int):
+    return t
+  else:
+    y = 1
+    for x in t:
+      y *= x
+    return y
+
+
+def _eval_indexed_slices(a):
+  """Converts IndexedSlices to IndexedSlicesValue with numpy indices/values.
+
+  When eager execution is enabled, converts IndexedSlices
+  to IndexedSlicesValue with numpy indices/values.
+
+  Args:
+    a: any value.
+
+  Returns:
+    If a is IndexedSlices and eager execution is enabled, calls numpy() on a's
+    fields. Otherwise returns a unchanged.
+  """
+  if (isinstance(a, indexed_slices.IndexedSlices) and
+      context.executing_eagerly()):
+    return indexed_slices.IndexedSlicesValue(
+        indices=[x.numpy() for x in a.indices],
+        values=[x.numpy() for x in a.values],
+        dense_shape=a.dense_shape)
+  return a
+
+
+def _to_numpy(a):
+  """Converts Tensors, EagerTensors, and IndexedSlicesValue to numpy arrays.
+
+  Args:
+    a: any value.
+
+  Returns:
+    If a is EagerTensor or Tensor, returns the evaluation of a by calling
+    numpy() or run(). If a is IndexedSlicesValue, constructs the corresponding
+    dense numpy array. Otherwise returns a unchanged.
+  """
+  if isinstance(a, ops.EagerTensor):
+    return a.numpy()
+  if isinstance(a, tensor.Tensor):
+    sess = ops.get_default_session()
+    return sess.run(a)
+  if isinstance(a, indexed_slices.IndexedSlicesValue):
+    arr = np.zeros(a.dense_shape)
+    assert len(a.values) == len(a.indices), (
+        "IndexedSlicesValue has %s value slices but %s indices\n%s" %
+        (a.values, a.indices, a))
+    for values_slice, index in zip(a.values, a.indices):
+      assert 0 <= index < len(arr), (
+          "IndexedSlicesValue has invalid index %s\n%s" % (index, a))
+      arr[index] += values_slice
+    return arr
+  return a
+
+
+def _prepare(f, xs_dtypes, xs_shapes):
+  """Return a function that executes 'f'.
+
+    In TF 2.x, this is the same as `f`.
+    In TF 1.x, returns a Python function that executes the graph defined by `f`
+    in a Session.
+
+  Args:
+    f: the function.
+    xs_dtypes: dtypes of f's arguments.
+    xs_shapes: shapes of f's arguments.
+
+  Returns:
+  """
+  if context.executing_eagerly():
+
+    def decorated_eager(*xs_data):
+      return f(*map(ops.convert_to_tensor, xs_data))
+
+    return decorated_eager
+  xs = [
+      array_ops.placeholder(x_dtype, shape=x_shape)
+      for x_dtype, x_shape in zip(xs_dtypes, xs_shapes)
+  ]
+  y = f(*xs)
+  sess = ops.get_default_session()
+
+  def decorated_graph(*xs_data):
+    xs_data = [_to_numpy(a) for a in xs_data]
+    return sess.run(y, feed_dict=dict(zip(xs, xs_data)))
+
+  return decorated_graph
+
+
+def _compute_theoretical_jacobian(f, y_shape, y_dtype, xs, param):
+  """Computes the theoretical Jacobian for f regarding xs[param].
+
+  One can think of the relation among f, xs and y as y = f(xs).
+
+  Args:
+    f: the function.
+    y_shape: the shape of the result.
+    y_dtype: the dtype of the result.
+    xs: a list of tensors.
+    param: the index of the target parameter.
+
+  Returns:
+    A 2-d numpy array representing the Jacobian. It has "y_size" rows
+    and "x_size" columns where "x_size" is the number of elements in xs[param]
+    and "y_size" is the number of elements in the result.
+
+  Raises:
+    ValueError: If result is empty but the gradient is nonzero.
+  """
+  x = xs[param]
+  # Complex vectors are treated as vectors of twice as many reals.
+  x_shape = tuple(x.shape) + (2,) if x.dtype.is_complex else x.shape
+  y_factor = 2 if y_dtype.is_complex else 1
+
+  # To compute the jacobian, we treat x and y as one-dimensional vectors.
+  x_size = _product(x_shape)
+  x_val_size = _product(x_shape[1:])  # This is used for sparse gradients
+  y_size = _product(y_shape) * y_factor
+
+  # Allocate 2-D Jacobian, with y dimensions smashed into the first
+  # dimension and x dimensions smashed into the second.
+  jacobian = np.zeros((y_size, x_size), dtype=x.dtype.real_dtype.as_numpy_dtype)
+
+  # For each of the entry of dy, we set this to be 1 and
+  # everything else to be 0 and compute the gradients -- this will give us one
+  # row of the Jacobian matrix.
+  dy_data = np.zeros(y_shape, dtype=y_dtype.as_numpy_dtype)
+  dy_data_flat = dy_data.ravel().view(y_dtype.real_dtype.as_numpy_dtype)
+  grad_fn_unprep = backprop.gradients_function(f, [param])
+  grad_fn = _prepare(lambda dy, *xs: grad_fn_unprep(*xs, dy=dy),
+                     [y_dtype] + [z.dtype for z in xs],
+                     [None] + [z.shape for z in xs])
+  for row in range(y_size):
+    dy_data_flat[row] = 1
+    grad = _to_numpy(grad_fn(dy_data, *xs)[0])
+    grad = _eval_indexed_slices(grad)
+    if isinstance(grad, indexed_slices.IndexedSlicesValue):
+      for i, v in zip(grad.indices, grad.values):
+        c_begin = i * x_val_size
+        c_end = c_begin + x_val_size
+        jacobian[row, c_begin:c_end] += v.flat
+    elif grad is not None:
+      jacobian[row, :] = grad.ravel().view(jacobian.dtype)
+    # This reset of `dy_data_flat` needs to happen after `grad` is copied to
+    # `jacobian` because `grad` and `dy_data_flat` may share memory.
+    dy_data_flat[row] = 0
+
+  # If the output is empty, run the gradients at least once and make sure
+  # they produce zeros.
+  if y_size == 0:  # don't use 'not y_size', because y_size may not be an int
+    grad = _to_numpy(grad_fn(dy_data, *xs)[0])
+    if grad.shape != x.shape:
+      raise ValueError("Empty gradient has wrong shape: expected %s, got %s" %
+                       (x.shape, grad.shape))
+    if np.any(grad):
+      raise ValueError("Empty tensor with nonzero gradients")
+
+  logging.vlog(1, "Theoretical Jacobian =\n%s", jacobian)
+  return jacobian
+
+
+def _compute_numeric_jacobian(f, y_size, y_dtype, xs, param, delta):
+  """Computes the numeric Jacobian for f regarding xs[param].
+
+  One can think of the relation among f, xs and y as y = f(xs).
+
+  Args:
+    f: the function.
+    y_size: the number of elements of the result.
+    y_dtype: the dtype of the result.
+    xs: a list of tensors.
+    param: the index of the target parameter.
+    delta: the amount of perturbation we give to the input.
+
+  Returns:
+    A 2-d numpy array representing the Jacobian. It has "y_size" rows
+    and "x_size" columns where "x_size" is the number of elements in xs[param]
+    and "y_size" is the number of elements in the result.
+  """
+  x_shape = xs[param].shape
+  x_dtype = xs[param].dtype
+
+  # To compute the jacobian, we treat x and y as one-dimensional vectors
+  x_size = _product(x_shape) * (2 if x_dtype.is_complex else 1)
+  y_size = y_size * (2 if y_dtype.is_complex else 1)
+  x_dtype = x_dtype.real_dtype.as_numpy_dtype
+  y_dtype = y_dtype.real_dtype.as_numpy_dtype
+
+  xs_dtypes = [x.dtype for x in xs]
+  xs_shapes = [x.shape for x in xs]
+  # Converts xs to numpy arrays to do in-place perturbation.
+  # Calls asarray() to avoid copying in ravel() later.
+  xs = [numpy_compat.np_asarray(_to_numpy(x)) for x in xs]
+  x = xs[param]
+
+  # Make sure we have the right types
+  scale = numpy_compat.np_asarray(2 * delta, dtype=y_dtype)[()]
+
+  jacobian = np.zeros((y_size, x_size), dtype=x_dtype)
+
+  # For each of the entry of x, we slightly perturbs this by adding and
+  # subtracting a delta and then compute difference between the outputs. This
+  # will give us one column of the Jacobian matrix.
+  f = _prepare(f, xs_dtypes, xs_shapes)
+  for col in range(x_size):
+    original = x.ravel().view(x_dtype)[col]
+    x.ravel().view(x_dtype)[col] += delta
+    y_pos = _to_numpy(f(*xs))
+    x.ravel().view(x_dtype)[col] = original
+    x.ravel().view(x_dtype)[col] -= delta
+    y_neg = _to_numpy(f(*xs))
+    x.ravel().view(x_dtype)[col] = original
+    diff = (y_pos - y_neg) / scale
+    jacobian[:, col] = diff.ravel().view(y_dtype)
+
+  logging.vlog(1, "Numeric Jacobian =\n%s", jacobian)
+  return jacobian
+
+
+def _compute_gradient(f, y_shape, y_dtype, xs, param, delta):
+  """Computes the theoretical and numerical jacobian."""
+  x = xs[param]
+  t = x.dtype
+  allowed_types = [
+      dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64,
+      dtypes.complex64, dtypes.complex128
+  ]
+  assert t.base_dtype in allowed_types, ("Cannot compute gradient for "
+                                         "unsupported type %s of argument %s" %
+                                         (t.name, param))
+  t2 = y_dtype
+  assert t2.base_dtype in allowed_types, ("Cannot compute gradient for "
+                                          "unsupported type %s of y" % t2.name)
+  y_size = _product(y_shape)
+  jacob_t = _compute_theoretical_jacobian(f, y_shape, y_dtype, xs, param)
+  jacob_n = _compute_numeric_jacobian(f, y_size, y_dtype, xs, param, delta)
+  return jacob_t, jacob_n
+
+
+def _compute_gradient_list(f, xs, delta):
+  """Compute gradients for a list of x values."""
+  # convert xs to tensors so that dtype and shape have uniform types
+  xs = [ops.convert_to_tensor(x) for x in xs]
+  # run the function to get info of the result
+  xs_dtypes = [x.dtype for x in xs]
+  xs_shapes = [x.shape for x in xs]
+  f_temp = _prepare(f, xs_dtypes, xs_shapes)
+  y = f_temp(*xs)
+  return tuple(
+      zip(*[
+          _compute_gradient(f, y.shape, dtypes.as_dtype(y.dtype), xs, i, delta)
+          for i in range(len(xs))
+      ]))
+
+
+@tf_export("test.compute_gradient", v1=[])
+def compute_gradient(f, x, delta=None):
+  """Computes the theoretical and numeric Jacobian of `f`.
+
+  With y = f(x), computes the theoretical and numeric Jacobian dy/dx.
+
+  Args:
+    f: the function.
+    x: the arguments for the function as a list or tuple of values convertible
+      to a Tensor.
+    delta: (optional) perturbation used to compute numeric Jacobian.
+
+  Returns:
+    A pair of lists, where the first is a list of 2-d numpy arrays representing
+    the theoretical Jacobians for each argument, and the second list is the
+    numerical ones. Each 2-d array has "y_size" rows
+    and "x_size" columns where "x_size" is the number of elements in the
+    corresponding argument and "y_size" is the number of elements in f(x).
+
+  Raises:
+    ValueError: If result is empty but the gradient is nonzero.
+    ValueError: If x is not list, but any other type.
+
+  Example:
+
+  >>> @tf.function
+  ... def test_func(x):
+  ...   return x*x
+  ...
+  >>>
+  >>> class MyTest(tf.test.TestCase):
+  ...
+  ...   def test_gradient_of_test_func(self):
+  ...     theoretical, numerical = tf.test.compute_gradient(test_func, [1.0])
+  ...     # ((array([[2.]], dtype=float32),),
+  ...     #  (array([[2.000004]], dtype=float32),))
+  ...     self.assertAllClose(theoretical, numerical)
+
+  """
+  if not isinstance(x, (list, tuple)):
+    raise ValueError(
+        "`x` must be a list or tuple of values convertible to a Tensor "
+        "(arguments to `f`), not a %s" % type(x))
+  if delta is None:
+    # By default, we use a step size for the central finite difference
+    # approximation that is exactly representable as a binary floating
+    # point number, since this reduces the amount of noise due to rounding
+    # in the approximation of some functions.
+    delta = 1.0 / 1024
+  return _compute_gradient_list(f, x, delta)
+
+
+def max_error(grad1, grad2):
+  """Computes maximum elementwise gap.
+
+  Computes the maximum elementwise gap between two lists of tensors of the same
+  shape.
+
+  Args:
+    grad1: a lists of tensors.
+    grad2: a lists of tensors with the same shape as grad1.
+
+  Returns:
+    The maximum elementwise gap between the two.
+  """
+  error = 0
+  for j_t, j_n in zip(grad1, grad2):
+    if j_t.size or j_n.size:  # Handle zero size tensors correctly
+      error = np.maximum(error, np.fabs(j_t - j_n).max())
+  return error

tensorflow/python/ops/gradient_checker_v2_test.py

@@ -0,0 +1,378 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for compute_gradient."""
+
+import numpy as np
+
+from tensorflow.python.eager import backprop
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import custom_gradient
+from tensorflow.python.ops import \
+gradient_checker_v2 as gradient_checker
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import sparse_ops
+# needs this to register gradient for SoftmaxCrossEntropyWithLogits:
+import tensorflow.python.ops.nn_grad  # pylint: disable=unused-import
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging
+
+
+def _random_complex(shape, dtype):
+  data = np.random.random_sample(shape).astype(dtype.as_numpy_dtype)
+  if dtype.is_complex:
+    data.imag = np.random.random_sample(shape)
+  return data
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class GradientCheckerTest(test.TestCase):
+
+  def testSparseTensorReshape(self):
+    x = constant_op.constant(2.0, shape=(2,))
+
+    def sparse_tensor_reshape(values):
+      sparse = sparse_tensor.SparseTensor(
+          indices=[[0, 0], [1, 2]], values=values, dense_shape=[3, 4])
+      sparse = sparse_ops.sparse_reshape(sparse, shape=(12,))
+      return sparse.values
+
+    error = gradient_checker.max_error(
+        *gradient_checker.compute_gradient(sparse_tensor_reshape, [x]))
+
+    self.assertLess(error, 1e-4)
+
+  def testWithStaticShape(self):
+    size = (2, 3)
+    constant = constant_op.constant(2.0, shape=size, name="const")
+
+    def add_constant_with_static_shape_check(x):
+      self.assertAllEqual(x.shape.as_list(), constant.shape.as_list())
+      return x + constant
+
+    x = constant_op.constant(3.0, shape=size, name="x")
+
+    error = gradient_checker.max_error(*gradient_checker.compute_gradient(
+        add_constant_with_static_shape_check, [x]))
+
+    self.assertLess(error, 1e-4)
+
+  def testWithArgumentsAsTuple(self):
+    size = (2, 3)
+    x1 = constant_op.constant(2.0, shape=size, name="x1")
+    x2 = constant_op.constant(3.0, shape=size, name="x2")
+
+    error = gradient_checker.max_error(*gradient_checker.compute_gradient(
+        lambda x1: math_ops.add(x1, x2), (x1,)))
+
+    tf_logging.info("x1 error = %f", error)
+    self.assertLess(error, 1e-4)
+
+  def testAddSimple(self):
+    size = (2, 3)
+    x1 = constant_op.constant(2.0, shape=size, name="x1")
+    x2 = constant_op.constant(3.0, shape=size, name="x2")
+    error = gradient_checker.max_error(*gradient_checker.compute_gradient(
+        lambda x1: math_ops.add(x1, x2), [x1]))
+    tf_logging.info("x1 error = %f", error)
+    self.assertLess(error, 1e-4)
+
+  def testBfloat16(self):
+    x1 = constant_op.constant(2.0, dtype="bfloat16")
+    x2 = constant_op.constant(3.0, dtype="bfloat16")
+    # bfloat16 is very imprecise, so we use very large delta and error bar here.
+    error = gradient_checker.max_error(*gradient_checker.compute_gradient(
+        lambda x1: math_ops.add(x1, x2), [x1], delta=0.1))
+    tf_logging.info("x1 error = %f", error)
+    self.assertLess(error, 0.07)
+
+  def testAddCustomized(self):
+    size = (2, 3)
+    x1 = constant_op.constant(2.0, shape=size, dtype=dtypes.float64, name="x1")
+    x2 = np.asarray(np.arange(6, dtype=np.float64).reshape(2, 3))
+    # checkint gradients for x2 using a special delta
+    error = gradient_checker.max_error(*gradient_checker.compute_gradient(
+        lambda x2: math_ops.add(x1, x2), [x2], delta=1e-2))
+    tf_logging.info("x2 error = %f", error)
+    self.assertLess(error, 1e-10)
+
+  def testGather(self):
+
+    def f(params):
+      index_values = [1, 3]
+      indices = constant_op.constant(index_values, name="i")
+      return array_ops.gather(params, indices, name="y")
+
+    p_shape = (4, 2)
+    p_size = 8
+    params = constant_op.constant(
+        np.arange(p_size).astype(np.float64), shape=p_shape, name="p")
+    error = gradient_checker.max_error(
+        *gradient_checker.compute_gradient(f, [params]))
+    tf_logging.info("gather error = %f", error)
+    self.assertLess(error, 1e-4)
+
+  def testNestedGather(self):
+
+    def f(params):
+      index_values = [1, 3, 5, 6]
+      indices = constant_op.constant(index_values, name="i")
+      y = array_ops.gather(params, indices, name="y")
+      index_values2 = [0, 2]
+      indices2 = constant_op.constant(index_values2, name="i2")
+      return array_ops.gather(y, indices2, name="y2")
+
+    p_shape = (8, 2)
+    p_size = 16
+    params = constant_op.constant(
+        np.arange(p_size).astype(np.float64), shape=p_shape, name="p")
+    error = gradient_checker.max_error(
+        *gradient_checker.compute_gradient(f, [params]))
+    tf_logging.info("nested gather error = %f", error)
+    self.assertLess(error, 1e-4)
+
+  def testComplexMul(self):
+    c = constant_op.constant(5 + 7j, dtype=dtypes.complex64)
+
+    def f(x):
+      return c * x
+
+    x_shape = c.shape
+    x_dtype = c.dtype
+    x = constant_op.constant(_random_complex(x_shape, x_dtype))
+    analytical, numerical = gradient_checker.compute_gradient(f, [x])
+    correct = np.array([[5, -7], [7, 5]])
+    self.assertAllEqual(correct, analytical[0])
+    self.assertAllClose(correct, numerical[0], rtol=1e-4)
+    x = constant_op.constant(_random_complex(x_shape, x_dtype))
+    self.assertLess(
+        gradient_checker.max_error(*gradient_checker.compute_gradient(f, [x])),
+        3e-4)
+
+  def testComplexConj(self):
+
+    def f(x):
+      return math_ops.conj(x)
+
+    x_shape = ()
+    x_dtype = dtypes.complex64
+    x = constant_op.constant(_random_complex(x_shape, x_dtype))
+    analytical, numerical = gradient_checker.compute_gradient(f, [x])
+    correct = np.array([[1, 0], [0, -1]])
+    self.assertAllEqual(correct, analytical[0])
+    self.assertAllClose(correct, numerical[0], rtol=2e-5)
+    x = constant_op.constant(_random_complex(x_shape, x_dtype))
+    self.assertLess(
+        gradient_checker.max_error(*gradient_checker.compute_gradient(f, [x])),
+        2e-5)
+
+  def testEmptySucceeds(self):
+
+    def f(x):
+      return array_ops.identity(x)
+
+    x = constant_op.constant(
+        np.random.random_sample((0, 3)), dtype=dtypes.float32)
+    for grad in gradient_checker.compute_gradient(f, [x]):
+      self.assertEqual(grad[0].shape, (0, 0))
+    error = gradient_checker.max_error(
+        *gradient_checker.compute_gradient(f, [x]))
+    self.assertEqual(error, 0)
+
+  def testEmptyMatMul(self):
+
+    def f(x, y):
+      return math_ops.matmul(x, y)
+
+    x = constant_op.constant(
+        np.random.random_sample((0, 3)), dtype=dtypes.float32)
+    y = constant_op.constant(
+        np.random.random_sample((3, 4)), dtype=dtypes.float32)
+    for grad in gradient_checker.compute_gradient(f, [x, y]):
+      self.assertEqual(grad[0].shape, (0, 0))
+      self.assertEqual(grad[1].shape, (0, 12))
+    error = gradient_checker.max_error(
+        *gradient_checker.compute_gradient(f, [x, y]))
+    self.assertEqual(error, 0)
+
+  def testEmptyFails(self):
+
+    @custom_gradient.custom_gradient
+    def id_bad_grad(x):
+      y = array_ops.identity(x)
+
+      def grad_fn(dy):
+        # dx = constant_op.constant(np.zeros((1, 4)), dtype=dtypes.float32)
+        dx = array_ops.transpose(dy)
+        return dx
+
+      return y, grad_fn
+
+    def f(x):
+      return id_bad_grad(x)
+
+    x = constant_op.constant(
+        np.random.random_sample((0, 3)), dtype=dtypes.float32)
+    bad = r"Empty gradient has wrong shape: expected \(0, 3\), got \(3, 0\)"
+    with self.assertRaisesRegex(ValueError, bad):
+      gradient_checker.compute_gradient(f, [x])
+
+  def testNaNGradFails(self):
+
+    @custom_gradient.custom_gradient
+    def id_nan_grad(x):
+      y = array_ops.identity(x)
+
+      def grad_fn(dy):
+        dx = np.nan * dy
+        # dx = dy
+        return dx
+
+      return y, grad_fn
+
+    def f(x):
+      return id_nan_grad(x)
+
+    x = constant_op.constant(
+        np.random.random_sample((1, 1)), dtype=dtypes.float32)
+    error = gradient_checker.max_error(
+        *gradient_checker.compute_gradient(f, [x]))
+    # Typical test would assert error < max_err, so assert this test would
+    # raise AssertionError, since NaN is not < 1.0.
+    error_msg = r"(nan|np.float32\(nan\)) not less than 1.0"
+    with self.assertRaisesRegex(AssertionError, error_msg):
+      self.assertLess(error, 1.0)
+
+  def testGradGrad(self):
+
+    def f(x):
+      with backprop.GradientTape() as tape:
+        tape.watch(x)
+        y = math_ops.square(x)
+        z = math_ops.square(y)
+      return tape.gradient(z, x)
+
+    analytical, numerical = gradient_checker.compute_gradient(f, [2.0])
+    self.assertAllEqual([[[48.]]], analytical)
+    self.assertAllClose([[[48.]]], numerical, rtol=1e-4)
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class MiniMNISTTest(test.TestCase):
+
+  # Gradient checker for MNIST.
+  def _BuildAndTestMiniMNIST(self, param_index, tag):
+    # Fix seed to avoid occasional flakiness
+    np.random.seed(6)
+
+    # Hyperparameters
+    batch = 3
+    inputs = 16
+    features = 32
+    classes = 10
+
+    # Define the parameters
+    inp_data = np.random.random_sample(inputs * batch)
+    hidden_weight_data = np.random.randn(inputs * features) / np.sqrt(inputs)
+    hidden_bias_data = np.random.random_sample(features)
+    sm_weight_data = np.random.randn(features * classes) / np.sqrt(features)
+    sm_bias_data = np.random.random_sample(classes)
+
+    # special care for labels since they need to be normalized per batch
+    label_data = np.random.random(batch * classes).reshape((batch, classes))
+    s = label_data.sum(axis=1)
+    label_data /= s[:, None]
+
+    # We treat the inputs as "parameters" here
+    inp = constant_op.constant(
+        inp_data.tolist(),
+        shape=[batch, inputs],
+        dtype=dtypes.float64,
+        name="inp")
+    hidden_weight = constant_op.constant(
+        hidden_weight_data.tolist(),
+        shape=[inputs, features],
+        dtype=dtypes.float64,
+        name="hidden_weight")
+    hidden_bias = constant_op.constant(
+        hidden_bias_data.tolist(),
+        shape=[features],
+        dtype=dtypes.float64,
+        name="hidden_bias")
+    softmax_weight = constant_op.constant(
+        sm_weight_data.tolist(),
+        shape=[features, classes],
+        dtype=dtypes.float64,
+        name="softmax_weight")
+    softmax_bias = constant_op.constant(
+        sm_bias_data.tolist(),
+        shape=[classes],
+        dtype=dtypes.float64,
+        name="softmax_bias")
+
+    # List all the parameter so that we can test them one at a time
+    all_params = [inp, hidden_weight, hidden_bias, softmax_weight, softmax_bias]
+
+    # Now, Building MNIST
+    def f(inp, hidden_weight, hidden_bias, softmax_weight, softmax_bias):
+      features = nn_ops.relu(
+          nn_ops.xw_plus_b(inp, hidden_weight, hidden_bias), name="features")
+      logits = nn_ops.xw_plus_b(
+          features, softmax_weight, softmax_bias, name="logits")
+      labels = constant_op.constant(
+          label_data.tolist(),
+          shape=[batch, classes],
+          dtype=dtypes.float64,
+          name="labels")
+      cost = nn_ops.softmax_cross_entropy_with_logits(
+          labels=labels, logits=logits, name="cost")
+      return cost
+
+    def f_restricted(x):
+      xs = all_params
+      i = param_index
+      # use x for the i-th parameter
+      xs = xs[0:i] + [x] + xs[i + 1:]
+      return f(*xs)
+
+    # Test the gradients.
+    err = gradient_checker.max_error(*gradient_checker.compute_gradient(
+        f_restricted, [all_params[param_index]], delta=1e-5))
+
+    tf_logging.info("Mini MNIST: %s gradient error = %g", tag, err)
+    return err
+
+  def testInputGradient(self):
+    self.assertLess(self._BuildAndTestMiniMNIST(0, "input"), 1e-8)
+
+  def testHiddenWeightGradient(self):
+    self.assertLess(self._BuildAndTestMiniMNIST(1, "hidden_weight"), 1e-8)
+
+  def testHiddenBiasGradient(self):
+    self.assertLess(self._BuildAndTestMiniMNIST(2, "hidden_bias"), 1e-8)
+
+  def testSoftmaxWeightGradient(self):
+    self.assertLess(self._BuildAndTestMiniMNIST(3, "softmax_weight"), 1e-8)
+
+  def testSoftmaxBiasGradient(self):
+    self.assertLess(self._BuildAndTestMiniMNIST(4, "softmax_bias"), 1e-8)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/ops/gradients.py

@@ -0,0 +1,26 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Implements the graph generation for computation of gradients."""
+
+# pylint: disable=unused-import
+from tensorflow.python.eager import function
+from tensorflow.python.eager.backprop import GradientTape
+from tensorflow.python.eager.forwardprop import ForwardAccumulator
+from tensorflow.python.ops.custom_gradient import custom_gradient
+from tensorflow.python.ops.gradients_util import AggregationMethod
+from tensorflow.python.ops.gradients_impl import gradients
+from tensorflow.python.ops.gradients_impl import hessians
+from tensorflow.python.ops.unconnected_gradients import UnconnectedGradients
+# pylint: enable=unused-import

tensorflow/python/ops/gradients_impl.py

@@ -0,0 +1,491 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Implements the graph generation for computation of gradients."""
+
+from tensorflow.compiler.jit.ops import xla_ops_grad  # pylint: disable=unused-import
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_grad  # pylint: disable=unused-import
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops  # pylint: disable=unused-import
+from tensorflow.python.ops import control_flow_grad  # pylint: disable=unused-import
+from tensorflow.python.ops import cudnn_rnn_grad  # pylint: disable=unused-import
+from tensorflow.python.ops import gradients_util
+from tensorflow.python.ops import image_grad  # pylint: disable=unused-import
+from tensorflow.python.ops import io_ops  # pylint: disable=unused-import
+from tensorflow.python.ops import linalg_grad  # pylint: disable=unused-import
+from tensorflow.python.ops import linalg_ops  # pylint: disable=unused-import
+from tensorflow.python.ops import logging_ops  # pylint: disable=unused-import
+from tensorflow.python.ops import lookup_grad  # pylint: disable=unused-import
+from tensorflow.python.ops import manip_grad  # pylint: disable=unused-import
+from tensorflow.python.ops import math_grad  # pylint: disable=unused-import
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nccl_ops  # pylint: disable=unused-import
+from tensorflow.python.ops import nn_grad  # pylint: disable=unused-import
+from tensorflow.python.ops import optional_grad  # pylint: disable=unused-import
+from tensorflow.python.ops import parsing_grad  # pylint: disable=unused-import
+from tensorflow.python.ops import proto_ops  # pylint: disable=unused-import
+from tensorflow.python.ops import random_grad  # pylint: disable=unused-import
+from tensorflow.python.ops import rnn_grad  # pylint: disable=unused-import
+from tensorflow.python.ops import sdca_ops  # pylint: disable=unused-import
+from tensorflow.python.ops import sets  # pylint: disable=unused-import
+from tensorflow.python.ops import sparse_grad  # pylint: disable=unused-import
+from tensorflow.python.ops import tensor_array_grad  # pylint: disable=unused-import
+from tensorflow.python.ops import tensor_array_ops
+from tensorflow.python.ops import while_loop
+from tensorflow.python.ops.linalg.sparse import sparse_csr_matrix_grad  # pylint: disable=unused-import
+from tensorflow.python.ops.signal import fft_ops  # pylint: disable=unused-import
+from tensorflow.python.ops.unconnected_gradients import UnconnectedGradients
+from tensorflow.python.training import checkpoint_ops  # pylint: disable=unused-import
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["gradients"])
+def gradients(ys,
+              xs,
+              grad_ys=None,
+              name="gradients",
+              colocate_gradients_with_ops=False,
+              gate_gradients=False,
+              aggregation_method=None,
+              stop_gradients=None,
+              unconnected_gradients=UnconnectedGradients.NONE):
+  """Constructs symbolic derivatives of sum of `ys` w.r.t. x in `xs`.
+
+  `ys` and `xs` are each a `Tensor` or a list of tensors.  `grad_ys`
+  is a list of `Tensor`, holding the gradients received by the
+  `ys`. The list must be the same length as `ys`.
+
+  `gradients()` adds ops to the graph to output the derivatives of `ys` with
+  respect to `xs`.  It returns a list of `Tensor` of length `len(xs)` where
+  each tensor is the `sum(dy/dx)` for y in `ys` and for x in `xs`.
+
+  `grad_ys` is a list of tensors of the same length as `ys` that holds
+  the initial gradients for each y in `ys`.  When `grad_ys` is None,
+  we fill in a tensor of '1's of the shape of y for each y in `ys`.  A
+  user can provide their own initial `grad_ys` to compute the
+  derivatives using a different initial gradient for each y (e.g., if
+  one wanted to weight the gradient differently for each value in
+  each y).
+
+  `stop_gradients` is a `Tensor` or a list of tensors to be considered constant
+  with respect to all `xs`. These tensors will not be backpropagated through,
+  as though they had been explicitly disconnected using `stop_gradient`.  Among
+  other things, this allows computation of partial derivatives as opposed to
+  total derivatives. For example:
+
+  ```python
+  a = tf.constant(0.)
+  b = 2 * a
+  g = tf.gradients(a + b, [a, b], stop_gradients=[a, b])
+  ```
+
+  Here the partial derivatives `g` evaluate to `[1.0, 1.0]`, compared to the
+  total derivatives `tf.gradients(a + b, [a, b])`, which take into account the
+  influence of `a` on `b` and evaluate to `[3.0, 1.0]`.  Note that the above is
+  equivalent to:
+
+  ```python
+  a = tf.stop_gradient(tf.constant(0.))
+  b = tf.stop_gradient(2 * a)
+  g = tf.gradients(a + b, [a, b])
+  ```
+
+  `stop_gradients` provides a way of stopping gradient after the graph has
+  already been constructed, as compared to `tf.stop_gradient` which is used
+  during graph construction.  When the two approaches are combined,
+  backpropagation stops at both `tf.stop_gradient` nodes and nodes in
+  `stop_gradients`, whichever is encountered first.
+
+  All integer tensors are considered constant with respect to all `xs`, as if
+  they were included in `stop_gradients`.
+
+  `unconnected_gradients` determines the value returned for each x in xs if it
+  is unconnected in the graph to ys. By default this is None to safeguard
+  against errors. Mathematically these gradients are zero which can be requested
+  using the `'zero'` option. `tf.UnconnectedGradients` provides the
+  following options and behaviors:
+
+  ```python
+  a = tf.ones([1, 2])
+  b = tf.ones([3, 1])
+  g1 = tf.gradients([b], [a], unconnected_gradients='none')
+  sess.run(g1)  # [None]
+
+  g2 = tf.gradients([b], [a], unconnected_gradients='zero')
+  sess.run(g2)  # [array([[0., 0.]], dtype=float32)]
+  ```
+
+  Let us take one practical example which comes during the back propogation
+  phase. This function is used to evaluate the derivatives of the cost function
+  with respect to Weights `Ws` and Biases `bs`. Below sample implementation
+  provides the exaplantion of what it is actually used for :
+
+  ```python
+  Ws = tf.constant(0.)
+  bs = 2 * Ws
+  cost = Ws + bs  # This is just an example. So, please ignore the formulas.
+  g = tf.gradients(cost, [Ws, bs])
+  dCost_dW, dCost_db = g
+  ```
+
+
+  Args:
+    ys: A `Tensor` or list of tensors to be differentiated.
+    xs: A `Tensor` or list of tensors to be used for differentiation.
+    grad_ys: Optional. A `Tensor` or list of tensors the same size as
+      `ys` and holding the gradients computed for each y in `ys`.
+    name: Optional name to use for grouping all the gradient ops together.
+      defaults to 'gradients'.
+    colocate_gradients_with_ops: If True, try colocating gradients with
+      the corresponding op.
+    gate_gradients: If True, add a tuple around the gradients returned
+      for an operations.  This avoids some race conditions.
+    aggregation_method: Specifies the method used to combine gradient terms.
+      Accepted values are constants defined in the class `AggregationMethod`.
+    stop_gradients: Optional. A `Tensor` or list of tensors not to differentiate
+      through.
+    unconnected_gradients: Optional. Specifies the gradient value returned when
+      the given input tensors are unconnected. Accepted values are constants
+      defined in the class `tf.UnconnectedGradients` and the default value is
+      `none`.
+
+  Returns:
+    A list of `Tensor` of length `len(xs)` where each tensor is the `sum(dy/dx)`
+    for y in `ys` and for x in `xs`.
+
+  Raises:
+    LookupError: if one of the operations between `x` and `y` does not
+      have a registered gradient function.
+    ValueError: if the arguments are invalid.
+    RuntimeError: if called in Eager mode.
+
+  """
+  # Creating the gradient graph for control flow mutates Operations.
+  # _mutation_lock ensures a Session.run call cannot occur between creating and
+  # mutating new ops.
+  # pylint: disable=protected-access
+  with ops.get_default_graph()._mutation_lock():
+    return gradients_util._GradientsHelper(
+        ys, xs, grad_ys, name, colocate_gradients_with_ops,
+        gate_gradients, aggregation_method, stop_gradients,
+        unconnected_gradients)
+  # pylint: enable=protected-access
+
+
+@tf_export("gradients", v1=[])
+def gradients_v2(ys,  # pylint: disable=invalid-name
+                 xs,
+                 grad_ys=None,
+                 name="gradients",
+                 gate_gradients=False,
+                 aggregation_method=None,
+                 stop_gradients=None,
+                 unconnected_gradients=UnconnectedGradients.NONE):
+  """Constructs symbolic derivatives of sum of `ys` w.r.t. x in `xs`.
+
+  `tf.gradients` is only valid in a graph context. In particular,
+  it is valid in the context of a `tf.function` wrapper, where code
+  is executing as a graph.
+
+  `ys` and `xs` are each a `Tensor` or a list of tensors.  `grad_ys`
+  is a list of `Tensor`, holding the gradients received by the
+  `ys`. The list must be the same length as `ys`.
+
+  `gradients()` adds ops to the graph to output the derivatives of `ys` with
+  respect to `xs`.  It returns a list of `Tensor` of length `len(xs)` where
+  each tensor is the `sum(dy/dx)` for y in `ys` and for x in `xs`.
+
+  `grad_ys` is a list of tensors of the same length as `ys` that holds
+  the initial gradients for each y in `ys`.  When `grad_ys` is None,
+  we fill in a tensor of '1's of the shape of y for each y in `ys`.  A
+  user can provide their own initial `grad_ys` to compute the
+  derivatives using a different initial gradient for each y (e.g., if
+  one wanted to weight the gradient differently for each value in
+  each y).
+
+  `stop_gradients` is a `Tensor` or a list of tensors to be considered constant
+  with respect to all `xs`. These tensors will not be backpropagated through,
+  as though they had been explicitly disconnected using `stop_gradient`.  Among
+  other things, this allows computation of partial derivatives as opposed to
+  total derivatives. For example:
+
+  >>> @tf.function
+  ... def example():
+  ...   a = tf.constant(0.)
+  ...   b = 2 * a
+  ...   return tf.gradients(a + b, [a, b], stop_gradients=[a, b])
+  >>> example()
+  [<tf.Tensor: shape=(), dtype=float32, numpy=1.0>,
+  <tf.Tensor: shape=(), dtype=float32, numpy=1.0>]
+
+  Here the partial derivatives `g` evaluate to `[1.0, 1.0]`, compared to the
+  total derivatives `tf.gradients(a + b, [a, b])`, which take into account the
+  influence of `a` on `b` and evaluate to `[3.0, 1.0]`.  Note that the above is
+  equivalent to:
+
+  >>> @tf.function
+  ... def example():
+  ...   a = tf.stop_gradient(tf.constant(0.))
+  ...   b = tf.stop_gradient(2 * a)
+  ...   return tf.gradients(a + b, [a, b])
+  >>> example()
+  [<tf.Tensor: shape=(), dtype=float32, numpy=1.0>,
+  <tf.Tensor: shape=(), dtype=float32, numpy=1.0>]
+
+  `stop_gradients` provides a way of stopping gradient after the graph has
+  already been constructed, as compared to `tf.stop_gradient` which is used
+  during graph construction.  When the two approaches are combined,
+  backpropagation stops at both `tf.stop_gradient` nodes and nodes in
+  `stop_gradients`, whichever is encountered first.
+
+  All integer tensors are considered constant with respect to all `xs`, as if
+  they were included in `stop_gradients`.
+
+  `unconnected_gradients` determines the value returned for each x in xs if it
+  is unconnected in the graph to ys. By default this is None to safeguard
+  against errors. Mathematically these gradients are zero which can be requested
+  using the `'zero'` option. `tf.UnconnectedGradients` provides the
+  following options and behaviors:
+
+  >>> @tf.function
+  ... def example(use_zero):
+  ...   a = tf.ones([1, 2])
+  ...   b = tf.ones([3, 1])
+  ...   if use_zero:
+  ...     return tf.gradients([b], [a], unconnected_gradients='zero')
+  ...   else:
+  ...     return tf.gradients([b], [a], unconnected_gradients='none')
+  >>> example(False)
+  [None]
+  >>> example(True)
+  [<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[0., 0.]], ...)>]
+
+  Let us take one practical example which comes during the back propogation
+  phase. This function is used to evaluate the derivatives of the cost function
+  with respect to Weights `Ws` and Biases `bs`. Below sample implementation
+  provides the exaplantion of what it is actually used for :
+
+  >>> @tf.function
+  ... def example():
+  ...   Ws = tf.constant(0.)
+  ...   bs = 2 * Ws
+  ...   cost = Ws + bs  # This is just an example. Please ignore the formulas.
+  ...   g = tf.gradients(cost, [Ws, bs])
+  ...   dCost_dW, dCost_db = g
+  ...   return dCost_dW, dCost_db
+  >>> example()
+  (<tf.Tensor: shape=(), dtype=float32, numpy=3.0>,
+  <tf.Tensor: shape=(), dtype=float32, numpy=1.0>)
+
+  Args:
+    ys: A `Tensor` or list of tensors to be differentiated.
+    xs: A `Tensor` or list of tensors to be used for differentiation.
+    grad_ys: Optional. A `Tensor` or list of tensors the same size as
+      `ys` and holding the gradients computed for each y in `ys`.
+    name: Optional name to use for grouping all the gradient ops together.
+      defaults to 'gradients'.
+    gate_gradients: If True, add a tuple around the gradients returned
+      for an operations.  This avoids some race conditions.
+    aggregation_method: Specifies the method used to combine gradient terms.
+      Accepted values are constants defined in the class `AggregationMethod`.
+    stop_gradients: Optional. A `Tensor` or list of tensors not to differentiate
+      through.
+    unconnected_gradients: Optional. Specifies the gradient value returned when
+      the given input tensors are unconnected. Accepted values are constants
+      defined in the class `tf.UnconnectedGradients` and the default value is
+      `none`.
+
+  Returns:
+    A list of `Tensor` of length `len(xs)` where each tensor is the `sum(dy/dx)`
+    for y in `ys` and for x in `xs`.
+
+  Raises:
+    LookupError: if one of the operations between `x` and `y` does not
+      have a registered gradient function.
+    ValueError: if the arguments are invalid.
+    RuntimeError: if called in Eager mode.
+
+  """
+  # Creating the gradient graph for control flow mutates Operations.
+  # _mutation_lock ensures a Session.run call cannot occur between creating and
+  # mutating new ops.
+  # pylint: disable=protected-access
+  with ops.get_default_graph()._mutation_lock():
+    return gradients_util._GradientsHelper(
+        ys, xs, grad_ys, name, True, gate_gradients,
+        aggregation_method, stop_gradients,
+        unconnected_gradients)
+  # pylint: enable=protected-access
+
+
+# TODO(vrv): Make this available when we want to make it public.
+def _hessian_vector_product(ys, xs, v):
+  """Multiply the Hessian of `ys` wrt `xs` by `v`.
+
+  This is an efficient construction that uses a backprop-like approach
+  to compute the product between the Hessian and another vector. The
+  Hessian is usually too large to be explicitly computed or even
+  represented, but this method allows us to at least multiply by it
+  for the same big-O cost as backprop.
+
+  Implicit Hessian-vector products are the main practical, scalable way
+  of using second derivatives with neural networks. They allow us to
+  do things like construct Krylov subspaces and approximate conjugate
+  gradient descent.
+
+  Example: if `y` = 1/2 `x`^T A `x`, then `hessian_vector_product(y,
+  x, v)` will return an expression that evaluates to the same values
+  as (A + A.T) `v`.
+
+  Args:
+    ys: A scalar value, or a tensor or list of tensors to be summed to
+        yield a scalar.
+    xs: A list of tensors that we should construct the Hessian over.
+    v: A list of tensors, with the same shapes as xs, that we want to
+       multiply by the Hessian.
+
+  Returns:
+    A list of tensors (or if the list would be length 1, a single tensor)
+    containing the product between the Hessian and `v`.
+
+  Raises:
+    ValueError: `xs` and `v` have different length.
+
+  """
+
+  # Validate the input
+  length = len(xs)
+  if len(v) != length:
+    raise ValueError("xs and v must have the same length.")
+
+  # First backprop
+  grads = gradients(ys, xs)
+
+  assert len(grads) == length
+  elemwise_products = [
+      math_ops.multiply(grad_elem, array_ops.stop_gradient(v_elem))
+      for grad_elem, v_elem in zip(grads, v)
+      if grad_elem is not None
+  ]
+
+  # Second backprop
+  return gradients(elemwise_products, xs)
+
+
+@tf_export(v1=["hessians"])
+def hessians(ys,
+             xs,
+             name="hessians",
+             colocate_gradients_with_ops=False,
+             gate_gradients=False,
+             aggregation_method=None):
+  """Constructs the Hessian of sum of `ys` with respect to `x` in `xs`.
+
+  `hessians()` adds ops to the graph to output the Hessian matrix of `ys`
+  with respect to `xs`.  It returns a list of `Tensor` of length `len(xs)`
+  where each tensor is the Hessian of `sum(ys)`.
+
+  The Hessian is a matrix of second-order partial derivatives of a scalar
+  tensor (see https://en.wikipedia.org/wiki/Hessian_matrix for more details).
+
+  Args:
+    ys: A `Tensor` or list of tensors to be differentiated.
+    xs: A `Tensor` or list of tensors to be used for differentiation.
+    name: Optional name to use for grouping all the gradient ops together.
+      defaults to 'hessians'.
+    colocate_gradients_with_ops: See `gradients()` documentation for details.
+    gate_gradients: See `gradients()` documentation for details.
+    aggregation_method: See `gradients()` documentation for details.
+
+  Returns:
+    A list of Hessian matrices of `sum(ys)` for each `x` in `xs`.
+
+  Raises:
+    LookupError: if one of the operations between `xs` and `ys` does not
+      have a registered gradient function.
+  """
+  xs = gradients_util._AsList(xs)  # pylint: disable=protected-access
+  kwargs = {
+      "colocate_gradients_with_ops": colocate_gradients_with_ops,
+      "gate_gradients": gate_gradients,
+      "aggregation_method": aggregation_method
+  }
+  # Compute first-order derivatives and iterate for each x in xs.
+  hessians = []
+  _gradients = gradients(ys, xs, **kwargs)
+  for gradient, x in zip(_gradients, xs):
+    # change shape to one-dimension without graph branching
+    gradient = array_ops.reshape(gradient, [-1])
+
+    # Declare an iterator and tensor array loop variables for the gradients.
+    n = array_ops.size(x)
+    loop_vars = [
+        array_ops.constant(0, dtypes.int32),
+        tensor_array_ops.TensorArray(x.dtype, n)
+    ]
+    # Iterate over all elements of the gradient and compute second order
+    # derivatives.
+    _, hessian = while_loop.while_loop(
+        lambda j, _: j < n,
+        lambda j, result: (j + 1,
+                           result.write(j, gradients(gradient[j], x)[0])),
+        loop_vars
+    )
+
+    _shape = array_ops.shape(x)
+    _reshaped_hessian = array_ops.reshape(hessian.stack(),
+                                          array_ops.concat((_shape, _shape), 0))
+    hessians.append(_reshaped_hessian)
+  return hessians
+
+
+@tf_export("hessians", v1=[])
+def HessiansV2(ys,
+               xs,
+               gate_gradients=False,
+               aggregation_method=None,
+               name="hessians"):
+  """Constructs the Hessian of sum of `ys` with respect to `x` in `xs`.
+
+  `hessians()` adds ops to the graph to output the Hessian matrix of `ys`
+  with respect to `xs`.  It returns a list of `Tensor` of length `len(xs)`
+  where each tensor is the Hessian of `sum(ys)`.
+
+  The Hessian is a matrix of second-order partial derivatives of a scalar
+  tensor (see https://en.wikipedia.org/wiki/Hessian_matrix for more details).
+
+  Args:
+    ys: A `Tensor` or list of tensors to be differentiated.
+    xs: A `Tensor` or list of tensors to be used for differentiation.
+    gate_gradients: See `gradients()` documentation for details.
+    aggregation_method: See `gradients()` documentation for details.
+    name: Optional name to use for grouping all the gradient ops together.
+      defaults to 'hessians'.
+
+  Returns:
+    A list of Hessian matrices of `sum(ys)` for each `x` in `xs`.
+
+  Raises:
+    LookupError: if one of the operations between `xs` and `ys` does not
+      have a registered gradient function.
+  """
+  return hessians(
+      ys,
+      xs,
+      name=name,
+      colocate_gradients_with_ops=True,
+      gate_gradients=gate_gradients,
+      aggregation_method=aggregation_method)

tensorflow/python/ops/handle_data_util.py

@@ -0,0 +1,91 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Decorator to overrides the gradient for a function."""
+
+from tensorflow.python.client import pywrap_tf_session
+from tensorflow.python.framework import cpp_shape_inference_pb2
+from tensorflow.python.framework import dtypes
+from tensorflow.python.types import core
+from tensorflow.python.util import compat
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("__internal__.ops.get_resource_handle_data", v1=[])
+def get_resource_handle_data(graph_op):
+  assert (isinstance(graph_op, core.Symbol)
+          and not isinstance(graph_op, core.Value))
+
+  with graph_op.graph._c_graph.get() as c_graph:  # pylint: disable=protected-access
+    handle_data = pywrap_tf_session.GetHandleShapeAndType(
+        c_graph, graph_op._as_tf_output())  # pylint: disable=protected-access
+
+  return cpp_shape_inference_pb2.CppShapeInferenceResult.HandleData.FromString(
+      compat.as_bytes(handle_data))
+
+
+def get_handle_data(source_t):
+  """Obtains HandleData from a tensor."""
+  if isinstance(source_t, core.Value):
+    return source_t._handle_data  # pylint: disable=protected-access
+  return get_resource_handle_data(source_t)
+
+
+def copy_handle_data(source_t, target_t):
+  """Copies HandleData for variant and resource type tensors if available.
+
+  The CppShapeInferenceResult::HandleData proto contains information about the
+  shapes and types of the element tensors of resource/variant type tensors.
+  We need to copy this across function boundaries, i.e., when capturing a
+  placeholder or when returning a function tensor as output. If we don't do this
+  the element tensors will have unknown shapes, e.g., if a TensorList variant
+  tensor is captured as a placeholder, elements popped from that list would have
+  unknown shape.
+
+  Args:
+    source_t: The tensor to copy HandleData from.
+    target_t: The tensor to copy HandleData to.
+  """
+  if (target_t.dtype == dtypes.resource or
+      target_t.dtype == dtypes.variant):
+    handle_data = get_handle_data(source_t)
+    set_handle_data(target_t, handle_data)
+
+
+def set_handle_data(target_t, handle_data):
+  """Sets handle data on the giver tensor."""
+  if (
+      handle_data is None
+      or not handle_data.is_set
+      or not handle_data.shape_and_type
+  ):
+    return
+
+  # pylint: disable=protected-access
+  if isinstance(target_t, core.Value):
+    target_t._handle_data = handle_data
+    return
+  with target_t.graph._c_graph.get() as c_graph:
+    pywrap_tf_session.SetHandleShapeAndType(c_graph, target_t._as_tf_output(),
+                                            handle_data.SerializeToString())
+  # pylint: enable=protected-access
+
+
+def create_handle_data(shape, dtype):
+  handle_data = cpp_shape_inference_pb2.CppShapeInferenceResult.HandleData()
+  handle_data.is_set = True
+  handle_data.shape_and_type.append(
+      cpp_shape_inference_pb2.CppShapeInferenceResult.HandleShapeAndType(
+          shape=shape.as_proto(), dtype=dtype.as_datatype_enum))
+  return handle_data

tensorflow/python/ops/histogram_ops.py

@@ -0,0 +1,149 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# pylint: disable=g-short-docstring-punctuation
+"""Histograms.
+"""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import clip_ops
+from tensorflow.python.ops import control_flow_assert
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import gen_math_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.util import dispatch
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export('histogram_fixed_width_bins')
+@dispatch.add_dispatch_support
+def histogram_fixed_width_bins(values,
+                               value_range,
+                               nbins=100,
+                               dtype=dtypes.int32,
+                               name=None):
+  """Bins the given values for use in a histogram.
+
+  Given the tensor `values`, this operation returns a rank 1 `Tensor`
+  representing the indices of a histogram into which each element
+  of `values` would be binned. The bins are equal width and
+  determined by the arguments `value_range` and `nbins`.
+
+  Args:
+    values:  Numeric `Tensor`.
+    value_range:  Shape [2] `Tensor` of same `dtype` as `values`.
+      values <= value_range[0] will be mapped to hist[0],
+      values >= value_range[1] will be mapped to hist[-1].
+    nbins:  Scalar `int32 Tensor`.  Number of histogram bins.
+    dtype:  dtype for returned histogram.
+    name:  A name for this operation (defaults to 'histogram_fixed_width').
+
+  Returns:
+    A `Tensor` holding the indices of the binned values whose shape matches
+    `values`.
+
+  Raises:
+    TypeError: If any unsupported dtype is provided.
+    tf.errors.InvalidArgumentError: If value_range does not
+        satisfy value_range[0] < value_range[1].
+
+  Examples:
+
+  >>> # Bins will be:  (-inf, 1), [1, 2), [2, 3), [3, 4), [4, inf)
+  ...
+  >>> nbins = 5
+  >>> value_range = [0.0, 5.0]
+  >>> new_values = [-1.0, 0.0, 1.5, 2.0, 5.0, 15]
+  >>> indices = tf.histogram_fixed_width_bins(new_values, value_range, nbins=5)
+  >>> indices.numpy()
+  array([0, 0, 1, 2, 4, 4], dtype=int32)
+  """
+  with ops.name_scope(name, 'histogram_fixed_width_bins',
+                      [values, value_range, nbins]):
+    values = ops.convert_to_tensor(values, name='values')
+    shape = array_ops.shape(values)
+
+    values = array_ops.reshape(values, [-1])
+    value_range = ops.convert_to_tensor(value_range, name='value_range')
+    nbins = ops.convert_to_tensor(nbins, dtype=dtypes.int32, name='nbins')
+    check = control_flow_assert.Assert(
+        math_ops.greater(nbins, 0), ['nbins %s must > 0' % nbins])
+    nbins = control_flow_ops.with_dependencies([check], nbins)
+    nbins_float = math_ops.cast(nbins, values.dtype)
+
+    # Map tensor values that fall within value_range to [0, 1].
+    scaled_values = math_ops.truediv(
+        values - value_range[0],
+        value_range[1] - value_range[0],
+        name='scaled_values')
+
+    # map tensor values within the open interval value_range to {0,.., nbins-1},
+    # values outside the open interval will be zero or less, or nbins or more.
+    indices = math_ops.floor(nbins_float * scaled_values, name='indices')
+
+    # Clip edge cases (e.g. value = value_range[1]) or "outliers."
+    indices = math_ops.cast(
+        clip_ops.clip_by_value(indices, 0, nbins_float - 1), dtypes.int32)
+    return array_ops.reshape(indices, shape)
+
+
+@tf_export('histogram_fixed_width')
+@dispatch.add_dispatch_support
+def histogram_fixed_width(values,
+                          value_range,
+                          nbins=100,
+                          dtype=dtypes.int32,
+                          name=None):
+  """Return histogram of values.
+
+  Given the tensor `values`, this operation returns a rank 1 histogram counting
+  the number of entries in `values` that fell into every bin.  The bins are
+  equal width and determined by the arguments `value_range` and `nbins`.
+
+  Args:
+    values:  Numeric `Tensor`.
+    value_range:  Shape [2] `Tensor` of same `dtype` as `values`.
+      values <= value_range[0] will be mapped to hist[0],
+      values >= value_range[1] will be mapped to hist[-1].
+    nbins:  Scalar `int32 Tensor`.  Number of histogram bins.
+    dtype:  dtype for returned histogram.
+    name:  A name for this operation (defaults to 'histogram_fixed_width').
+
+  Returns:
+    A 1-D `Tensor` holding histogram of values.
+
+  Raises:
+    TypeError: If any unsupported dtype is provided.
+    tf.errors.InvalidArgumentError: If value_range does not
+        satisfy value_range[0] < value_range[1].
+
+  Examples:
+
+  >>> # Bins will be:  (-inf, 1), [1, 2), [2, 3), [3, 4), [4, inf)
+  ...
+  >>> nbins = 5
+  >>> value_range = [0.0, 5.0]
+  >>> new_values = [-1.0, 0.0, 1.5, 2.0, 5.0, 15]
+  >>> hist = tf.histogram_fixed_width(new_values, value_range, nbins=5)
+  >>> hist.numpy()
+  array([2, 1, 1, 0, 2], dtype=int32)
+  """
+  with ops.name_scope(name, 'histogram_fixed_width',
+                      [values, value_range, nbins]) as name:
+    # pylint: disable=protected-access
+    return gen_math_ops._histogram_fixed_width(
+        values, value_range, nbins, dtype=dtype, name=name)
+    # pylint: enable=protected-access

tensorflow/python/ops/image_grad.py

@@ -0,0 +1,381 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Contains Gradient functions for image ops."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import gen_image_ops
+from tensorflow.python.ops import math_ops
+
+
+@ops.RegisterGradient("ResizeNearestNeighbor")
+def _ResizeNearestNeighborGrad(op: ops.Operation, grad):
+  """The derivatives for nearest neighbor resizing.
+
+  Args:
+    op: The ResizeNearestNeighbor op.
+    grad: The tensor representing the gradient w.r.t. the output.
+
+  Returns:
+    The gradients w.r.t. the input and the output.
+  """
+  image = op.inputs[0]
+  if image.get_shape()[1:3].is_fully_defined():
+    image_shape = image.get_shape()[1:3]
+  else:
+    image_shape = array_ops.shape(image)[1:3]
+
+  grads = gen_image_ops.resize_nearest_neighbor_grad(
+      grad,
+      image_shape,
+      align_corners=op.get_attr("align_corners"),
+      half_pixel_centers=op.get_attr("half_pixel_centers"))
+  return [grads, None]
+
+
+@ops.RegisterGradient("ResizeBilinear")
+def _ResizeBilinearGrad(op: ops.Operation, grad):
+  """The derivatives for bilinear resizing.
+
+  Args:
+    op: The ResizeBilinear op.
+    grad: The tensor representing the gradient w.r.t. the output.
+
+  Returns:
+    The gradients w.r.t. the input.
+  """
+  grad0 = gen_image_ops.resize_bilinear_grad(
+      grad,
+      op.inputs[0],
+      align_corners=op.get_attr("align_corners"),
+      half_pixel_centers=op.get_attr("half_pixel_centers"))
+  return [grad0, None]
+
+
+ops.NotDifferentiable("ResizeBilinearGrad")
+
+
+@ops.RegisterGradient("ScaleAndTranslate")
+def _ScaleAndTranslateGrad(op, grad):
+  """The derivatives for ScaleAndTranslate transformation op.
+
+  Args:
+    op: The ScaleAndTranslate op.
+    grad: The tensor representing the gradient w.r.t. the output.
+
+  Returns:
+    The gradients w.r.t. the input.
+  """
+
+  grad0 = gen_image_ops.scale_and_translate_grad(
+      grad,
+      op.inputs[0],
+      op.inputs[2],
+      op.inputs[3],
+      kernel_type=op.get_attr("kernel_type"),
+      antialias=op.get_attr("antialias"))
+  return [grad0, None, None, None]
+
+
+@ops.RegisterGradient("ResizeBicubic")
+def _ResizeBicubicGrad(op: ops.Operation, grad):
+  """The derivatives for bicubic resizing.
+
+  Args:
+    op: The ResizeBicubic op.
+    grad: The tensor representing the gradient w.r.t. the output.
+
+  Returns:
+    The gradients w.r.t. the input.
+  """
+  allowed_types = [dtypes.float32, dtypes.float64]
+  grad0 = None
+  if op.inputs[0].dtype in allowed_types:
+    grad0 = gen_image_ops.resize_bicubic_grad(
+        grad,
+        op.inputs[0],
+        align_corners=op.get_attr("align_corners"),
+        half_pixel_centers=op.get_attr("half_pixel_centers"))
+  return [grad0, None]
+
+
+@ops.RegisterGradient("CropAndResize")
+def _CropAndResizeGrad(op: ops.Operation, grad):
+  """The derivatives for crop_and_resize.
+
+  We back-propagate to the image only when the input image tensor has floating
+  point dtype but we always back-propagate to the input boxes tensor.
+
+  Args:
+    op: The CropAndResize op.
+    grad: The tensor representing the gradient w.r.t. the output.
+
+  Returns:
+    The gradients w.r.t. the input image, boxes, as well as the always-None
+    gradients w.r.t. box_ind and crop_size.
+  """
+  image = op.inputs[0]
+  if image.get_shape().is_fully_defined():
+    image_shape = image.get_shape().as_list()
+  else:
+    image_shape = array_ops.shape(image)
+
+  allowed_types = [dtypes.float16, dtypes.float32, dtypes.float64]
+  if op.inputs[0].dtype in allowed_types:
+    # pylint: disable=protected-access
+    grad0 = gen_image_ops.crop_and_resize_grad_image(
+        grad, op.inputs[1], op.inputs[2], image_shape, T=op.get_attr("T"),
+        method=op.get_attr("method"))
+    # pylint: enable=protected-access
+  else:
+    grad0 = None
+
+  # `grad0` is the gradient to the input image pixels and it
+  # has been implemented for nearest neighbor and bilinear sampling
+  # respectively. `grad1` is the gradient to the input crop boxes' coordinates.
+  # When using nearest neighbor sampling, the gradient to crop boxes'
+  # coordinates are not well defined. In practice, we still approximate
+  # grad1 using the gradient derived from bilinear sampling.
+  grad1 = gen_image_ops.crop_and_resize_grad_boxes(
+      grad, op.inputs[0], op.inputs[1], op.inputs[2])
+
+  return [grad0, grad1, None, None]
+
+
+def _CustomReciprocal(x):
+  """Wrapper function around `math_ops.div_no_nan()` to perform a "safe" reciprocal incase the input is zero. Avoids divide by zero and NaNs.
+
+  Input:
+    x -> input tensor to be reciprocat-ed.
+  Returns:
+    x_reciprocal -> reciprocal of x without NaNs.
+  """
+  return math_ops.div_no_nan(math_ops.cast(1.0, x.dtype), x)
+
+
+@ops.RegisterGradient("RGBToHSV")
+def _RGBToHSVGrad(op: ops.Operation, grad):
+  """The gradients for `rgb_to_hsv` operation.
+
+  This function is a piecewise continuous function as defined here:
+  https://en.wikipedia.org/wiki/HSL_and_HSV#From_RGB
+  We perform the multivariate derivative and compute all partial derivatives
+  separately before adding them in the end. Formulas are given before each
+  partial derivative calculation.
+
+  Args:
+    op: The `rgb_to_hsv` `Operation` that we are differentiating.
+    grad: Gradient with respect to the output of the `rgb_to_hsv` op.
+
+  Returns:
+    Gradients with respect to the input of `rgb_to_hsv`.
+  """
+  # Input Channels
+  reds = op.inputs[0][..., 0]
+  greens = op.inputs[0][..., 1]
+  blues = op.inputs[0][..., 2]
+  # Output Channels
+  saturation = op.outputs[0][..., 1]
+  value = op.outputs[0][..., 2]
+
+  dtype = op.inputs[0].dtype
+
+  # Mask/Indicator for max and min values of each pixel.
+  # Arbitrary assignment in case of tie breakers with R>G>B.
+  # Max values
+  red_biggest = math_ops.cast((reds >= blues) & \
+                 (reds >= greens), dtype)
+  green_biggest = math_ops.cast((greens > reds) & \
+                   (greens >= blues), dtype)
+  blue_biggest = math_ops.cast((blues > reds) & \
+                  (blues > greens), dtype)
+  # Min values
+  red_smallest = math_ops.cast((reds < blues) & \
+                  (reds < greens), dtype)
+  green_smallest = math_ops.cast((greens <= reds) & \
+                    (greens < blues), dtype)
+  blue_smallest = math_ops.cast((blues <= reds) & \
+                   (blues <= greens), dtype)
+
+  # Derivatives of R, G, B wrt Value slice
+  dv_dr = red_biggest
+  dv_dg = green_biggest
+  dv_db = blue_biggest
+
+  # Derivatives of R, G, B wrt Saturation slice
+
+  # The first term in the addition is the case when the corresponding color
+  # from (r,g,b) was "MAX"
+  # -> derivative = MIN/square(MAX), MIN could be one of the other two colors
+  # The second term is the case when the corresponding color from
+  # (r,g,b) was "MIN"
+  # -> derivative = -1/MAX, MAX could be one of the other two colours.
+  ds_dr = math_ops.cast(reds > 0, dtype) * math_ops.add(
+      red_biggest * math_ops.add(green_smallest * greens, blue_smallest * blues)
+      * _CustomReciprocal(math_ops.square(reds)), red_smallest * -1 *
+      _CustomReciprocal((green_biggest * greens) + (blue_biggest * blues)))
+  ds_dg = math_ops.cast(greens > 0, dtype) * math_ops.add(
+      green_biggest * math_ops.add(red_smallest * reds, blue_smallest * blues) *
+      _CustomReciprocal(math_ops.square(greens)), green_smallest * -1 *
+      _CustomReciprocal((red_biggest * reds) + (blue_biggest * blues)))
+  ds_db = math_ops.cast(blues > 0, dtype) * math_ops.add(
+      blue_biggest * math_ops.add(green_smallest * greens, red_smallest * reds)
+      * _CustomReciprocal(math_ops.square(blues)), blue_smallest * -1 *
+      _CustomReciprocal((green_biggest * greens) + (red_biggest * reds)))
+
+  # Derivatives of R, G, B wrt Hue slice
+
+  # Need to go case by case for each color.
+  # for red, dh_dr -> dh_dr_1 + dh_dr_2 + dh_dr_3 + dh_dr_4 + dh_dr_5
+  # dh_dr_1 ->
+  # if red was MAX, then derivative = 60 * -1 * (G-B)/square(MAX-MIN) == 60 *\
+  #  -1 * (greens-blues) * reciprocal(square(saturation)) * \
+  #  reciprocal(square(value))
+  # elif green was MAX, there are two subcases
+  # ie when red was MIN and when red was NOT MIN
+  #   dh_dr_2 ->
+  #   if red was MIN (use UV rule) ->  60 * ((1 * -1/(MAX-MIN)) +\
+  #     (B-R)*(-1/square(MAX-MIN) * -1)) == 60 * (blues - greens) *\
+  #     reciprocal(square(reds - greens))
+  #   dh_dr_3 ->
+  #   if red was NOT MIN -> 60 * -1/MAX-MIN == -60 * reciprocal(greens-blues)
+  # elif blue was MAX, there are two subcases
+  #   dh_dr_4 ->
+  #   if red was MIN (similarly use the UV rule) -> 60 * (blues - greens) *\
+  #     reciprocal(square(blues - reds))
+  #   dh_dr_5 ->
+  #   if red was NOT MIN -> 60 * 1/MAX-MIN == 60 * reciprocal(blues-greens)
+  dh_dr_1 = 60 * (
+      math_ops.cast(reds > 0, dtype) * red_biggest * -1 *
+      (greens - blues) * _CustomReciprocal(math_ops.square(saturation)) *
+      _CustomReciprocal(math_ops.square(value)))
+  dh_dr_2 = 60 * (
+      math_ops.cast(greens > 0, dtype) * green_biggest * red_smallest *
+      (blues - greens) * _CustomReciprocal(math_ops.square(reds - greens)))
+  dh_dr_3 = 60 * (
+      math_ops.cast(greens > 0, dtype) * green_biggest * blue_smallest * -1 *
+      _CustomReciprocal(greens - blues))
+  dh_dr_4 = 60 * (
+      math_ops.cast(blues > 0, dtype) * blue_biggest * red_smallest *
+      (blues - greens) * _CustomReciprocal(math_ops.square(blues - reds)))
+  dh_dr_5 = 60 * (
+      math_ops.cast(blues > 0, dtype) * blue_biggest * green_smallest *
+      _CustomReciprocal(blues - greens))
+
+  dh_dr = dh_dr_1 + dh_dr_2 + dh_dr_3 + dh_dr_4 + dh_dr_5
+  # Converting from degrees to [0,1] scale as specified in
+  # https://www.tensorflow.org/api_docs/python/tf/image/rgb_to_hsv
+  dh_dr = dh_dr / 360
+
+  # for green, dh_dg -> dh_dg_1 + dh_dg_2 + dh_dg_3 + dh_dg_4 + dh_dg_5
+  # dh_dg_1 ->
+  # if green was MAX, then derivative = 60 * -1 * (B-R)/square(MAX-MIN) == 60 *\
+  #   -1 * (blues - reds) * reciprocal(square(saturation)) * \
+  #   reciprocal(square(value))
+  # elif red was MAX, there are two subcases ie
+  # when green was MIN and when green was NOT MIN
+  #   dh_dg_2 ->
+  #   if green was MIN (use UV rule) ->  60 * ((1 * 1/(MAX-MIN)) + \
+  #     (greens-blues) * (-1/square(MAX-MIN) * -1)) == 60 * \
+  #     ((reciprocal(reds-greens) + (greens-blues) * \
+  #     reciprocal(square(reds-greens))))
+  #   dh_dg_3 ->
+  #   if green was NOT MIN -> 60 * 1/MAX-MIN == 60 * reciprocal(reds - blues)
+  # elif blue was MAX, there are two subcases
+  #   dh_dg_4 ->
+  #   if green was MIN (similarly use the UV rule) -> 60 * -1 * \
+  #     (reciprocal(blues - greens) + (reds-greens)* -1 * \
+  #     reciprocal(square(blues-greens)))
+  #   dh_dr_5 ->
+  #   if green was NOT MIN -> 60 * -1/MAX-MIN == -60 * reciprocal(blues - reds)
+  dh_dg_1 = 60 * (
+      math_ops.cast(greens > 0, dtype) * green_biggest * -1 *
+      (blues - reds) * _CustomReciprocal(math_ops.square(saturation)) *
+      _CustomReciprocal(math_ops.square(value)))
+  dh_dg_2 = 60 * (
+      math_ops.cast(reds > 0, dtype) * red_biggest * green_smallest *
+      (reds - blues) * _CustomReciprocal(math_ops.square(reds - greens)))
+  dh_dg_3 = 60 * (
+      math_ops.cast(reds > 0, dtype) * red_biggest * blue_smallest *
+      _CustomReciprocal(reds - blues))
+  dh_dg_4 = 60 * (
+      math_ops.cast(blues > 0, dtype) * blue_biggest * green_smallest *
+      (reds - blues) * _CustomReciprocal(math_ops.square(blues - greens)))
+  dh_dg_5 = 60 * (
+      math_ops.cast(blues > 0, dtype) * blue_biggest * red_smallest * -1 *
+      _CustomReciprocal(blues - reds))
+
+  dh_dg = dh_dg_1 + dh_dg_2 + dh_dg_3 + dh_dg_4 + dh_dg_5
+  # Converting from degrees to [0,1] scale as specified in
+  # https://www.tensorflow.org/api_docs/python/tf/image/rgb_to_hsv
+  dh_dg = dh_dg / 360
+
+  # for blue, dh_db -> dh_db_1 + dh_db_2 + dh_db_3 + dh_db_4 + dh_db_5
+  # dh_db_1 ->
+  # if blue was MAX, then derivative = 60 * -1 * (R-G)/square(MAX-MIN) == 60 *\
+  #   -1 * reciprocal(square(saturation)) * reciprocal(square(value))
+  # elif red was MAX, there are two subcases
+  # ie when blue was MIN and when blue was NOT MIN
+  #   dh_dg_2 ->
+  #   if blue was MIN (use UV rule) ->  60 * ((1 * -1/(MAX-MIN)) + \
+  #     (greens-blues) * (-1/square(MAX-MIN) * -1)) == 60 * (greens - reds) *\
+  #     reciprocal(square(reds - blues))
+  #   dh_dg_3 ->
+  #   if blue was NOT MIN -> 60 * -1/MAX-MIN == 60 * -1 * \
+  #     reciprocal(reds - greens)
+  # elif green was MAX, there are two subcases
+  #   dh_dg_4 ->
+  #   if blue was MIN (similarly use the UV rule) -> 60 * -1 * \
+  #     (reciprocal(greens - blues) + (blues - reds) * -1 * \
+  #     reciprocal(square(greens - blues)))
+  #   dh_dr_5 ->
+  #   if blue was NOT MIN -> 60 * 1/MAX-MIN == 60 * reciprocal(greens - reds)
+  dh_db_1 = 60 * (
+      math_ops.cast(blues > 0, dtype) * blue_biggest * -1 *
+      (reds - greens) * _CustomReciprocal(math_ops.square(saturation)) *
+      _CustomReciprocal(math_ops.square(value)))
+  dh_db_2 = 60 * (
+      math_ops.cast(reds > 0, dtype) * red_biggest * blue_smallest *
+      (greens - reds) * _CustomReciprocal(math_ops.square(reds - blues)))
+  dh_db_3 = 60 * (
+      math_ops.cast(reds > 0, dtype) * red_biggest * green_smallest * -1 *
+      _CustomReciprocal(reds - greens))
+  dh_db_4 = 60 * (
+      math_ops.cast(greens > 0, dtype) * green_biggest * blue_smallest *
+      (greens - reds) * _CustomReciprocal(math_ops.square(greens - blues)))
+  dh_db_5 = 60 * (
+      math_ops.cast(greens > 0, dtype) * green_biggest * red_smallest *
+      _CustomReciprocal(greens - reds))
+
+  dh_db = dh_db_1 + dh_db_2 + dh_db_3 + dh_db_4 + dh_db_5
+  # Converting from degrees to [0,1] scale as specified in
+  # https://www.tensorflow.org/api_docs/python/tf/image/rgb_to_hsv
+  dh_db = dh_db / 360
+
+  # Gradients wrt to inputs
+  dv_drgb = array_ops_stack.stack(
+      [grad[..., 2] * dv_dr, grad[..., 2] * dv_dg, grad[..., 2] * dv_db],
+      axis=-1)
+  ds_drgb = array_ops_stack.stack(
+      [grad[..., 1] * ds_dr, grad[..., 1] * ds_dg, grad[..., 1] * ds_db],
+      axis=-1)
+  dh_drgb = array_ops_stack.stack(
+      [grad[..., 0] * dh_dr, grad[..., 0] * dh_dg, grad[..., 0] * dh_db],
+      axis=-1)
+
+  gradient_input = math_ops.add(math_ops.add(dv_drgb, ds_drgb), dh_drgb)
+  return gradient_input

tensorflow/python/ops/image_grad_d9m_test.py

@@ -0,0 +1,374 @@
+# Copyright 2020-2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Functional tests for deterministic image op gradient functions."""
+
+import numpy as np
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import context
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import random_seed
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gradients_impl
+from tensorflow.python.ops import image_grad_test_base as test_base
+from tensorflow.python.ops import image_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.platform import test
+
+
+class ResizeNearestNeighborOpDeterminismExceptionsTest(test.TestCase,
+                                                       parameterized.TestCase):
+  """Test d9m-unimplemented exceptions from ResizeNearestNeighborOpGrad.
+
+  Test that tf.errors.UnimplementedError is thrown, as appropriate, by the
+  GPU-specific code-path through ResizeNearestNeighborOpGrad when deterministic
+  ops are enabled.
+
+  This test assumes that image_grad_test.py runs equivalent test cases when
+  deterministic ops are not enabled and will therefore detect erroneous
+  exception throwing in those cases.
+  """
+
+  @parameterized.parameters(
+      {
+          'align_corners': False,
+          'half_pixel_centers': False,
+          'data_type': dtypes.float16
+      }, {
+          'align_corners': False,
+          'half_pixel_centers': False,
+          'data_type': dtypes.float32
+      }, {
+          'align_corners': False,
+          'half_pixel_centers': False,
+          'data_type': dtypes.float64
+      }, {
+          'align_corners': True,
+          'half_pixel_centers': False,
+          'data_type': dtypes.float32
+      }, {
+          'align_corners': False,
+          'half_pixel_centers': True,
+          'data_type': dtypes.float32
+      })
+  @test_util.run_gpu_only
+  @test_util.run_all_in_graph_and_eager_modes
+  def testExceptionThrowing(self, align_corners, half_pixel_centers, data_type):
+    with self.session(), test_util.force_gpu():
+      input_image = array_ops.zeros((1, 2, 2, 1), dtype=data_type)
+      with backprop.GradientTape() as tape:
+        tape.watch(input_image)
+        output_image = image_ops.resize_nearest_neighbor(
+            input_image, (3, 3),
+            align_corners=align_corners,
+            half_pixel_centers=half_pixel_centers)
+      with self.assertRaisesRegex(
+          errors.UnimplementedError,
+          'A deterministic GPU implementation of ResizeNearestNeighborGrad' +
+          ' is not currently available.'):
+        gradient = tape.gradient(output_image, input_image)
+        self.evaluate(gradient)
+
+
+class ResizeBilinearOpDeterministicTest(test_base.ResizeBilinearOpTestBase):
+  """Test that ResizeBilinearGrad operates reproducibly.
+
+  Inheriting from test_base.ResizeBilinearOpTestBase ensures that regular op
+  functionality is correct when the deterministic code-path is selected.
+  """
+
+  def _randomNDArray(self, shape):
+    return 2 * np.random.random_sample(shape) - 1
+
+  def _randomDataOp(self, shape, data_type):
+    return constant_op.constant(self._randomNDArray(shape), dtype=data_type)
+
+  @parameterized.parameters(
+      # Note that there is no 16-bit floating point format registered for GPU
+      {
+          'align_corners': False,
+          'half_pixel_centers': False,
+          'data_type': dtypes.float32
+      },
+      {
+          'align_corners': False,
+          'half_pixel_centers': False,
+          'data_type': dtypes.float64
+      },
+      {
+          'align_corners': True,
+          'half_pixel_centers': False,
+          'data_type': dtypes.float32
+      },
+      {
+          'align_corners': False,
+          'half_pixel_centers': True,
+          'data_type': dtypes.float32
+      })
+  @test_util.run_in_graph_and_eager_modes
+  @test_util.run_gpu_only
+  def testDeterministicGradients(self, align_corners, half_pixel_centers,
+                                 data_type):
+    if not align_corners and test_util.is_xla_enabled():
+      # Align corners is deprecated in TF2.0, but align_corners==False is not
+      # supported by XLA.
+      self.skipTest('align_corners==False not currently supported by XLA')
+    with self.session(force_gpu=True):
+      seed = (
+          hash(align_corners) % 256 + hash(half_pixel_centers) % 256 +
+          hash(data_type) % 256)
+      np.random.seed(seed)
+      input_shape = (1, 25, 12, 3)  # NHWC
+      output_shape = (1, 200, 250, 3)
+      input_image = self._randomDataOp(input_shape, data_type)
+      repeat_count = 3
+      if context.executing_eagerly():
+
+        def resize_bilinear_gradients(local_seed):
+          np.random.seed(local_seed)
+          upstream_gradients = self._randomDataOp(output_shape, dtypes.float32)
+          with backprop.GradientTape(persistent=True) as tape:
+            tape.watch(input_image)
+            output_image = image_ops.resize_bilinear(
+                input_image,
+                output_shape[1:3],
+                align_corners=align_corners,
+                half_pixel_centers=half_pixel_centers)
+            gradient_injector_output = output_image * upstream_gradients
+          return tape.gradient(gradient_injector_output, input_image)
+
+        for i in range(repeat_count):
+          local_seed = seed + i  # select different upstream gradients
+          result_a = resize_bilinear_gradients(local_seed)
+          result_b = resize_bilinear_gradients(local_seed)
+          self.assertAllEqual(result_a, result_b)
+      else:  # graph mode
+        upstream_gradients = array_ops.placeholder(
+            dtypes.float32, shape=output_shape, name='upstream_gradients')
+        output_image = image_ops.resize_bilinear(
+            input_image,
+            output_shape[1:3],
+            align_corners=align_corners,
+            half_pixel_centers=half_pixel_centers)
+        gradient_injector_output = output_image * upstream_gradients
+        # The gradient function behaves as if grad_ys is multiplied by the op
+        # gradient result, not passing the upstream gradients through the op's
+        # gradient generation graph. This is the reason for using the
+        # gradient injector
+        resize_bilinear_gradients = gradients_impl.gradients(
+            gradient_injector_output,
+            input_image,
+            grad_ys=None,
+            colocate_gradients_with_ops=True)[0]
+        for i in range(repeat_count):
+          feed_dict = {upstream_gradients: self._randomNDArray(output_shape)}
+          result_a = resize_bilinear_gradients.eval(feed_dict=feed_dict)
+          result_b = resize_bilinear_gradients.eval(feed_dict=feed_dict)
+          self.assertAllEqual(result_a, result_b)
+
+
+class CropAndResizeOpDeterminismExceptionsTest(test.TestCase):
+  """Test d9m-unimplemented exceptions from CropAndResizeBackprop{Image|Boxes}.
+
+  Test that tf.errors.UnimplementedError is thrown or not thrown, as
+  appropriate, by the GPU code-paths for CropAndResizeBackprop{Image|Boxes} when
+  deterministic ops are enabled.
+
+  This test assumes that test_base.CropAndResizeOpTestBase runs all the same
+  test cases when deterministic ops are not enabled and will therefore detect
+  erroneous exception throwing in those cases.
+  """
+
+  def _genParams(self, dtype=dtypes.float32):
+    batch_size = 1
+    image_height = 10
+    image_width = 10
+    channels = 1
+    image_shape = (batch_size, image_height, image_width, channels)
+    num_boxes = 3
+    boxes_shape = (num_boxes, 4)
+    random_seed.set_seed(123)
+    image = random_ops.random_normal(shape=image_shape, dtype=dtype)
+    boxes = random_ops.random_uniform(shape=boxes_shape, dtype=dtypes.float32)
+    box_indices = random_ops.random_uniform(
+        shape=(num_boxes,), minval=0, maxval=batch_size, dtype=dtypes.int32)
+    crop_size = constant_op.constant([3, 3], dtype=dtypes.int32)
+    return image, boxes, box_indices, crop_size
+
+  @test_util.run_in_graph_and_eager_modes
+  @test_util.run_gpu_only
+  def testExceptionThrowing(self):
+    for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]:
+      image, boxes, box_indices, crop_size = self._genParams(dtype)
+      with backprop.GradientTape(persistent=True) as tape:
+        tape.watch(image)
+        tape.watch(boxes)
+        op_output = image_ops.crop_and_resize_v2(image, boxes, box_indices,
+                                                 crop_size)
+      image_error_message = ('Deterministic GPU implementation of' +
+                             ' CropAndResizeBackpropImage not available')
+      with self.assertRaisesRegex(errors_impl.UnimplementedError,
+                                  image_error_message):
+        result = tape.gradient(op_output, image)
+        self.evaluate(result)
+      expected_error_message = ('Deterministic GPU implementation of' +
+                                ' CropAndResizeBackpropBoxes not available')
+      if context.executing_eagerly():
+        # With eager execution, the backprop-to-image code is apparently
+        # executed (first), even when its output is never used.
+        expected_error_message = image_error_message
+      with self.assertRaisesRegex(errors_impl.UnimplementedError,
+                                  expected_error_message):
+        result = tape.gradient(op_output, boxes)
+        self.evaluate(result)
+
+
+class CropAndResizeOpDeterministicTest(test_base.CropAndResizeOpTestBase):
+  """Test that CropAndResizeBackprop{Image|Boxes} operates reproducibly.
+
+  Inheriting from test_base.CropAndResizeOpTestBase ensures that regular op
+  functionality is correct when the deterministic code-path is selected.
+  """
+
+  def _randomFloats(self, shape, low=0.0, high=1.0, dtype=dtypes.float32):
+    """Generate a tensor of random floating-point values.
+
+    Values will be continuously distributed in the range [low, high).
+
+    Note that we use numpy to generate random numbers and then feed the result
+    through a constant op to avoid the re-rolling of TensorFlow random ops on
+    each run in graph mode.
+
+    Args:
+      shape: The output shape.
+      low: Lower bound of random numbers generated, inclusive.
+      high: Upper bound of random numbers generated, exclusive.
+      dtype: The output dtype.
+
+    Returns:
+      A random tensor
+    """
+    val = np.random.random_sample(
+        shape)  # float64 continuous uniform [0.0, 1.0)
+    diff = high - low
+    val *= diff
+    val += low
+    return constant_op.constant(val, dtype=dtype)
+
+  def _randomInts(self, shape, low, high):
+    """Generate a tensor of random 32-bit integer values.
+
+    Note that we use numpy to generate random numbers and then feed the result
+    through a constant op to avoid the re-rolling of TensorFlow random ops on
+    each run in graph mode.
+
+    Args:
+      shape: The output shape.
+      low: Lower bound of random numbers generated, inclusive.
+      high: Upper bound of random numbers generated, exclusive.
+
+    Returns:
+      A random tensor
+    """
+    val = np.random.randint(low=low, high=high, size=shape)
+    return constant_op.constant(val, dtype=dtypes.int32)
+
+  def _genParams(self, dtype=dtypes.float32):
+    batch_size = 16
+    input_height = 64
+    input_width = 64
+    depth = 1
+    input_shape = (batch_size, input_height, input_width, depth)
+    np.random.seed(456)
+    image = self._randomFloats(input_shape, low=-1.0, high=1.0, dtype=dtype)
+    box_count = 4 * batch_size
+    boxes = self._randomFloats((box_count, 4),
+                               low=0.0,
+                               high=1.01,
+                               dtype=dtypes.float32)
+    box_indices = self._randomInts((box_count,), low=0, high=batch_size)
+    crop_size = [input_height * 2, input_width * 2]
+    output_shape = (box_count, *crop_size, depth)
+    # The output of this op is always float32, regardless of image data type
+    injected_gradients = self._randomFloats(
+        output_shape, low=-0.001, high=0.001, dtype=dtypes.float32)
+    return image, boxes, box_indices, crop_size, injected_gradients
+
+  def _testReproducibleBackprop(self, test_image_not_boxes):
+    with test_util.force_cpu():
+      for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]:
+        params = self._genParams(dtype)
+        image, boxes, box_indices, crop_size, injected_gradients = params
+
+        with backprop.GradientTape(persistent=True) as tape:
+          tape.watch([image, boxes])
+          output = image_ops.crop_and_resize_v2(
+              image, boxes, box_indices, crop_size, method='bilinear')
+          upstream = output * injected_gradients
+
+        image_gradients_a, boxes_gradients_a = tape.gradient(
+            upstream, [image, boxes])
+        for _ in range(5):
+          image_gradients_b, boxes_gradients_b = tape.gradient(
+              upstream, [image, boxes])
+          if test_image_not_boxes:
+            self.assertAllEqual(image_gradients_a, image_gradients_b)
+          else:
+            self.assertAllEqual(boxes_gradients_a, boxes_gradients_b)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testReproducibleBackpropToImage(self):
+    """Test that backprop to image is reproducible.
+
+    With non-reproducible ordering of reduction operations, upsampling of a
+    crop, leading to three or more output pixels being derived from an input
+    pixel, can contribute to nondeterminism in the gradient associated with that
+    input pixel location.
+
+    Note that the number of boxes can be less than, equal to, or greater than
+    the batch size. With non-reproducible ordering of reduction operations,
+    three
+    or more crops overlapping on the same input image pixel can independently
+    contribute to nondeterminism in the image gradient associated with that
+    input pixel location. This is independent of contributions caused by the
+    upsampling of any given crop.
+    """
+    self._testReproducibleBackprop(test_image_not_boxes=True)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testReproducibleBackpropToBoxes(self):
+    """Test that backprop to boxes is reproducible.
+
+    If the input and output dimensions are the same, then the boxes gradients
+    will be deterministically zero. Otherwise, in the presence of
+    non-reproducible ordering of reduction operations, nondeterminism can be
+    introduced, whether there is upsampling or downsampling and whether or not
+    there are overlapping crops.
+    """
+    self._testReproducibleBackprop(test_image_not_boxes=False)
+
+
+if __name__ == '__main__':
+  # TODO(reedwm): Merge this file with image_grad_test.py and
+  # image_grad_test_base.py
+  config.enable_op_determinism()
+  test.main()