Python tensorflow.compat.v2.ones_like() Examples

def ones_like(a, dtype=None):
  """Returns an array of ones with the shape and type of the input array.

  Args:
    a: array_like. Could be an ndarray, a Tensor or any object that can be
      converted to a Tensor using `tf.convert_to_tensor`.
    dtype: Optional, defaults to dtype of the input array. The type of the
      resulting ndarray. Could be a python type, a NumPy type or a TensorFlow
      `DType`.

  Returns:
    An ndarray.
  """
  if isinstance(a, arrays_lib.ndarray):
    a = a.data
  if dtype is None:
    dtype = utils.result_type(a)
  else:
    dtype = utils.result_type(dtype)
  return arrays_lib.tensor_to_ndarray(tf.ones_like(a, dtype)) 

def labels_of_top_ranked_predictions_in_batch(labels, predictions):
  """Applying tf.metrics.mean to this gives precision at 1.

  Args:
    labels: minibatch of dense 0/1 labels, shape [batch_size rows, num_classes]
    predictions: minibatch of predictions of the same shape

  Returns:
    one-dimension tensor top_labels, where top_labels[i]=1.0 iff the
    top-scoring prediction for batch element i has label 1.0
  """
  indices_of_top_preds = tf.cast(tf.argmax(input=predictions, axis=1), tf.int32)
  batch_size = tf.reduce_sum(input_tensor=tf.ones_like(indices_of_top_preds))
  row_indices = tf.range(batch_size)
  thresholded_labels = tf.where(labels > 0.0, tf.ones_like(labels),
                                tf.zeros_like(labels))
  label_indices_to_gather = tf.transpose(
      a=tf.stack([row_indices, indices_of_top_preds]))
  return tf.gather_nd(thresholded_labels, label_indices_to_gather) 

def test_correctness_1d(self, use_analytic_pricing, error_tol):
    """Tests model with constant parameters in 1 dimension."""
    dtype = tf.float64

    discount_rate_fn = lambda x: 0.01 * tf.ones_like(x, dtype=dtype)
    expiries = np.array([1.0])
    maturities = np.array([5.0])
    strikes = np.exp(-0.01 * maturities) / np.exp(-0.01 * expiries)
    price = tff.models.hull_white.bond_option_price(
        strikes=strikes,
        expiries=expiries,
        maturities=maturities,
        dim=1,
        mean_reversion=self.mean_reversion_1d,
        volatility=self.volatility_1d,
        discount_rate_fn=discount_rate_fn,
        use_analytic_pricing=use_analytic_pricing,
        num_samples=500000,
        time_step=0.1,
        random_type=tff.math.random.RandomType.PSEUDO_ANTITHETIC,
        dtype=dtype)
    self.assertEqual(price.dtype, dtype)
    self.assertAllEqual(price.shape, [1, 1])
    price = self.evaluate(price)
    self.assertAllClose(price, [[0.02817777]], rtol=error_tol, atol=error_tol) 

def test_sample_paths_dtypes(self):
    """Sampled paths have the expected dtypes."""
    for dtype in [np.float32, np.float64]:
      drift_fn = lambda t, x: tf.sqrt(t) * tf.ones_like(x, dtype=t.dtype)
      vol_fn = lambda t, x: t * tf.ones([1, 1], dtype=t.dtype)
      process = GenericItoProcess(
          dim=1, drift_fn=drift_fn, volatility_fn=vol_fn, dtype=dtype)

      paths = self.evaluate(
          process.sample_paths(
              times=[0.1, 0.2],
              num_samples=10,
              initial_state=[0.1],
              time_step=0.01,
              seed=123))
      self.assertEqual(paths.dtype, dtype)

  # Several tests below are unit tests for GenericItoProcess.fd_solver_backward:
  # they mock out the pde solver and check only the conversion of SDE to PDE,
  # but not PDE solving. There are also integration tests further below. 

def test_correctness_2d(self, use_analytic_pricing, error_tol):
    """Tests model with constant parameters in 2 dimension."""
    dtype = tf.float64

    discount_rate_fn = lambda x: 0.01 * tf.ones_like(x, dtype=dtype)
    price = tff.models.hull_white.cap_floor_price(
        strikes=self.strikes,
        expiries=self.expiries,
        maturities=self.maturities,
        daycount_fractions=self.daycount_fractions,
        notional=100.0,
        dim=2,
        mean_reversion=self.mean_reversion_2d,
        volatility=self.volatility_2d,
        reference_rate_fn=discount_rate_fn,
        use_analytic_pricing=use_analytic_pricing,
        num_samples=500000,
        time_step=0.1,
        random_type=tff.math.random.RandomType.PSEUDO_ANTITHETIC,
        dtype=dtype)
    self.assertEqual(price.dtype, dtype)
    self.assertAllEqual(price.shape, [1, 2])
    price = self.evaluate(price)
    self.assertAllClose(price, [[0.4072088281493774, 0.2016075430673558]],
                        rtol=error_tol, atol=error_tol) 

