Python tensorflow.compat.v2.concat() Examples

def _torso(self, observation):
    conv_out = observation[streetview_constants.IMAGE_FEATURES]
    heading = observation[streetview_constants.HEADING]
    last_action = observation[streetview_constants.PREV_ACTION_IDX]

    conv_out = tf.cast(conv_out, tf.float32)

    img_encoding = self._dense_img_extra(self._dense_img(conv_out))
    img_encoding = tf.keras.layers.Flatten()(img_encoding)

    heading = tf.expand_dims(heading, -1)
    last_action_embedded = self._action_embedder(last_action)

    torso_output = tf.concat([heading, last_action_embedded, img_encoding],
                             axis=1)
    timestep_embedded = self._timestep_embedder(
        observation[streetview_constants.TIMESTEP])
    return {
        'neck_input': torso_output,
        streetview_constants.TIMESTEP: timestep_embedded,
    } 

def hessian_as_matrix(function: Callable[[Parameters], tf.Tensor],
                      parameters: Parameters) -> tf.Tensor:
  """Computes the Hessian of a given function.

  Same as `hessian`, although return a matrix of size [w_dim, w_dim], where
  `w_dim` is the number of parameters, which makes it easier to work with.

  Args:
    function: A function for which we want to compute the Hessian.
    parameters: Parameters with respect to the Hessian should be computed.

  Returns:
    A tensor of size [w_dim, w_dim] representing the Hessian.
  """
  hessian_as_tensor_list = hessian(function, parameters)
  hessian_as_tensor_list = [
      tf.reshape(e, [e.shape[0], -1]) for e in hessian_as_tensor_list]
  return tf.concat(hessian_as_tensor_list, axis=1) 

def generate_random_test_dual_quaternions():
  """Generates random test dual quaternions."""
  angles = generate_random_test_euler_angles()
  random_quaternion_real = quaternion.from_euler(angles)

  min_translation = -3.0
  max_translation = 3.0
  translations = np.random.uniform(min_translation, max_translation,
                                   angles.shape)

  translations_quaternion_shape = np.asarray(translations.shape)
  translations_quaternion_shape[-1] = 1
  translations = np.concatenate(
      (translations / 2.0, np.zeros(translations_quaternion_shape)), axis=-1)

  random_quaternion_translation = tf.convert_to_tensor(value=translations)

  random_quaternion_dual = quaternion.multiply(random_quaternion_translation,
                                               random_quaternion_real)

  random_dual_quaternion = tf.concat(
      (random_quaternion_real, random_quaternion_dual), axis=-1)

  return random_dual_quaternion 

def test_conjugate_preset(self):
    """Tests if the conjugate function is providing correct results."""
    x_init = test_helpers.generate_preset_test_dual_quaternions()
    x = tf.convert_to_tensor(value=x_init)
    y = tf.convert_to_tensor(value=x_init)

    x = dual_quaternion.conjugate(x)
    x_real, x_dual = tf.split(x, (4, 4), axis=-1)

    y_real, y_dual = tf.split(y, (4, 4), axis=-1)
    xyz_y_real, w_y_real = tf.split(y_real, (3, 1), axis=-1)
    xyz_y_dual, w_y_dual = tf.split(y_dual, (3, 1), axis=-1)
    y_real = tf.concat((-xyz_y_real, w_y_real), axis=-1)
    y_dual = tf.concat((-xyz_y_dual, w_y_dual), axis=-1)

    self.assertAllEqual(x_real, y_real)
    self.assertAllEqual(x_dual, y_dual) 

