Support dynamic shape in tfp.math.find_root_secant.

PiperOrigin-RevId: 379291021

Author: davmre

tensorflow_probability/python/math/root_search.py

@@ -218,18 +218,17 @@ def f(x):
   num_iterations = tf.zeros_like(position, dtype=tf.int32)
   max_iterations = tf.convert_to_tensor(max_iterations, dtype=tf.int32)
   max_iterations = tf.broadcast_to(
-      max_iterations, name='max_iterations', shape=position.shape)
+      max_iterations, name='max_iterations', shape=ps.shape(position))
 
   # Compute the step from `next_position` if present. This covers the case where
   # a user has two starting points, which bound the root or has a specific step
   # size in mind.
   if next_position is None:
-    epsilon = tf.constant(1e-4, dtype=position.dtype, shape=position.shape)
-    step = position * epsilon + tf.sign(position) * epsilon
+    step = (position + tf.sign(position)) * 1e-4
   else:
     step = next_position - initial_position
 
-  finished = tf.constant(False, shape=position.shape)
+  finished = tf.zeros(ps.shape(position), dtype=tf.bool)
 
   # Negate `stopping_condition` to determine if the search should continue.
   # This means, in particular, that tf.reduce_*all* will return only when the

tensorflow_probability/python/math/root_search_test.py

@@ -22,6 +22,7 @@
 import numpy as np
 import scipy.optimize as optimize
 
+import tensorflow.compat.v1 as tf1
 import tensorflow.compat.v2 as tf
 import tensorflow_probability as tfp
 
@@ -172,6 +173,17 @@ def test_secant_invalid_max_iterations(self):
           tfp.math.find_root_secant(
               f, guess, max_iterations=-1, validate_args=True))
 
+  def test_secant_non_static_shape(self):
+    if tf.executing_eagerly():
+      self.skipTest('Test uses dynamic shapes.')
+
+    f = lambda x: (x - 1.) * (x + 1)
+    initial_position = tf1.placeholder_with_default([1., 1., 1.], shape=None)
+    self.assertAllClose(
+        tfp.math.find_root_secant(
+            f, initial_position).objective_at_estimated_root,
+        [0., 0., 0.])
+
 
 @test_util.test_all_tf_execution_regimes
 class ChandrupatlaRootSearchTest(test_util.TestCase):
@@ -266,6 +278,18 @@ def test_chandrupatla_automatically_selects_bounds(self):
         position_tolerance=1e-8)
     self.assertAllClose(value_at_roots, tf.zeros_like(value_at_roots))
 
+  def test_chandrupatla_non_static_shape(self):
+    if tf.executing_eagerly():
+      self.skipTest('Test uses dynamic shapes.')
+
+    f = lambda x: (x - 1.) * (x + 1)
+    low = tf1.placeholder_with_default([-100., -100., -100.], shape=None)
+    high = tf1.placeholder_with_default([100., 100., 100.], shape=None)
+    self.assertAllClose(
+        tfp.math.find_root_chandrupatla(
+            f, low=low, high=high).objective_at_estimated_root,
+        [0., 0., 0.])
+
 
 @test_util.test_all_tf_execution_regimes
 class BracketRootTest(test_util.TestCase):