def test_expected_continuation(self):
    """Tests that expected continuation works in V=1 case.

    In particular this verifies that the regression done to get the expected
    continuation value is performed on those elements which have a positive
    exercise value.
    """
    for dtype in (np.float32, np.float64):
      a = tf.range(start=-2, limit=3, delta=1, dtype=dtype)
      design = tf.concat([a, a], axis=0)
      design = tf.concat([[tf.ones_like(design), design]], axis=1)

      # These values ensure that the expected continuation value is `(1,...,1).`
      exercise_now = tf.expand_dims(
          tf.concat([tf.ones_like(a), tf.zeros_like(a)], axis=0), -1)
      cashflow = tf.expand_dims(
          tf.concat([tf.ones_like(a), -tf.ones_like(a)], axis=0), -1)

      expected_exercise = lsm.expected_exercise_fn(
          design, cashflow, exercise_now)
      self.assertAllClose(expected_exercise, tf.ones_like(cashflow)) 

def test_sample_paths_dtypes(self):
    """Sampled paths have the expected dtypes."""
    for dtype in [np.float32, np.float64]:
      drift_fn = lambda t, x: tf.sqrt(t) * tf.ones_like(x, dtype=t.dtype)
      vol_fn = lambda t, x: t * tf.ones([1, 1], dtype=t.dtype)

      paths = self.evaluate(
          euler_sampling.sample(
              dim=1,
              drift_fn=drift_fn, volatility_fn=vol_fn,
              times=[0.1, 0.2],
              num_samples=10,
              initial_state=[0.1],
              time_step=0.01,
              seed=123,
              dtype=dtype))

      self.assertEqual(paths.dtype, dtype) 

def test_maybe_update_along_axis(self, dtype):
    """Tests that the values are updated correctly."""
    tensor = tf.ones([5, 4, 3, 2], dtype=dtype)
    new_tensor = tf.zeros([5, 4, 1, 2], dtype=dtype)
    @tf.function
    def maybe_update_along_axis(do_update):
      return utils.maybe_update_along_axis(
          tensor=tensor, new_tensor=new_tensor, axis=1, ind=2,
          do_update=do_update)
    updated_tensor = maybe_update_along_axis(True)
    with self.subTest(name='Shape'):
      self.assertEqual(updated_tensor.shape, tensor.shape)
    with self.subTest(name='UpdatedVals'):
      self.assertAllEqual(updated_tensor[:, 2, :, :],
                          tf.zeros_like(updated_tensor[:, 2, :, :]))
    with self.subTest(name='NotUpdatedVals'):
      self.assertAllEqual(updated_tensor[:, 1, :, :],
                          tf.ones_like(updated_tensor[:, 2, :, :]))
    with self.subTest(name='DoNotUpdateVals'):
      not_updated_tensor = maybe_update_along_axis(False)
      self.assertAllEqual(not_updated_tensor, tensor) 

def test_correctness_time_dep_1d(self, use_analytic_pricing, error_tol):
    """Tests model with piecewise constant volatility in 1 dimension."""
    dtype = tf.float64

    discount_rate_fn = lambda x: 0.01 * tf.ones_like(x, dtype=dtype)
    expiries = np.array([1.0])
    maturities = np.array([5.0])
    strikes = np.exp(-0.01 * maturities) / np.exp(-0.01 * expiries)
    volatility = tff.math.piecewise.PiecewiseConstantFunc(
        jump_locations=[0.5], values=self.volatility_time_dep_1d,
        dtype=dtype)
    price = tff.models.hull_white.bond_option_price(
        strikes=strikes,
        expiries=expiries,
        maturities=maturities,
        dim=1,
        mean_reversion=self.mean_reversion_1d,
        volatility=volatility,
        discount_rate_fn=discount_rate_fn,
        use_analytic_pricing=use_analytic_pricing,
        num_samples=500000,
        time_step=0.1,
        random_type=tff.math.random.RandomType.PSEUDO_ANTITHETIC,
        dtype=dtype)
    self.assertEqual(price.dtype, dtype)
    self.assertAllEqual(price.shape, [1, 1])
    price = self.evaluate(price)
    self.assertAllClose(price, [[0.02237839]], rtol=error_tol, atol=error_tol) 