def generate_preset_test_dual_quaternions():
  """Generates pre-set test quaternions."""
  angles = generate_preset_test_euler_angles()
  preset_quaternion_real = quaternion.from_euler(angles)

  translations = generate_preset_test_translations()
  translations = np.concatenate(
      (translations / 2.0, np.zeros((np.ma.size(translations, 0), 1))), axis=1)
  preset_quaternion_translation = tf.convert_to_tensor(value=translations)

  preset_quaternion_dual = quaternion.multiply(preset_quaternion_translation,
                                               preset_quaternion_real)

  preset_dual_quaternion = tf.concat(
      (preset_quaternion_real, preset_quaternion_dual), axis=-1)

  return preset_dual_quaternion 

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 _exact_discretization_setup(self, dim):
    """Initial setup for efficient computations."""
    self._zero_padding = tf.zeros((dim, 1), dtype=self._dtype)
    self._jump_locations = tf.concat(
        [self._volatility.jump_locations(),
         self._mean_reversion.jump_locations()], axis=-1)
    self._jump_values_vol = self._volatility(self._jump_locations)
    self._jump_values_mr = self._mean_reversion(self._jump_locations)
    if dim == 1:
      self._padded_knots = tf.concat([
          self._zero_padding,
          tf.expand_dims(self._jump_locations[:-1], axis=0)
      ], axis=1)
      self._jump_values_vol = tf.expand_dims(self._jump_values_vol, axis=0)
      self._jump_values_mr = tf.expand_dims(self._jump_values_mr, axis=0)
      self._jump_locations = tf.expand_dims(self._jump_locations, axis=0)

    else:
      self._padded_knots = tf.concat(
          [self._zero_padding, self._jump_locations[:, :-1]], axis=1) 

def _compute_yt(self, t, mr_t, sigma_t):
    """Computes y(t) as described in [1], section 10.1.6.1."""
    t = tf.repeat(tf.expand_dims(t, axis=0), self._dim, axis=0)
    time_index = tf.searchsorted(self._jump_locations, t)
    y_between_vol_knots = self._y_integral(
        self._padded_knots, self._jump_locations, self._jump_values_vol,
        self._jump_values_mr)
    y_at_vol_knots = tf.concat(
        [self._zero_padding,
         _cumsum_using_matvec(y_between_vol_knots)], axis=1)

    vn = tf.concat(
        [self._zero_padding, self._jump_locations], axis=1)
    y_t = self._y_integral(
        tf.gather(vn, time_index, batch_dims=1), t, sigma_t, mr_t)
    y_t = y_t + tf.gather(y_at_vol_knots, time_index, batch_dims=1)
    return tf.math.exp(-2 * mr_t * t) * y_t 

def _conditional_variance_x(self, t, mr_t, sigma_t):
    """Computes the variance of x(t), see [1], Eq. 10.41."""
    t = tf.repeat(tf.expand_dims(t, axis=0), self._dim, axis=0)
    var_x_between_vol_knots = self._variance_int(self._padded_knots,
                                                 self._jump_locations,
                                                 self._jump_values_vol,
                                                 self._jump_values_mr)
    varx_at_vol_knots = tf.concat(
        [self._zero_padding,
         _cumsum_using_matvec(var_x_between_vol_knots)],
        axis=1)

    time_index = tf.searchsorted(self._jump_locations, t)
    vn = tf.concat(
        [self._zero_padding,
         self._jump_locations], axis=1)

    var_x_t = self._variance_int(
        tf.gather(vn, time_index, batch_dims=1), t, sigma_t, mr_t)
    var_x_t = var_x_t + tf.gather(varx_at_vol_knots, time_index, batch_dims=1)

    var_x_t = (var_x_t[:, 1:] - var_x_t[:, :-1]) * tf.math.exp(
        -2 * tf.broadcast_to(mr_t, t.shape)[:, 1:] * t[:, 1:])
    return var_x_t 

