* Extend Geometric Brownian Motion to work with dynamically shaped nu…

…mber of times.

* Add algorithm type to mutilvariate_normal sampler

PiperOrigin-RevId: 487269015

tf_quant_finance/math/random_ops/multivariate_normal.py

@@ -266,7 +266,7 @@ def _mvnormal_pseudo(sample_shape,
       raise ValueError('`seed` should be specified if the `random_type` is '
                        '`STATELESS` or `STATELESS_ANTITHETIC`')
     samples = tf.random.stateless_normal(
-        shape=output_shape, dtype=dtype, seed=seed)
+        shape=output_shape, dtype=dtype, seed=seed, alg='philox')
   if scale_matrix is None:
     return mean + samples
   else:

tf_quant_finance/models/geometric_brownian_motion/BUILD

@@ -30,6 +30,7 @@ py_library(
         "//tf_quant_finance/math/pde",
         "//tf_quant_finance/models:ito_process",
         "//tf_quant_finance/models:utils",
+        "//tf_quant_finance/utils",
         # tensorflow dep,
     ],
 )
@@ -41,6 +42,7 @@ py_library(
         "//tf_quant_finance/math/pde",
         "//tf_quant_finance/models:ito_process",
         "//tf_quant_finance/models:utils",
+        "//tf_quant_finance/utils",
         # tensorflow dep,
     ],
 )

tf_quant_finance/models/geometric_brownian_motion/geometric_brownian_motion_test.py

@@ -873,26 +873,83 @@ def sample_fn():
                      * np.array([0.1, 0.5, 1.0]) + np.log(2.))
     self.assertAllClose(tf.squeeze(mean), expected_mean, atol=1e-2, rtol=1e-2)
 
-  def test_multiivariate_xla_compatible(self):
-    """Tests that multiivariate GBM sampling is XLA-compatible."""
+  def test_multivariate_xla_compatible(self):
+    """Tests that multivariate GBM sampling is XLA-compatible."""
     corr_matrix = [[1, 0.1], [0.1, 1]]
     process = tff.models.MultivariateGeometricBrownianMotion(
         dim=2, means=0.05, volatilities=[0.1, 0.2], corr_matrix=corr_matrix,
         dtype=tf.float64)
     times = [0.1, 0.5, 1.0]
     initial_state = [1.0, 2.0]
-    @tf.function
+    @tf.function(jit_compile=True)
     def sample_fn():
       return process.sample_paths(
           times=times, initial_state=initial_state, num_samples=10000)
-    samples = tf.xla.experimental.compile(sample_fn)[0]
+    samples = sample_fn()
     log_s = tf.math.log(samples)
     mean = tf.reduce_mean(log_s, axis=0)
     expected_mean = ((process._means - process._vols**2 / 2)
                      * np.array(np.expand_dims(times, -1))
                      + np.log(initial_state))
     self.assertAllClose(mean, expected_mean, atol=1e-2, rtol=1e-2)
 
+  def test_multivariate_dynamic_inputs(self):
+    """Tests that GBM sampling can accept dynamically shaped inputs."""
+    corr_matrix = [[1, 0.1], [0.1, 1]]
+    times = [0.1, 0.5, 1.0]
+    initial_state = [1.0, 2.0]
+    @tf.function(input_signature=[
+        tf.TensorSpec([None], dtype=tf.float64, name='initial_state'),
+        tf.TensorSpec([None], dtype=tf.float64, name='means'),
+        tf.TensorSpec([None], dtype=tf.float64, name='volatilities'),
+        tf.TensorSpec([2, 2], dtype=tf.float64, name='corr_matrix'),
+        tf.TensorSpec([None], dtype=tf.float64, name='times'),
+    ])
+    def sample_fn(initial_state, means, volatilities, corr_matrix, times):
+      process = tff.models.MultivariateGeometricBrownianMotion(
+          dim=2, means=means, volatilities=volatilities,
+          corr_matrix=corr_matrix, dtype=tf.float64)
+      return process.sample_paths(
+          times=times, initial_state=initial_state, num_samples=10000)
+    means = np.array([0.05], dtype=np.float64)
+    volatilities = np.array([0.1, 0.2], dtype=np.float64)
+    samples = sample_fn(initial_state=initial_state,
+                        means=means, volatilities=volatilities,
+                        corr_matrix=corr_matrix, times=times)
+    log_s = tf.math.log(samples)
+    expected_means = ((means - volatilities**2 / 2)
+                      * np.array(np.expand_dims(times, -1))
+                      + np.log(initial_state))
+    actual_means = tf.reduce_mean(log_s, axis=0)
+    self.assertAllClose(actual_means, expected_means, atol=1e-2, rtol=1e-2)
+
+  def test_univariate_dynamic_inputs(self):
+    """Tests that GBM sampling can accept dynamically shaped inputs."""
+    times = [0.1, 0.5, 1.0]
+    initial_state = [[1.0], [2.0]]
+    @tf.function(input_signature=[
+        tf.TensorSpec([None, None], dtype=tf.float64, name='initial_state'),
+        tf.TensorSpec([None, 1], dtype=tf.float64, name='mean'),
+        tf.TensorSpec([None, 1], dtype=tf.float64, name='volatility'),
+        tf.TensorSpec([None], dtype=tf.float64, name='times'),
+    ])
+    def sample_fn(initial_state, mean, volatility, times):
+      process = tff.models.GeometricBrownianMotion(
+          mean=mean, volatility=volatility, dtype=tf.float64)
+      return process.sample_paths(
+          times=times, initial_state=initial_state, num_samples=10000)
+    mean = np.array([[0.05], [0.1]], dtype=np.float64)
+    volatility = np.array([[0.1]], dtype=np.float64)
+    samples = sample_fn(initial_state=initial_state,
+                        mean=mean, volatility=volatility, times=times)
+    log_s = tf.math.log(samples)
+    expected_mean = ((mean - volatility**2 / 2)
+                     * np.array(np.expand_dims(times, 0))
+                     + np.log(initial_state))
+    actual_mean = tf.reduce_mean(log_s, axis=1)
+    self.assertAllClose(actual_mean, expected_mean[..., np.newaxis],
+                        atol=1e-2, rtol=1e-2)
+
   def test_normal_draws_shape_mismatch_2d(self):
     """Error is raised if `dim` is mismatched with the one from normal_draws."""
     dtype = tf.float64