def sinc(x):
  def f(x):
    pi_x = x * np.pi
    return tf.where(x == 0, tf.ones_like(x), tf.math.sin(pi_x) / pi_x)
  return _scalar(f, x, True) 

def test_2d_batch(self, use_analytic_pricing, error_tol):
    """Tests model with 2d batch of options."""
    dtype = tf.float64

    discount_rate_fn = lambda x: 0.01 * tf.ones_like(x, dtype=dtype)
    expiries = np.array([[1.0, 1.0], [2.0, 2.0]])
    maturities = np.array([[5.0, 5.0], [4.0, 4.0]])
    strikes = np.exp(-0.01 * maturities) / np.exp(-0.01 * expiries)
    price = tff.models.hull_white.bond_option_price(
        strikes=strikes,
        expiries=expiries,
        maturities=maturities,
        dim=1,
        mean_reversion=self.mean_reversion_1d,
        volatility=self.volatility_1d,
        discount_rate_fn=discount_rate_fn,
        use_analytic_pricing=use_analytic_pricing,
        num_samples=500000,
        time_step=0.1,
        random_type=tff.math.random.RandomType.PSEUDO_ANTITHETIC,
        dtype=dtype)
    self.assertEqual(price.dtype, dtype)
    self.assertAllEqual(price.shape, [2, 2, 1])
    price = self.evaluate(price)
    expected = [[[0.02817777], [0.02817777]], [[0.02042677], [0.02042677]]]
    self.assertAllClose(price, expected, rtol=error_tol, atol=error_tol) 

def test_correctness_2d(self, use_analytic_pricing, error_tol):
    """Tests model with constant parameters in 2 dimension."""
    dtype = tf.float64

    discount_rate_fn = lambda x: 0.01 * tf.ones_like(x, dtype=dtype)
    expiries = np.array([1.0])
    maturities = np.array([5.0])
    strikes = np.exp(-0.01 * maturities) / np.exp(-0.01 * expiries)
    price = tff.models.hull_white.bond_option_price(
        strikes=strikes,
        expiries=expiries,
        maturities=maturities,
        dim=2,
        mean_reversion=self.mean_reversion_2d,
        volatility=self.volatility_2d,
        discount_rate_fn=discount_rate_fn,
        use_analytic_pricing=use_analytic_pricing,
        num_samples=500000,
        time_step=0.1,
        random_type=tff.math.random.RandomType.PSEUDO_ANTITHETIC,
        dtype=dtype)
    self.assertEqual(price.dtype, dtype)
    self.assertAllEqual(price.shape, [1, 2])
    price = self.evaluate(price)
    self.assertAllClose(price, [[0.02817777, 0.01309971]],
                        rtol=error_tol, atol=error_tol) 

def test_mixed_1d_batch_2d(self, use_analytic_pricing, error_tol):
    """Tests mixed 1d batch with constant parameters in 2 dimension."""
    dtype = tf.float64

    discount_rate_fn = lambda x: 0.01 * tf.ones_like(x, dtype=dtype)
    expiries = np.array([1.0, 1.0, 2.0])
    maturities = np.array([5.0, 6.0, 4.0])
    strikes = np.exp(-0.01 * maturities) / np.exp(-0.01 * expiries)
    price = tff.models.hull_white.bond_option_price(
        strikes=strikes,
        expiries=expiries,
        maturities=maturities,
        dim=2,
        mean_reversion=self.mean_reversion_2d,
        volatility=self.volatility_2d,
        discount_rate_fn=discount_rate_fn,
        use_analytic_pricing=use_analytic_pricing,
        num_samples=500000,
        time_step=0.1,
        random_type=tff.math.random.RandomType.PSEUDO_ANTITHETIC,
        dtype=dtype)
    self.assertEqual(price.dtype, dtype)
    self.assertAllEqual(price.shape, [3, 2])
    price = self.evaluate(price)
    expected = [[0.02817777, 0.01309971], [0.03436008, 0.01575343],
                [0.02042677, 0.00963248]]
    self.assertAllClose(price, expected, rtol=error_tol, atol=error_tol) 