def test_interpolation_differentiable(self):
    dtype = tf.float64
    interval_times = tf.constant([0.25, 0.5, 1.0, 2.0, 3.0], dtype=dtype)
    knot_1y = tf.constant([0.052], dtype=dtype)
    interval_values = tf.concat([
        tf.constant([0.05, 0.051], dtype=dtype), knot_1y,
        tf.constant([0.053, 0.055], dtype=dtype)
    ],
                                axis=0)
    test_time = tf.constant([1.1, 2.7], dtype=dtype)
    interpolated, _ = monotone_convex.interpolate(test_time, interval_values,
                                                  interval_times)
    gradient_1y = self.evaluate(tf.convert_to_tensor(
        tf.gradients(interpolated[0], knot_1y)[0]))
    gradient_zero = self.evaluate(tf.convert_to_tensor(
        tf.gradients(interpolated[1], knot_1y)[0]))

    self.assertAlmostEqual(gradient_1y[0], 0.42)
    self.assertAlmostEqual(gradient_zero[0], 0.0) 

def test_interpolated_yields_consistency(self):
    dtypes = [np.float32, np.float64]
    for dtype in dtypes:
      reference_times = np.array([1.0, 2.0, 3.0, 4.0], dtype=dtype)
      yields = np.array([5.0, 4.75, 4.53333333, 4.775], dtype=dtype)

      # Times for which the interpolated values are required.
      interpolation_times_1 = tf.constant([0.25, 0.5, 1.0, 2.0, 3.0],
                                          dtype=dtype)
      interpolation_times_2 = tf.constant([1.1, 2.5, 2.9, 3.6, 4.0],
                                          dtype=dtype)
      expected = np.array([
          5.1171875, 5.09375, 5.0, 4.75, 4.533333, 4.9746, 4.624082, 4.535422,
          4.661777, 4.775
      ],
                          dtype=dtype)
      actual_1 = monotone_convex.interpolate_yields(
          interpolation_times_1, reference_times, yields=yields)
      actual_2 = monotone_convex.interpolate_yields(
          interpolation_times_2, reference_times, yields=yields)
      actual = self.evaluate(tf.concat([actual_1, actual_2], axis=0))
      np.testing.assert_allclose(actual, expected, rtol=1e-5) 

def _initial_discount_rates(bond_cashflows,
                            bond_cashflow_times,
                            present_values,
                            name='initial_discount_rates'):
  """Constructs a guess for the initial rates as the yields to maturity."""
  n = len(bond_cashflows)
  groups = []
  for i in range(n):
    groups.append(tf.fill(tf.shape(bond_cashflows[i]), i))
  bond_cashflows = tf.concat(bond_cashflows, axis=0)
  bond_cashflow_times = tf.concat(bond_cashflow_times, axis=0)
  groups = tf.concat(groups, axis=0)
  return cashflows.yields_from_pv(
      bond_cashflows,
      bond_cashflow_times,
      present_values,
      groups=groups,
      name=name) 

def _apply_sequence_to_tensor_op(cls, op_fn, tensor_wrappers):
    """Applies given sequence-to-tensor op.

    This method is used for implementing ops that take a sequence of tensors and
    return a new tensor, such as tf.concat and tf.stack. Implementing wrappers
    should apply `op_fn` to the backing tensor(s) and return an new wrapper
    instance with the combined backing tensor.

    Args:
     op_fn: Callable that applies sequence-to-tensor op to the given sequence
       of Tensors. E.g. applies tf.concat.
     tensor_wrappers: a sequence of tensor wrappers to be transformed. All
       elements have the type of the implementing TensorWrapper class.

    Returns:
      A TensorWrapper instance with combined backing tensor(s).
    """
    raise NotImplementedError() 

def _reshape_inner_dims(
    tensor: tf.Tensor,
    shape: tf.TensorShape,
    new_shape: tf.TensorShape) -> tf.Tensor:
  """Reshapes tensor to: shape(tensor)[:-len(shape)] + new_shape."""
  tensor_shape = tf.shape(tensor)
  ndims = shape.rank
  tensor.shape[-ndims:].assert_is_compatible_with(shape)
  new_shape_inner_tensor = tf.cast(
      [-1 if d is None else d for d in new_shape.as_list()], tf.int64)
  new_shape_outer_tensor = tf.cast(
      tensor_shape[:-ndims], tf.int64)
  full_new_shape = tf.concat(
      (new_shape_outer_tensor, new_shape_inner_tensor), axis=0)
  new_tensor = tf.reshape(tensor, full_new_shape)
  new_tensor.set_shape(tensor.shape[:-ndims] + new_shape)
  return new_tensor 

