code: def test_batch_with_nans(self):
      idxs = np.arange(20, dtype=np.float32)
      bounds = np.reshape(np.exp(idxs), [4, -1])
      roots = np.reshape(1.0 / (20.0 - idxs), [4, -1])
  
      def objective_fn(x):
          return tf.where(
              x < -bounds,
              np.nan,
              tf.where(x > bounds, np.inf, (x - roots) ** 3),
          )
  
      low, high = self.evaluate(root_search.bracket_root(objective_fn))
      f_low, f_high = self.evaluate(
          (objective_fn(low), objective_fn(high))
      )
      self.assertAllFinite(f_low)
      self.assertAllFinite(f_high)
      self.assertAllTrue(low < roots)
      self.assertAllTrue(high > roots)
  
imports: import numpy
  import scipy.optimize
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import special_math
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import root_search
label_desc: Use image node_id 1 for calling the BracketRootTest obj's underlying member method code with example usage: obj.test_batch_with_nans() and returns: tf