def test_call_put(self, use_analytic_pricing, error_tol):
    """Tests mixed 1d batch with constant parameters in 2 dimension."""
    dtype = tf.float64

    discount_rate_fn = lambda x: 0.01 * tf.ones_like(x, dtype=dtype)
    expiries = np.array([1.0, 1.0, 2.0])
    maturities = np.array([5.0, 6.0, 4.0])
    strikes = np.exp(-0.01 * maturities) / np.exp(-0.01 * expiries) -0.01
    price = tff.models.hull_white.bond_option_price(
        strikes=strikes,
        expiries=expiries,
        maturities=maturities,
        is_call_options=[True, False, False],
        dim=2,
        mean_reversion=self.mean_reversion_2d,
        volatility=self.volatility_2d,
        discount_rate_fn=discount_rate_fn,
        use_analytic_pricing=use_analytic_pricing,
        num_samples=500000,
        time_step=0.1,
        random_type=tff.math.random.RandomType.PSEUDO_ANTITHETIC,
        dtype=dtype)
    self.assertEqual(price.dtype, dtype)
    self.assertAllEqual(price.shape, [3, 2])
    price = self.evaluate(price)
    expected = [[0.03325893, 0.01857569], [0.02945686, 0.01121538],
                [0.01579646, 0.0054692]]
    self.assertAllClose(price, expected, rtol=error_tol, atol=error_tol) 

def _prepare_grid(times, *params):
  """Prepares grid of times for path generation.

  Args:
    times:  Rank 1 `Tensor` of increasing positive real values. The times at
      which the path points are to be evaluated.
    *params: Parameters of the Heston model. Either scalar `Tensor`s of the
      same `dtype` or instances of `PiecewiseConstantFunc`.

  Returns:
    Tuple `(all_times, mask)`.
    `all_times` is a 1-D real `Tensor` containing all points from 'times`, the
    uniform grid of points between `[0, times[-1]]` with grid size equal to
    `time_step`, and jump locations of piecewise constant parameters The
    `Tensor` is sorted in ascending order and may contain duplicates.
    `mask` is a boolean 1-D `Tensor` of the same shape as 'all_times', showing
    which elements of 'all_times' correspond to THE values from `times`.
    Guarantees that times[0]=0 and mask[0]=False.
  """
  additional_times = []
  for param in params:
    if hasattr(param, 'is_piecewise_constant'):
      if param.is_piecewise_constant:
        # Flatten all jump locations
        additional_times.append(tf.reshape(param.jump_locations(), [-1]))
  zeros = tf.constant([0], dtype=times.dtype)
  all_times = tf.concat([zeros] + [times] + additional_times, axis=0)
  additional_times_mask = [
      tf.zeros_like(times, dtype=tf.bool) for times in additional_times]
  mask = tf.concat([
      tf.cast(zeros, dtype=tf.bool),
      tf.ones_like(times, dtype=tf.bool)
  ] + additional_times_mask, axis=0)
  perm = tf.argsort(all_times, stable=True)
  all_times = tf.gather(all_times, perm)
  mask = tf.gather(mask, perm)
  return all_times, mask 

def test_time_dependent_1d(self, random_type, seed):
    """Tests model with time dependent vol in 1 dimension."""
    for dtype in [tf.float32, tf.float64]:
      def discount_fn(x):
        return 0.01 * tf.ones_like(x, dtype=dtype)  # pylint: disable=cell-var-from-loop
      mean_reversion = tff.math.piecewise.PiecewiseConstantFunc(
          [], values=[self.mean_reversion[0]], dtype=dtype)
      volatility = tff.math.piecewise.PiecewiseConstantFunc(
          [0.1, 2.0], values=self.volatility_time_dep_1d, dtype=dtype)
      process = tff.models.hull_white.HullWhiteModel1F(
          mean_reversion=mean_reversion,
          volatility=volatility,
          initial_discount_rate_fn=discount_fn,
          dtype=dtype)
      times = np.array([0.1, 1.0, 2.0, 3.0])
      paths = process.sample_paths(
          times,
          num_samples=500000,
          random_type=random_type,
          seed=seed,
          skip=1000000)
      self.assertEqual(paths.dtype, dtype)
      self.assertAllEqual(paths.shape, [500000, 4, 1])
      paths = self.evaluate(paths)
      r_std = np.squeeze(np.std(paths, axis=0))
      expected_std = self.true_std_time_dep(
          times, np.array([0.0, 0.1, 2.0]),
          np.array(self.volatility_time_dep_1d), 0.1)
      self.assertAllClose(r_std, expected_std, rtol=1e-4, atol=1e-4) 