def _get_bi_lstm_encoder(self, num_hidden_layers, hidden_dim):
    """Get Bi-LSTM encoder.

    Args:
      num_hidden_layers: Number of stacked layers.
      hidden_dim: The hidden size of LSTM.

    Returns:
      A list of 2N+1 elements. The first element is output of all timesteps
        and the others are LSTM state. The first N are forward states and the
        last N are backward states.
    """
    self._cells = []
    for layer_id in range(num_hidden_layers):

      self._cells.append(
          tf.keras.layers.LSTMCell(
              hidden_dim, name='lstm_layer_{}'.format(layer_id)))

    self._cells_rnn = tf.keras.layers.RNN(
        self._cells, return_sequences=True, return_state=True)
    return tf.keras.layers.Bidirectional(self._cells_rnn, merge_mode='concat') 

def _get_initial_state(self, observation, batch_size):
    text_token_ids = observation[constants.INS_TOKEN_IDS]
    text_enc_outputs, final_state = self._instruction_encoder(text_token_ids)
    if self._ndh_instruction_encoder is not None:
      ndh_text_enc_outputs, ndh_final_state = self._ndh_instruction_encoder(
          text_token_ids)
      mask = tf.equal(observation[constants.PROBLEM_TYPE],
                      constants.PROBLEM_VLN)
      text_enc_outputs = tf.nest.map_structure(
          lambda x, y: tf.compat.v1.where(mask, x, y), text_enc_outputs,
          ndh_text_enc_outputs)
      final_state = tf.nest.map_structure(
          lambda x, y: tf.compat.v1.where(mask, x, y), final_state,
          ndh_final_state)

    if self._ins_classifier is not None:
      # Concatenate all hidden layers' state vectors. Use state.h
      ins_classifier_logits = self._ins_classifier(
          tf.concat([s[0] for s in final_state], axis=1))
    else:
      ins_classifier_logits = tf.zeros(shape=(batch_size, 2))
    return (final_state, text_enc_outputs, ins_classifier_logits) 

def conjugate(dual_quaternion, name=None):
  """Computes the conjugate of a dual quaternion.

  Note:
    In the following, A1 to An are optional batch dimensions.

  Args:
    dual_quaternion: A tensor of shape `[A1, ..., An, 8]`, where the last
    dimension represents a normalized dual quaternion.
    name: A name for this op that defaults to "dual_quaternion_conjugate".

  Returns:
    A tensor of shape `[A1, ..., An, 8]`, where the last dimension represents
    a normalized dual quaternion.

  Raises:
    ValueError: If the shape of `dual_quaternion` is not supported.
  """
  with tf.compat.v1.name_scope(name, "dual_quaternion_conjugate",
                               [dual_quaternion]):
    dual_quaternion = tf.convert_to_tensor(value=dual_quaternion)

    shape.check_static(
        tensor=dual_quaternion, tensor_name="dual_quaternion",
        has_dim_equals=(-1, 8))

    quaternion_real, quaternion_dual = tf.split(
        dual_quaternion, (4, 4), axis=-1)

    quaternion_real = asserts.assert_normalized(quaternion_real)

    return tf.concat((quaternion.conjugate(quaternion_real),
                      quaternion.conjugate(quaternion_dual)),
                     axis=-1) 

def while_with_variable_shape_growing_matrix_rows(n):
  m = tf.constant([[0]])
  i = 0
  while i < n:
    tf.autograph.experimental.set_loop_options(
        shape_invariants=[(m, tf.TensorShape([None, 1]))])
    m = tf.concat((m, [[i]]), 0)
    i += 1
  return m 