tf_quant_finance/models/geometric_brownian_motion/multivariate_geometric_brownian_motion.py

@@ -16,6 +16,7 @@
 
 import tensorflow.compat.v2 as tf
 
+from tf_quant_finance import utils as tff_utils
 from tf_quant_finance.math.pde import fd_solvers
 from tf_quant_finance.models import ito_process
 from tf_quant_finance.models import utils
@@ -217,7 +218,7 @@ def sample_paths(self,
             + tf.zeros([num_samples, 1], dtype=self._dtype))
       # Shape [num_samples, 1, dim]
       initial_state = tf.expand_dims(initial_state, axis=1)
-      num_requested_times = times.shape[0]
+      num_requested_times = tff_utils.get_shape(times)[0]
       return self._sample_paths(
           times=times, num_requested_times=num_requested_times,
           initial_state=initial_state,
@@ -246,8 +247,8 @@ def _sample_paths(self,
     else:
       # Shape [num_time_points, num_samples, dim]
       normal_draws = tf.transpose(normal_draws, [1, 0, 2])
-      num_samples = tf.shape(normal_draws)[1]
-      draws_dim = normal_draws.shape[2]
+      num_samples = tff_utils.get_shape(normal_draws)[1]
+      draws_dim = tff_utils.get_shape(normal_draws)[2]
       if self._dim != draws_dim:
         raise ValueError(
             "`dim` should be equal to `normal_draws.shape[2]` but are "
@@ -653,7 +654,7 @@ def second_order_coeff_fn(t, coord_grid):
 
     # We currently have [dim, dim] as innermost dimensions, but the returned
     # tensor must have [dim, dim] as outermost dimensions.
-    rank = len(sigma.shape.as_list())
+    rank = sigma.shape.rank
     perm = [rank - 2, rank - 1] + list(range(rank - 2))
     sigma_times_sigma_t = tf.transpose(sigma_times_sigma_t, perm)
     return sigma_times_sigma_t / 2
@@ -663,7 +664,7 @@ def first_order_coeff_fn(t, coord_grid):
 
     # We currently have [dim] as innermost dimension, but the returned
     # tensor must have [dim] as outermost dimension.
-    rank = len(mu.shape.as_list())
+    rank = mu.shape.rank
     perm = [rank - 1] + list(range(rank - 1))
     mu = tf.transpose(mu, perm)
     return mu

tf_quant_finance/models/geometric_brownian_motion/univariate_geometric_brownian_motion.py

@@ -16,6 +16,7 @@
 
 import tensorflow.compat.v2 as tf
 
+from tf_quant_finance import utils as tff_utils
 from tf_quant_finance.math import piecewise as pw
 from tf_quant_finance.math.pde import fd_solvers
 from tf_quant_finance.models import ito_process
@@ -225,7 +226,7 @@ def sample_paths(self,
           dtype=self._dtype,
           name='initial_state')
 
-      num_requested_times = times.shape[-1]
+      num_requested_times = tff_utils.get_shape(times)[-1]
       return self._sample_paths(
           times=times,
           num_requested_times=num_requested_times,