def test_mean_variance_correlation_generic_2d(self):
    """Tests model with generic parameters in 2 dimensions."""
    for dtype in [tf.float32, tf.float64]:
      # Mean reversion without batch dimesnion
      mean_reversion = tff.math.piecewise.PiecewiseConstantFunc(
          [0.1, 2.0], values=3 * [self.mean_reversion], dtype=dtype)
      # Volatility with batch dimesnion
      volatility = tff.math.piecewise.PiecewiseConstantFunc(
          [[0.1, 0.2, 0.5], [0.1, 2.0, 3.0]],
          values=[4 * [self.volatility[0]],
                  4 * [self.volatility[1]]], dtype=dtype)
      def corr_matrix(t):
        one = tf.ones_like(t)
        row1 = tf.stack([one, 0.5 * t], axis=-1)
        row2 = tf.reverse(row1, [0])
        corr_matrix = tf.stack([row1, row2], axis=-1)
        return corr_matrix
      process = tff.models.hull_white.VectorHullWhiteModel(
          dim=2,
          mean_reversion=mean_reversion,
          volatility=volatility,
          corr_matrix=corr_matrix,
          initial_discount_rate_fn=self.instant_forward_rate_2d_fn,
          dtype=dtype)
      times = [0.1, 0.5]
      paths = process.sample_paths(
          times,
          num_samples=50000,
          random_type=tff.math.random.RandomType.SOBOL,
          skip=100000,
          time_step=0.01)
      self.assertEqual(paths.dtype, dtype)
      self.assertAllEqual(paths.shape, [50000, 2, 2])
      paths = self.evaluate(paths)
      paths = paths[:, -1, :]  # Extract paths values for the terminal time
      mean = np.mean(paths, axis=0)
      variance = np.var(paths, axis=0)
      self.assertAllClose(mean, self.true_mean(times[-1]), rtol=1e-3, atol=1e-3)
      self.assertAllClose(variance,
                          self.true_var(times[-1]), rtol=1e-3, atol=1e-3) 

def setUp(self):
    self.mean_reversion_1d = [0.03]
    self.volatility_1d = [0.02]
    self.volatility_time_dep_1d = [0.01, 0.02]
    self.mean_reversion_2d = [0.03, 0.06]
    self.volatility_2d = [0.02, 0.01]

    self.expiries = np.array([0.0, 0.25, 0.5, 0.75])
    self.maturities = np.array([0.25, 0.5, 0.75, 1.0])
    self.strikes = 0.01 * np.ones_like(self.expiries)
    self.daycount_fractions = 0.25 * np.ones_like(self.expiries)

    super(HullWhiteCapFloorTest, self).setUp() 

def test_correctness_1d(self, use_analytic_pricing, error_tol):
    """Tests model with constant parameters in 1 dimension."""
    # 1 year cap with quarterly resets.
    dtype = tf.float64

    discount_rate_fn = lambda x: 0.01 * tf.ones_like(x, dtype=dtype)
    price = tff.models.hull_white.cap_floor_price(
        strikes=self.strikes,
        expiries=self.expiries,
        maturities=self.maturities,
        daycount_fractions=self.daycount_fractions,
        notional=100.0,
        dim=1,
        mean_reversion=self.mean_reversion_1d,
        volatility=self.volatility_1d,
        reference_rate_fn=discount_rate_fn,
        use_analytic_pricing=use_analytic_pricing,
        num_samples=500000,
        time_step=0.1,
        random_type=tff.math.random.RandomType.PSEUDO_ANTITHETIC,
        dtype=dtype)
    self.assertEqual(price.dtype, dtype)
    self.assertAllEqual(price.shape, [1, 1])
    price = self.evaluate(price)
    self.assertAllClose(price, [[0.4072088281493774]],
                        rtol=error_tol, atol=error_tol) 