def for_with_variable_shape_growing_vector(l):
  v = tf.constant([0, 0])
  for i in l:
    tf.autograph.experimental.set_loop_options(
        shape_invariants=[(v, tf.TensorShape([None]))])
    v = tf.concat((v, [i]), 0)
  return v 

def call(self, inputs, lengths, training):
    seq_mask = tf.sequence_mask(
        lengths, inputs.shape[1], dtype=tf.dtypes.float32)
    forward_outputs = self.forward_rnn(inputs, training=training)
    reversed_inputs = tf.reverse_sequence(inputs, lengths, seq_axis=1)
    backward_outputs = self.backward_rnn(reversed_inputs, training=training)
    backward_outputs = tf.reverse_sequence(
        backward_outputs, lengths, seq_axis=1)
    outputs = tf.concat([forward_outputs, backward_outputs], axis=-1)
    outputs = outputs * tf.expand_dims(seq_mask, -1)

    if self.reduce_states:
      outputs = tf.reduce_mean(outputs, axis=1)

    return outputs 

def while_with_variable_shape_growing_vector(n):
  v = tf.constant([0, 0])
  i = 0
  while i < n:
    tf.autograph.experimental.set_loop_options(
        shape_invariants=[(v, tf.TensorShape([None]))])
    v = tf.concat((v, [i]), 0)
    i += 1
  return v 

def _compute_is_bus_day_table(self):
    """Computes and caches "is business day" table."""
    if self._table_cache.is_bus_day is not None:
      return self._table_cache.is_bus_day

    with tf.init_scope():
      ordinals = tf.range(self._ordinal_offset,
                          self._ordinal_offset + self._calendar_size)
      # Apply weekend mask
      week_days = (ordinals - 1) % 7
      is_holiday = tf.gather(self._weekend_mask, week_days)

      # Apply holidays
      if self._holidays is not None:
        indices = self._holidays.ordinal() - self._ordinal_offset
        ones_at_indices = tf.scatter_nd(
            tf.expand_dims(indices, axis=-1), tf.ones_like(indices),
            is_holiday.shape)
        is_holiday = tf.bitwise.bitwise_or(is_holiday, ones_at_indices)

      # Add a business day at the beginning and at the end, i.e. at 31 Dec of
      # start_year-1 and at 1 Jan of end_year+1. This trick is to avoid dealing
      # with special cases on boundaries.
      # For example, for Following and Preceding conventions we'd need a special
      # value that means "unknown" in the tables. More complicated conventions
      # then combine the Following and Preceding tables, and would need special
      # treatment of the "unknown" values.
      # With these "fake" business days, all computations are automatically
      # correct, unless we land on those extra days - for this reason we add
      # assertions in all API calls before returning.
      is_bus_day_table = tf.concat([[1], 1 - is_holiday, [1]], axis=0)
      self._table_cache.is_bus_day = is_bus_day_table
    return is_bus_day_table 

def diagonal(a, offset=0, axis1=0, axis2=1):  # pylint: disable=missing-docstring
  a = asarray(a).data

  maybe_rank = a.shape.rank
  if maybe_rank is not None and offset == 0 and (
      axis1 == maybe_rank - 2 or axis1 == -2) and (axis2 == maybe_rank - 1 or
                                                   axis2 == -1):
    return utils.tensor_to_ndarray(tf.linalg.diag_part(a))

  a = moveaxis(utils.tensor_to_ndarray(a), (axis1, axis2), (-2, -1)).data

  a_shape = tf.shape(a)

  def _zeros():  # pylint: disable=missing-docstring
    return (tf.zeros(tf.concat([a_shape[:-1], [0]], 0), dtype=a.dtype), 0)

  # All zeros since diag_part doesn't handle all possible k (aka offset).
  # Written this way since cond will run shape inference on both branches,
  # and diag_part shape inference will fail when offset is out of bounds.
  a, offset = utils.cond(
      utils.logical_or(
          utils.less_equal(offset, -1 * utils.getitem(a_shape, -2)),
          utils.greater_equal(offset, utils.getitem(a_shape, -1)),
      ), _zeros, lambda: (a, offset))

  a = utils.tensor_to_ndarray(tf.linalg.diag_part(a, k=offset))
  return a 

