<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Copyright 2019 Google LLC.
Licensed under the Apache License, Version 2.0 (the "License");
Step1:  Root Search
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https
Step2:  This notebook demonstrates the use of low level Tensorflow Quant Finance tools for root finding using Brent's method with emphasis on the following aspects
Step3:  Brent's Method
Find the risk free rate for an interest rate swap
Step4:  Speedup from XLA
<ASSISTANT_TASK:>
Python Code:
#@title Licensed under the Apache License, Version 2.0 (the "License"); { display-mode: "form" }
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://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.
Explanation: Copyright 2019 Google LLC.
Licensed under the Apache License, Version 2.0 (the "License");
End of explanation
#@title Upgrade to TensorFlow 2.5+
!pip install --upgrade tensorflow
#@title Install TF Quant Finance
!pip install tf-quant-finance
Explanation: Root Search
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/google/tf-quant-finance/blob/master/tf_quant_finance/examples/jupyter_notebooks/Root_Search.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/google/tf-quant-finance/blob/master/tf_quant_finance/examples/jupyter_notebooks/Root_Search.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
  </td>
</table>
End of explanation
#@title Imports { display-mode: "form" }
import tensorflow as tf
 # tff for Tensorflow Finance
import tf_quant_finance as tff 
root_search = tff.math.root_search
import warnings
warnings.filterwarnings("ignore",
                        category=FutureWarning)  # suppress printing warnings
!nvidia-smi
Explanation: This notebook demonstrates the use of low level Tensorflow Quant Finance tools for root finding using Brent's method with emphasis on the following aspects:
Write Once: Tensorflow supports GPUs without any significant code changes. The same model can be run on CPU and GPU
XLA Acceleration: The XLA compiler can reduce overhead associated with the Tensorflow graph by fusing operators.
End of explanation
#@title Search Range
number_of_tenors = 100  #@param
swap_value = 0.5  #@param
dtype = tf.float64
tenors = range(1, number_of_tenors + 1)
t = tf.constant(tenors, dtype=dtype)
v = tf.constant(swap_value, dtype=dtype)
def valuation_fn(x):
  return tf.reduce_logsumexp(-x * t) - v
# Wrap with TF function for better performance
root_search_tf = tf.function(root_search.brentq)
def run_on_device(device):
  with tf.device(device):
    return root_search_tf(
      valuation_fn, tf.constant(0, dtype=dtype), tf.constant(1, dtype=dtype))
## TFF on CPU
brent_result = run_on_device('/cpu:0')
estimated_root, objective_at_root, num_iterations, converged = brent_result
print("------------------------")
print("Tensorflow CPU (with auto-threading)")
print("Converged:", converged)
print("Estimated root:", estimated_root)
print("Objective at root:", objective_at_root)
print("Number of search steps:", num_iterations)
print("Timing:")
%timeit -n 100 run_on_device('/cpu:0')
print("------------------------")
## TFF on GPU
brent_result = run_on_device('/gpu:0')
estimated_root, objective_at_root, num_iterations, converged = brent_result
print("------------------------")
print("Tensorflow GPU")
print("Converged:", converged)
print("Estimated root:", estimated_root)
print("Objective at root:", objective_at_root)
print("Number of search steps:", num_iterations)
print("Timing:")
%timeit -n 100 run_on_device('/gpu:0')
print("------------------------")
Explanation: Brent's Method
Find the risk free rate for an interest rate swap:
$$f(x) = \log(\sum_{i=0}^{N} e^{-r_i t_i}) - V_{swap}$$
End of explanation
#@title Search Range
number_of_tenors = 100  #@param
swap_value = 0.5  #@param
dtype = tf.float64
tenors = range(1, number_of_tenors + 1)
tf.compat.v1.reset_default_graph()
t = tf.constant(tenors, dtype=dtype)
v = tf.constant(swap_value, dtype=dtype)
def valuation_fn(x):
  return tf.reduce_logsumexp(-x * t) - v
# Wrap with TF function for better performance
root_search_xla = tf.function(root_search.brentq, jit_compile=True)
def run_on_device(device):
  with tf.device(device):
    return root_search_xla(
      valuation_fn, tf.constant(0, dtype=dtype), tf.constant(1, dtype=dtype))
## TFF on CPU compiled with XLA
brent_result = run_on_device('/cpu:0')
estimated_root, objective_at_root, num_iterations, converged = brent_result
print("------------------------")
print("Tensorflow CPU (compiled with XLA)")
print("Converged:", converged)
print("Estimated root:", estimated_root)
print("Objective at root:", objective_at_root)
print("Number of search steps:", num_iterations)
print("Timing:")
%timeit -n 100 brent_result = run_on_device('/cpu:0')
print("------------------------")
## TFF on GPU compiled with XLA
brent_result = run_on_device('/gpu:0')
estimated_root, objective_at_root, num_iterations, converged = brent_result
print("------------------------")
print("Tensorflow GPU (compiled with XLA)")
print("Converged:", converged)
print("Estimated root:", estimated_root)
print("Objective at root:", objective_at_root)
print("Number of search steps:", num_iterations)
print("Timing:")
%timeit -n 100 brent_result = run_on_device('/gpu:0')
print("------------------------")
Explanation: Speedup from XLA
End of explanation
<END_TASK>