def test_correctness_time_dep_1d(self, use_analytic_pricing, error_tol):
    """Tests model with piecewise constant volatility in 1 dimension."""
    dtype = tf.float64

    discount_rate_fn = lambda x: 0.01 * tf.ones_like(x, dtype=dtype)
    volatility = tff.math.piecewise.PiecewiseConstantFunc(
        jump_locations=[0.5], values=self.volatility_time_dep_1d,
        dtype=dtype)
    price = tff.models.hull_white.cap_floor_price(
        strikes=self.strikes,
        expiries=self.expiries,
        maturities=self.maturities,
        daycount_fractions=self.daycount_fractions,
        notional=100.0,
        dim=1,
        mean_reversion=self.mean_reversion_1d,
        volatility=volatility,
        reference_rate_fn=discount_rate_fn,
        use_analytic_pricing=use_analytic_pricing,
        num_samples=500000,
        time_step=0.1,
        random_type=tff.math.random.RandomType.PSEUDO_ANTITHETIC,
        dtype=dtype)
    self.assertEqual(price.dtype, dtype)
    self.assertAllEqual(price.shape, [1, 1])
    price = self.evaluate(price)
    self.assertAllClose(price, [[0.2394242699989869]], rtol=error_tol,
                        atol=error_tol) 

def test_2d_batch(self, use_analytic_pricing, error_tol):
    """Tests model with 2d batch of options."""
    dtype = tf.float64

    discount_rate_fn = lambda x: 0.01 * tf.ones_like(x, dtype=dtype)
    expiries = np.array([[self.expiries, self.expiries],
                         [self.expiries, self.expiries]])
    maturities = np.array([[self.maturities, self.maturities],
                           [self.maturities, self.maturities]])
    strikes = np.array([[self.strikes, self.strikes],
                        [self.strikes, self.strikes]])
    daycount_fractions = np.array(
        [[self.daycount_fractions, self.daycount_fractions],
         [self.daycount_fractions, self.daycount_fractions]])
    price = tff.models.hull_white.cap_floor_price(
        strikes=strikes,
        expiries=expiries,
        maturities=maturities,
        daycount_fractions=daycount_fractions,
        notional=100.0,
        dim=1,
        mean_reversion=self.mean_reversion_1d,
        volatility=self.volatility_1d,
        reference_rate_fn=discount_rate_fn,
        use_analytic_pricing=use_analytic_pricing,
        num_samples=500000,
        time_step=0.1,
        random_type=tff.math.random.RandomType.PSEUDO_ANTITHETIC,
        dtype=dtype)
    self.assertEqual(price.dtype, dtype)
    self.assertAllEqual(price.shape, [2, 2, 1])
    price = self.evaluate(price)
    expected = [[[0.4072088281493774], [0.4072088281493774]],
                [[0.4072088281493774], [0.4072088281493774]]]
    self.assertAllClose(price, expected, rtol=error_tol, atol=error_tol) 

def test_mixed_1d_batch_2d(self, use_analytic_pricing, error_tol):
    """Tests mixed 1d batch with constant parameters in 2 dimension."""
    dtype = tf.float64

    discount_rate_fn = lambda x: 0.01 * tf.ones_like(x, dtype=dtype)
    expiries = np.array([self.expiries, self.expiries, self.expiries])
    maturities = np.array([self.maturities, self.maturities, self.maturities])
    strikes = np.array([self.strikes, self.strikes, self.strikes])
    daycount_fractions = np.array([
        self.daycount_fractions, self.daycount_fractions,
        self.daycount_fractions
    ])
    price = tff.models.hull_white.cap_floor_price(
        strikes=strikes,
        expiries=expiries,
        maturities=maturities,
        daycount_fractions=daycount_fractions,
        notional=100.0,
        dim=2,
        mean_reversion=self.mean_reversion_2d,
        volatility=self.volatility_2d,
        reference_rate_fn=discount_rate_fn,
        use_analytic_pricing=use_analytic_pricing,
        num_samples=500000,
        time_step=0.1,
        random_type=tff.math.random.RandomType.PSEUDO_ANTITHETIC,
        dtype=dtype)
    self.assertEqual(price.dtype, dtype)
    self.assertAllEqual(price.shape, [3, 2])
    price = self.evaluate(price)
    expected = [[0.4072088281493774, 0.2016075430673558],
                [0.4072088281493774, 0.2016075430673558],
                [0.4072088281493774, 0.2016075430673558]]
    self.assertAllClose(price, expected, rtol=error_tol, atol=error_tol) 