def concat(cls, tensor_wrappers, axis):
    """See tf.concat."""
    cls._validate_tensor_types(tensor_wrappers, "concat")
    return cls._apply_sequence_to_tensor_op(
        lambda ts: tf.concat(ts, axis), tensor_wrappers) 

def _interpolate_adjacent(times, values, name=None):
  """Interpolates linearly between adjacent values.

  Suppose `times` are `[t_1, t_2, ..., t_n]` an array of length `n` and
  values are `[f_1, ... f_n]` of length `n`. This function associates
  each of the values to the midpoint of the interval i.e. `f_i` is associated
  to the midpoint of the interval `[t_i, t_{i+1}]`. Then it calculates the
  values at the interval boundaries by linearly interpolating between adjacent
  intervals. The first interval is considered to be `[0, t_1]`. The values at
  the endpoints (i.e. result[0] and result[n]) are computed as follows:
  `result[0] = values[0] - 0.5 * (result[1] - values[0])` and
  `result[n] = values[n-1] - 0.5 * (result[n-1] - values[n-1])`.
  The rationale for these specific values is discussed in Ref. [1].

  Args:
    times: A rank 1 `Tensor` of real dtype. The times at which the interpolated
      values are to be computed. The values in the array should be positive and
      monotonically increasing.
    values: A rank 1 `Tensor` of the same dtype and shape as `times`. The values
      assigned to the midpoints of the time intervals.
    name: Python `str` name prefixed to Ops created by this class.
      Default value: None which is mapped to the default name
        'interpolate_adjacent'.

  Returns:
    interval_values: The values interpolated from the supplied midpoint values
      as described above. A `Tensor` of the same dtype as `values` but shape
      `[n+1]` where `[n]` is the shape of `values`. The `i`th component of the
      is the value associated to the time point `t_{i+1}` with `t_0 = 0`.
  """
  with tf.compat.v1.name_scope(
      name, default_name='interpolate_adjacent', values=[times, values]):
    dt1 = diff(times, order=1, exclusive=False)
    dt2 = diff(times, order=2, exclusive=False)[1:]
    weight_right = dt1[:-1] / dt2
    weight_left = dt1[1:] / dt2
    interior_values = weight_right * values[1:] + weight_left * values[:-1]
    value_0 = values[0] - 0.5 * (interior_values[0] - values[0])
    value_n = values[-1] - 0.5 * (interior_values[-1] - values[-1])
    return tf.concat([[value_0], interior_values, [value_n]], axis=0) 

def _create_curve_building_tensors(float_leg_start_times,
                                   float_leg_end_times,
                                   fixed_leg_end_times,
                                   pv_settlement_times):
  """Helper function to create tensors needed for curve construction."""
  calc_groups_float = []
  calc_groups_fixed = []
  expiry_times = []
  settle_times_float = []
  settle_times_fixed = []
  num_instruments = len(float_leg_start_times)
  for i in range(num_instruments):
    expiry_times.append(
        tf.math.maximum(float_leg_end_times[i][-1], fixed_leg_end_times[i][-1]))

    calc_groups_float.append(
        tf.fill(tf.shape(float_leg_start_times[i]), i))
    calc_groups_fixed.append(tf.fill(tf.shape(fixed_leg_end_times[i]), i))
    settle_times_float.append(tf.fill(tf.shape(float_leg_start_times[i]),
                                      pv_settlement_times[i]))
    settle_times_fixed.append(tf.fill(tf.shape(fixed_leg_end_times[i]),
                                      pv_settlement_times[i]))

  expiry_times = tf.stack(expiry_times, axis=0)
  calc_groups_float = tf.concat(calc_groups_float, axis=0)
  calc_groups_fixed = tf.concat(calc_groups_fixed, axis=0)
  settle_times_float = tf.concat(settle_times_float, axis=0)
  settle_times_fixed = tf.concat(settle_times_fixed, axis=0)

  return CurveFittingVars(expiry_times=expiry_times,
                          calc_groups_float=calc_groups_float,
                          calc_groups_fixed=calc_groups_fixed,
                          settle_times_float=settle_times_float,
                          settle_times_fixed=settle_times_fixed) 

