Add custom_gradient function.

PiperOrigin-RevId: 188765271

5 files changed, 0 insertions, 351 deletions

diff --git a/tensorflow/contrib/bayesflow/BUILD b/tensorflow/contrib/bayesflow/BUILD
index 88956f0512..c6feec68e0 100644
--- a/tensorflow/contrib/bayesflow/BUILD
+++ b/tensorflow/contrib/bayesflow/BUILD
@@ -57,26 +57,6 @@ cuda_py_test(
 )
 
 cuda_py_test(
-    name = "custom_grad_test",
-    size = "small",
-    srcs = ["python/kernel_tests/custom_grad_test.py"],
-    additional_deps = [
-        ":bayesflow_py",
-        "//third_party/py/numpy",
-        "//tensorflow/contrib/layers:layers_py",
-        "//tensorflow/python:array_ops",
-        "//tensorflow/python:client_testlib",
-        "//tensorflow/python:framework_for_generated_wrappers",
-        "//tensorflow/python:framework_test_lib",
-        "//tensorflow/python:gradients",
-        "//tensorflow/python:init_ops",
-        "//tensorflow/python:platform_test",
-        "//tensorflow/python:variable_scope",
-        "//tensorflow/python:variables",
-    ],
-)
-
-cuda_py_test(
     name = "monte_carlo_test",
     size = "small",
     srcs = ["python/kernel_tests/monte_carlo_test.py"],
diff --git a/tensorflow/contrib/bayesflow/__init__.py b/tensorflow/contrib/bayesflow/__init__.py
index 89dfa583a4..f868203826 100644
--- a/tensorflow/contrib/bayesflow/__init__.py
+++ b/tensorflow/contrib/bayesflow/__init__.py
@@ -21,7 +21,6 @@ from __future__ import division
 from __future__ import print_function
 
 # pylint: disable=unused-import,line-too-long
-from tensorflow.contrib.bayesflow.python.ops import custom_grad
 from tensorflow.contrib.bayesflow.python.ops import hmc
 from tensorflow.contrib.bayesflow.python.ops import metropolis_hastings
 from tensorflow.contrib.bayesflow.python.ops import monte_carlo
@@ -31,7 +30,6 @@ from tensorflow.python.util.all_util import remove_undocumented
 
 
 _allowed_symbols = [
-    'custom_grad',
     'entropy',
     'hmc',
     'metropolis_hastings',
diff --git a/tensorflow/contrib/bayesflow/python/kernel_tests/custom_grad_test.py b/tensorflow/contrib/bayesflow/python/kernel_tests/custom_grad_test.py
deleted file mode 100644
index 1250765d09..0000000000
--- a/tensorflow/contrib/bayesflow/python/kernel_tests/custom_grad_test.py
+++ /dev/null
@@ -1,157 +0,0 @@
-# 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 Custom Gradient Ops."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as np
-
-from tensorflow.contrib.bayesflow.python.ops import custom_grad_impl
-from tensorflow.python.framework import constant_op
-from tensorflow.python.framework import dtypes
-from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import gradients_impl
-from tensorflow.python.ops import init_ops
-from tensorflow.python.ops import math_ops
-from tensorflow.python.ops import variable_scope
-from tensorflow.python.ops import variables
-from tensorflow.python.platform import test
-
-
-cg = custom_grad_impl
-
-
-class CustomGradientTest(test.TestCase):
-
-  def test_works_correctly(self):
-    with self.test_session() as sess:
-      f = lambda x: x**2 / 2
-      g = lambda x: (x - 1)**3 / 3
-      x_ = np.linspace(-100, 100, int(1e4)) + [0.]
-
-      x = constant_op.constant(x_)
-      fx = cg.custom_gradient(f(x), g(x), x)
-      gx = gradients_impl.gradients(fx, x)[0]
-      [fx_, gx_] = sess.run([fx, gx])
-
-      self.assertAllClose(f(x_), fx_)
-      self.assertAllClose(g(x_), gx_)
-
-  def test_works_correctly_both_f_g_zero(self):
-    with self.test_session() as sess:
-      f = lambda x: x**2 / 2
-      g = lambda x: x**3 / 3
-      x_ = np.linspace(-100, 100, int(1e4)) + [0.]
-
-      x = constant_op.constant(x_)
-      fx = cg.custom_gradient(f(x), g(x), x)
-      gx = gradients_impl.gradients(fx, x)[0]
-      [fx_, gx_] = sess.run([fx, gx])
-
-      self.assertAllClose(f(x_), fx_)
-      self.assertAllClose(g(x_), gx_)
-
-  def test_works_correctly_vector_of_vars(self):
-    with self.test_session() as sess:
-      x = variable_scope.get_variable(
-          name="x",
-          shape=[],
-          dtype=dtypes.float32,
-          initializer=init_ops.constant_initializer(2))
-      y = variable_scope.get_variable(
-          name="y",
-          shape=[],
-          dtype=dtypes.float32,
-          initializer=init_ops.constant_initializer(3))
-      sess.run([variables.global_variables_initializer()])
-
-      f = lambda z: z[0] * z[1]
-      g = lambda z: z[0]**2 * z[1]**2 / 2
-
-      z = array_ops.stack([x, y])
-      fz = cg.custom_gradient(f(z), g(z), z)
-      gz = gradients_impl.gradients(fz, variables.trainable_variables())
-      [z_, fz_, gx_, gy_] = sess.run([z, fz, gz[0], gz[1]])
-
-      self.assertEqual(f(z_), fz_)
-      self.assertEqual(g(z_), gx_)
-      self.assertEqual(g(z_), gy_)
-
-  def test_works_correctly_side_vars(self):
-    with self.test_session() as sess:
-      x_ = np.float32(2.1)  # Adding extra tenth to force imprecision.
-      y_ = np.float32(3.1)
-      x = variable_scope.get_variable(
-          name="x",
-          shape=[],
-          dtype=dtypes.float32,
-          initializer=init_ops.constant_initializer(x_))
-      y = variable_scope.get_variable(
-          name="y",
-          shape=[],
-          dtype=dtypes.float32,
-          initializer=init_ops.constant_initializer(y_))
-      sess.run([variables.global_variables_initializer()])
-
-      f = lambda x: x * y
-      g = lambda z: math_ops.square(x) * y
-
-      fx = cg.custom_gradient(f(x), g(x), x)
-      gx = gradients_impl.gradients(fx, variables.trainable_variables())
-      [x_, fx_, gx_] = sess.run([x, fx, gx[0]])
-      gy_ = gx[1]
-
-      self.assertEqual(x_ * y_, fx_)
-      self.assertEqual(np.square(x_) * y_, gx_)
-      self.assertEqual(None, gy_)
-
-  def test_works_correctly_fx_gx_manually_stopped(self):
-    with self.test_session() as sess:
-      x_ = np.float32(2.1)  # Adding extra tenth to force imprecision.
-      y_ = np.float32(3.1)
-      x = variable_scope.get_variable(
-          name="x",
-          shape=[],
-          dtype=dtypes.float32,
-          initializer=init_ops.constant_initializer(x_))
-      y = variable_scope.get_variable(
-          name="y",
-          shape=[],
-          dtype=dtypes.float32,
-          initializer=init_ops.constant_initializer(y_))
-      sess.run([variables.global_variables_initializer()])
-
-      stop = array_ops.stop_gradient  # For readability.
-
-      # Basically we need to stop the `x` portion of `f`. And when we supply the
-      # arg to `custom_gradient` we need to stop the complement, i.e., the `y`
-      # part.
-      f = lambda x: stop(x) * y
-      g = lambda x: stop(math_ops.square(x)) * y
-      fx = cg.custom_gradient(f(x), g(x), x + stop(y),
-                              fx_gx_manually_stopped=True)
-
-      gx = gradients_impl.gradients(fx, variables.trainable_variables())
-      [x_, fx_, gx_, gy_] = sess.run([x, fx, gx[0], gx[1]])
-
-      self.assertEqual(x_ * y_, fx_)
-      self.assertEqual(np.square(x_) * y_, gx_)
-      self.assertEqual(x_, gy_)
-
-
-if __name__ == "__main__":
-  test.main()
diff --git a/tensorflow/contrib/bayesflow/python/ops/custom_grad.py b/tensorflow/contrib/bayesflow/python/ops/custom_grad.py
deleted file mode 100644
index c8218c57cc..0000000000
--- a/tensorflow/contrib/bayesflow/python/ops/custom_grad.py
+++ /dev/null
@@ -1,34 +0,0 @@
-# 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.
-# ==============================================================================
-"""Functions for specifying custom gradients.
-
-See @{tf.contrib.bayesflow.custom_grad.custom_gradient}.
-"""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-# go/tf-wildcard-import
-# pylint: disable=wildcard-import
-from tensorflow.contrib.bayesflow.python.ops.custom_grad_impl import *
-# pylint: enable=wildcard-import
-from tensorflow.python.util.all_util import remove_undocumented
-
-_allowed_symbols = [
-    'custom_gradient',
-]
-
-remove_undocumented(__name__, _allowed_symbols)
diff --git a/tensorflow/contrib/bayesflow/python/ops/custom_grad_impl.py b/tensorflow/contrib/bayesflow/python/ops/custom_grad_impl.py
deleted file mode 100644
index 927cc28f67..0000000000
--- a/tensorflow/contrib/bayesflow/python/ops/custom_grad_impl.py
+++ /dev/null
@@ -1,138 +0,0 @@
-# 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.
-# ==============================================================================
-"""Functions for specifying custom gradients.
-
-@@custom_gradient
-
-"""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-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 math_ops
-
-__all__ = [
-    'custom_gradient',
-]
-
-
-def is_list_like(x):
-  return isinstance(x, (tuple, list))
-
-
-def identity(x, dtype=None, name=None):
-  return array_ops.identity(ops.convert_to_tensor(
-      x, dtype=dtype, name=name), name=name)
-
-
-def custom_gradient(fx, gx, x, fx_gx_manually_stopped=False, name=None):
-  """Embeds a custom gradient into a `Tensor`.
-
-  This function works by clever application of `stop_gradient`. I.e., observe
-  that:
-
-  ```none
-  h(x) = stop_gradient(f(x)) + stop_gradient(g(x)) * (x - stop_gradient(x))
-  ```
-
-  is such that `h(x) == stop_gradient(f(x))` and
-  `grad[h(x), x] == stop_gradient(g(x)).`
-
-  In addition to scalar-domain/scalar-range functions, this function also
-  supports tensor-domain/scalar-range functions.
-
-  Partial Custom Gradient:
-
-  Suppose `h(x) = htilde(x, y)`. Note that `dh/dx = stop(g(x))` but `dh/dy =
-  None`. This is because a `Tensor` cannot have only a portion of its gradient
-  stopped. To circumvent this issue, one must manually `stop_gradient` the
-  relevant portions of `f`, `g`. For example see the unit-test,
-  `test_works_correctly_fx_gx_manually_stopped`.
-
-  Args:
-    fx: `Tensor`. Output of function evaluated at `x`.
-    gx: `Tensor` or list of `Tensor`s. Gradient of function at (each) `x`.
-    x: `Tensor` or list of `Tensor`s. Args of evaluation for `f`.
-    fx_gx_manually_stopped: Python `bool` indicating that `fx`, `gx` manually
-      have `stop_gradient` applied.
-    name: Python `str` name prefixed to Ops created by this function.
-
-  Returns:
-    fx: Floating-type `Tensor` equal to `f(x)` but which has gradient
-      `stop_gradient(g(x))`.
-  """
-  def maybe_stop(x):
-    if fx_gx_manually_stopped:
-      return x
-    return array_ops.stop_gradient(x)
-  with ops.name_scope(name, 'custom_gradient', [fx, gx, x]):
-    fx = ops.convert_to_tensor(fx, name='fx')
-    # We don't want to bother eagerly computing `gx` since we may not even need
-    # it.
-    with ops.control_dependencies([fx]):
-      if is_list_like(x):
-        x = [identity(x_, name='x') for x_ in x]
-      else:
-        x = [identity(x, name='x')]
-
-      if is_list_like(gx):
-        gx = [identity(gx_, dtype=fx.dtype, name='gx')
-              for gx_ in gx]
-      else:
-        gx = [identity(gx, dtype=fx.dtype, name='gx')]
-
-      override_grad = []
-      for x_, gx_ in zip(x, gx):
-        # Observe: tf.gradients(f(x), x)[i].shape == x[i].shape
-        # thus we check that the user is supplying correct shapes.
-        equal_shape = check_ops.assert_equal(
-            array_ops.shape(x_),
-            array_ops.shape(gx_),
-            message='Each `x` must have the same shape as each `gx`.')
-        with ops.control_dependencies([equal_shape]):
-          # IEEE754 ensures `(x-x)==0.` and that `0.*x==0.` so we make sure to
-          # write the code this way, rather than, e.g.,
-          # `sum_x * stop(gx) + stop(fx - sum_x * gx)`.
-          # For more discussion regarding the relevant portions of the IEEE754
-          # standard, see the StackOverflow question,
-          # "Is there a floating point value of x, for which x-x == 0 is false?"
-          # http://stackoverflow.com/q/2686644
-          zeros_like_x_ = x_ - array_ops.stop_gradient(x_)
-          override_grad.append(math_ops.reduce_sum(
-              maybe_stop(gx_) * zeros_like_x_))
-      override_grad = sum(override_grad)
-      override_grad /= math_ops.cast(array_ops.size(fx),
-                                     dtype=fx.dtype.base_dtype)
-
-      # Proof of correctness:
-      #
-      #  f(x) = x * stop[gx] + stop[fx - x * gx]
-      #       = stop[fx]
-      #
-      #  g(x) = grad[fx]
-      #       = stop[gx] + grad[stop[fx - x * gx]]
-      #       = stop[gx] + 0
-      #
-      # Notice that when x is zero it still works:
-      # grad[x * stop(gx) + stop(fx - x * gx)] = 1 * stop[gx] + 0 = stop[gx]
-      #
-      # The proof is similar for the tensor-domain case, except that we
-      # `reduce_sum` the `stop[gx] * (x - stop[x])` then rescale by
-      # `tf.size(fx)` since this reduced version is broadcast to `fx`.
-      return maybe_stop(fx) + override_grad