def test_call_put(self, use_analytic_pricing, error_tol):
    """Tests mixed 1d batch with constant parameters in 2 dimension."""
    dtype = tf.float64

    discount_rate_fn = lambda x: 0.01 * tf.ones_like(x, dtype=dtype)
    expiries = np.array([self.expiries, self.expiries, self.expiries])
    maturities = np.array([self.maturities, self.maturities, self.maturities])
    strikes = np.array(
        [self.strikes - 0.005, self.strikes - 0.005, self.strikes - 0.005])
    daycount_fractions = np.array([
        self.daycount_fractions, self.daycount_fractions,
        self.daycount_fractions
    ])
    price = tff.models.hull_white.cap_floor_price(
        strikes=strikes,
        expiries=expiries,
        maturities=maturities,
        daycount_fractions=daycount_fractions,
        notional=100.0,
        dim=2,
        is_cap=[[True], [False], [False]],
        mean_reversion=self.mean_reversion_2d,
        volatility=self.volatility_2d,
        reference_rate_fn=discount_rate_fn,
        use_analytic_pricing=use_analytic_pricing,
        num_samples=500000,
        time_step=0.1,
        random_type=tff.math.random.RandomType.PSEUDO_ANTITHETIC,
        dtype=dtype)
    self.assertEqual(price.dtype, dtype)
    self.assertAllEqual(price.shape, [3, 2])
    price = self.evaluate(price)
    expected = [[0.7460080352561735, 0.5670824151746642],
                [0.2478780964427645, 0.0689524763612857],
                [0.2478780964427645, 0.0689524763612857]]
    self.assertAllClose(price, expected, rtol=error_tol, atol=error_tol) 

def test_gradients_are_straight_through(self):
    noisy = uniform_noise.NoisyNormal(loc=0, scale=1)
    em = ContinuousBatchedEntropyModel(noisy, 1)
    x = tf.range(-20., 20.)
    x_perturbed = x + tf.random.uniform(x.shape, -.49, .49)
    with tf.GradientTape() as tape:
      tape.watch(x_perturbed)
      x_quantized = em.quantize(x_perturbed)
    gradients = tape.gradient(x_quantized, x_perturbed)
    self.assertAllEqual(gradients, tf.ones_like(gradients)) 

def test_sample_paths_2d(self):
    """Tests path properties for 2-dimentional Ito process.

    We construct the following Ito processes.

    dX_1 = mu_1 sqrt(t) dt + s11 dW_1 + s12 dW_2
    dX_2 = mu_2 sqrt(t) dt + s21 dW_1 + s22 dW_2

    mu_1, mu_2 are constants.
    s_ij = a_ij t + b_ij

    For this process expected value at time t is (x_0)_i + 2/3 * mu_i * t^1.5.
    """
    mu = np.array([0.2, 0.7])
    a = np.array([[0.4, 0.1], [0.3, 0.2]])
    b = np.array([[0.33, -0.03], [0.21, 0.5]])

    def drift_fn(t, x):
      return mu * tf.sqrt(t) * tf.ones_like(x, dtype=t.dtype)

    def vol_fn(t, x):
      del x
      return (a * t + b) * tf.ones([2, 2], dtype=t.dtype)

    num_samples = 10000
    process = GenericItoProcess(dim=2, drift_fn=drift_fn, volatility_fn=vol_fn)
    times = np.array([0.1, 0.21, 0.32, 0.43, 0.55])
    x0 = np.array([0.1, -1.1])
    paths = self.evaluate(
        process.sample_paths(
            times,
            num_samples=num_samples,
            initial_state=x0,
            time_step=0.01,
            seed=12134))

    self.assertAllClose(paths.shape, (num_samples, 5, 2), atol=0)
    means = np.mean(paths, axis=0)
    times = np.reshape(times, [-1, 1])
    expected_means = x0 + (2.0 / 3.0) * mu * np.power(times, 1.5)
    self.assertAllClose(means, expected_means, rtol=1e-2, atol=1e-2) 

def __call__(self, x):
    non_sign_bits = self.bits - (self.negative_slope != 0)
    m = K.cast_to_floatx(pow(2, non_sign_bits))
    m_i = K.cast_to_floatx(pow(2, self.integer))
    x_uq = tf.where(
        x <= m_i, K.relu(x, alpha=self.negative_slope), tf.ones_like(x) * m_i)

    if self.use_sigmoid:
      p = _sigmoid(x / m_i) * m
      xq = m_i * tf.keras.backend.clip(
          2.0 * (_round_through(p, self.use_stochastic_rounding) / m) - 1.0,
          0.0, 1.0 - 1.0 / m)
      if self.negative_slope > 0:
        neg_factor = 1 / (self.negative_slope * m)
        xq = xq + m_i * self.negative_slope * tf.keras.backend.clip(
            2.0 * (_round_through(p * self.negative_slope,
            self.use_stochastic_rounding) * neg_factor) - 1.0,
            -1.0, 0.0)
    else:
      p = x * m / m_i
      xq = m_i * tf.keras.backend.clip(
          _round_through(p, self.use_stochastic_rounding) / m, 0.0,
          1.0 - 1.0 / m)
      if self.negative_slope > 0:
        neg_factor = 1 / (self.negative_slope * m)
        xq = xq + m_i * self.negative_slope * (tf.keras.backend.clip(
            _round_through(p * self.negative_slope,
                           self.use_stochastic_rounding) * neg_factor, -1.0, 0.0))
    return x_uq + tf.stop_gradient(-x_uq + xq) 