def _conditional_mean_x(self, t, mr_t, sigma_t):
    """Computes the drift term in [1], Eq. 10.39."""
    t = tf.repeat(tf.expand_dims(t, axis=0), self._dim, axis=0)
    time_index = tf.searchsorted(self._jump_locations, t)
    vn = tf.concat([self._zero_padding, self._jump_locations], axis=1)
    y_between_vol_knots = self._y_integral(self._padded_knots,
                                           self._jump_locations,
                                           self._jump_values_vol,
                                           self._jump_values_mr)

    y_at_vol_knots = tf.concat(
        [self._zero_padding,
         _cumsum_using_matvec(y_between_vol_knots)], axis=1)

    ex_between_vol_knots = self._ex_integral(self._padded_knots,
                                             self._jump_locations,
                                             self._jump_values_vol,
                                             self._jump_values_mr,
                                             y_at_vol_knots[:, :-1])

    ex_at_vol_knots = tf.concat(
        [self._zero_padding,
         _cumsum_using_matvec(ex_between_vol_knots)], axis=1)

    c = tf.gather(y_at_vol_knots, time_index, batch_dims=1)
    exp_x_t = self._ex_integral(
        tf.gather(vn, time_index, batch_dims=1), t, sigma_t, mr_t, c)
    exp_x_t = exp_x_t + tf.gather(ex_at_vol_knots, time_index, batch_dims=1)
    exp_x_t = (exp_x_t[:, 1:] - exp_x_t[:, :-1]) * tf.math.exp(
        -tf.broadcast_to(mr_t, t.shape)[:, 1:] * t[:, 1:])
    return exp_x_t 

def dstack(tup):
  arrays = [atleast_3d(a) for a in tup]
  arrays = _promote_dtype(*arrays)  # pylint: disable=protected-access
  unwrapped_arrays = [
      a.data if isinstance(a, arrays_lib.ndarray) else a for a in arrays
  ]
  return tf.concat(unwrapped_arrays, axis=2) 

def vstack(tup):
  arrays = [atleast_2d(a) for a in tup]
  arrays = _promote_dtype(*arrays)  # pylint: disable=protected-access
  unwrapped_arrays = [
      a.data if isinstance(a, arrays_lib.ndarray) else a for a in arrays
  ]
  return tf.concat(unwrapped_arrays, axis=0) 

def _prepare_grid(*, times, time_step, dtype):
  """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.
    time_step: Rank 0 real `Tensor`. Maximal distance between points in
      resulting grid.
    dtype: `tf.Dtype` of the input and output `Tensor`s.

  Returns:
    Tuple `(all_times, mask, time_points)`.
    `all_times` is a 1-D real `Tensor` containing all points from 'times` and
    the uniform grid of points between `[0, times[-1]]` with grid size equal to
    `time_step`. 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.
    `time_indices`. An integer `Tensor` of the same shape as `times` indicating
    `times` indices in `all_times`.
  """
  grid = tf.range(0.0, times[-1], time_step, dtype=dtype)
  all_times = tf.concat([grid, times], axis=0)
  mask = tf.concat([
      tf.zeros_like(grid, dtype=tf.bool),
      tf.ones_like(times, dtype=tf.bool)
  ],
                   axis=0)
  perm = tf.argsort(all_times, stable=True)
  all_times = tf.gather(all_times, perm)
  # Remove duplicate points
  all_times = tf.unique(all_times).y
  time_indices = tf.searchsorted(all_times, times)
  mask = tf.gather(mask, perm)
  return all_times, mask, time_indices