def __call__(self, x):
    if self.max_value is None:
      x = K.relu(x)
    else:
      x = tf.where(
          x <= self.max_value, K.relu(x), tf.ones_like(x) * self.max_value)

    x_clipped = _clip_power_of_two(x, self._min_exp, self._max_exp,
                                   self.max_value,
                                   self.quadratic_approximation,
                                   self.use_stochastic_rounding)
    return x + tf.stop_gradient(-x + pow(2.0, x_clipped)) 

def test_option_prices_neg_carries(self,
                                     discount_rates,
                                     volatilities,
                                     expiries,
                                     expected_prices):
    """Tests the prices for negative cost_of_carries."""
    spots = np.array([80.0, 90.0, 100.0, 110.0, 120.0] * 2)
    strikes = np.array([100.0] * 10)
    is_call_options = np.array([True] * 5 + [False] * 5)
    cost_of_carries = -0.04
    computed_prices, converged, failed = adesi_whaley(
        volatilities=volatilities,
        strikes=strikes,
        expiries=expiries,
        discount_rates=discount_rates,
        cost_of_carries=cost_of_carries,
        is_call_options=is_call_options,
        spots=spots,
        dtype=tf.float64)
    expected_prices = np.array(expected_prices)
    with self.subTest(name='ExpectedPrices'):
      self.assertAllClose(expected_prices, computed_prices,
                          rtol=5e-3, atol=5e-3)
    with self.subTest(name='AllConverged'):
      self.assertAllEqual(converged, tf.ones_like(computed_prices))
    with self.subTest(name='NonFailed'):
      self.assertAllEqual(failed, tf.zeros_like(computed_prices)) 

def test_option_prices_pos_carries(self,
                                     discount_rates,
                                     volatilities,
                                     expiries,
                                     expected_prices):
    """Tests the prices for positive cost_of_carries."""
    spots = np.array([80.0, 90.0, 100.0, 110.0, 120.0] * 2)
    strikes = np.array([100.0] * 10)
    is_call_options = [True] * 5 + [False] * 5
    cost_of_carries = 0.04
    computed_prices, converged, failed = adesi_whaley(
        volatilities=volatilities,
        strikes=strikes,
        expiries=expiries,
        discount_rates=discount_rates,
        cost_of_carries=cost_of_carries,
        spots=spots,
        is_call_options=is_call_options,
        dtype=tf.float64)
    with self.subTest(name='ExpectedPrices'):
      self.assertAllClose(expected_prices, computed_prices,
                          rtol=5e-3, atol=5e-3)
    with self.subTest(name='AllConverged'):
      self.assertAllEqual(converged, tf.ones_like(computed_prices))
    with self.subTest(name='NonFailed'):
      self.assertAllEqual(failed, tf.zeros_like(computed_prices)) 

def test_option_prices_zero_cost_of_carries(self,
                                              discount_rates,
                                              volatilities,
                                              expiries,
                                              expected_prices):
    """Tests the prices when cost_of_carries is zero."""
    forwards = np.array([80.0, 90.0, 100.0, 110.0, 120.0] * 2)
    strikes = np.array([100.0] * 10)
    is_call_options = [True] * 5 + [False] * 5
    cost_of_carries = 0.
    computed_prices, converged, failed = adesi_whaley(
        volatilities=volatilities,
        strikes=strikes,
        expiries=expiries,
        discount_rates=discount_rates,
        cost_of_carries=cost_of_carries,
        forwards=forwards,
        is_call_options=is_call_options,
        dtype=tf.float64)
    with self.subTest(name='ExpectedPrices'):
      self.assertAllClose(expected_prices, computed_prices,
                          rtol=5e-3, atol=5e-3)
    with self.subTest(name='AllConverged'):
      self.assertAllEqual(converged, tf.ones_like(computed_prices))
    with self.subTest(name='NonFailed'):
      self.assertAllEqual(failed, tf.zeros_like(computed_prices))