Python tensorflow.compat.v2.squeeze() Examples

def squeeze(a, axis=None):
  """Removes single-element axes from the 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`.
    axis: scalar or list/tuple of ints.

  TODO(srbs): tf.squeeze throws error when axis is a Tensor eager execution
  is enabled. So we cannot allow axis to be array_like here. Fix.

  Returns:
    An ndarray.
  """
  a = asarray(a)
  return utils.tensor_to_ndarray(tf.squeeze(a, axis)) 

def test_univariate_sample_mean_and_variance(self, dtype):
    """Tests the mean and vol of the univariate GBM sampled paths."""
    process = tff.models.GeometricBrownianMotion(0.05, 0.5, dtype=dtype)
    samples = process.sample_paths(
        times=[0.1, 0.5, 1.0], initial_state=2.0,
        random_type=tff.math.random.RandomType.SOBOL, num_samples=10000)
    log_s = tf.math.log(samples)
    mean = tf.reduce_mean(log_s, axis=0, keepdims=True)
    var = tf.reduce_mean((log_s - mean)**2, axis=0)
    expected_mean = ((process._mu - process._sigma**2 / 2)
                     * np.array([0.1, 0.5, 1.0]) + np.log(2.))
    expected_var = process._sigma**2 * np.array([0.1, 0.5, 1.0])
    with self.subTest("Drift"):
      self.assertAllClose(tf.squeeze(mean), expected_mean, atol=1e-3, rtol=1e-3)
    with self.subTest("Variance"):
      self.assertAllClose(tf.squeeze(var), expected_var, atol=1e-3, rtol=1e-3) 

def _get_parameters(times, *params):
  """Gets parameter values at at specified `times`."""
  res = []
  for param in params:
    if isinstance(param, piecewise.PiecewiseConstantFunc):
      jump_locations = param.jump_locations()
      if len(jump_locations.shape) > 1:
        # If `jump_locations` has batch dimension, transpose the result
        # Shape [num_times, dim]
        res.append(tf.transpose(param(times)))
      else:
        # Shape [num_times, dim]
        res.append(param(times))
    elif callable(param):
      # Used only in drift and volatility computation.
      # Here `times` is of shape [1]
      t = tf.squeeze(times)
      # The result has to have shape [1] + param.shape
      res.append(tf.expand_dims(param(t), 0))
    else:
      res.append(param + tf.zeros(times.shape + param.shape, dtype=times.dtype))
  return res 

def _apply_tridiag_matrix_explicitly(values, superdiag, diag, subdiag,
                                     dim, n_dims):
  """Applies tridiagonal matrix explicitly."""
  perm = _get_permutation(values, n_dims, dim)

  # Make the given dimension the last one in the tensors, treat all the
  # other spatial dimensions as batch dimensions.
  if perm is not None:
    values = tf.transpose(values, perm)
    superdiag, diag, subdiag = (
        tf.transpose(c, perm) for c in (superdiag, diag, subdiag))

  values = tf.squeeze(
      tf.linalg.tridiagonal_matmul((superdiag, diag, subdiag),
                                   tf.expand_dims(values, -1),
                                   diagonals_format='sequence'), -1)

  # Transpose back to how it was.
  if perm is not None:
    values = tf.transpose(values, perm)
  return values 

def testInhomogeneousBackwards(self, scheme, accuracy_order):
    # Tests solving du/dt = At + b for a backward time step.
    # Compares with exact solution u(0) = exp(-At) u(t)
    # + (exp(-At) - 1) A^(-1) b.
    time_step = 0.0001
    u = tf.constant([1, 2, -1, -2], dtype=tf.float64)
    matrix = tf.constant(
        [[1, -1, 0, 0], [3, 1, 2, 0], [0, -2, 1, 4], [0, 0, 3, 1]],
        dtype=tf.float64)
    b = tf.constant([1, -1, -2, 2], dtype=tf.float64)

    tridiag_form = self._convert_to_tridiagonal_format(matrix)
    actual = self.evaluate(scheme(u, time_step, 0, lambda t: (tridiag_form, b)))

    exponent = tf.linalg.expm(-matrix * time_step)
    eye = tf.eye(4, 4, dtype=tf.float64)
    u = tf.expand_dims(u, 1)
    b = tf.expand_dims(b, 1)
    expected = (
        tf.matmul(exponent, u) +
        tf.matmul(exponent - eye, tf.matmul(tf.linalg.inv(matrix), b)))
    expected = self.evaluate(tf.squeeze(expected))

    error_tolerance = 30 * time_step**(accuracy_order + 1)
    self.assertLess(np.max(np.abs(actual - expected)), error_tolerance) 

def testInhomogeneous(self, scheme, accuracy_order):
    # Tests solving du/dt = At + b for a time step.
    # Compares with exact solution u(t) = exp(At) u(0) + (exp(At) - 1) A^(-1) b.
    time_step = 0.0001
    u = tf.constant([1, 2, -1, -2], dtype=tf.float64)
    matrix = tf.constant(
        [[1, -1, 0, 0], [3, 1, 2, 0], [0, -2, 1, 4], [0, 0, 3, 1]],
        dtype=tf.float64)
    b = tf.constant([1, -1, -2, 2], dtype=tf.float64)

    tridiag_form = self._convert_to_tridiagonal_format(matrix)
    actual = self.evaluate(scheme(u, 0, time_step, lambda t: (tridiag_form, b)))

    exponent = tf.linalg.expm(matrix * time_step)
    eye = tf.eye(4, 4, dtype=tf.float64)
    u = tf.expand_dims(u, 1)
    b = tf.expand_dims(b, 1)
    expected = (
        tf.matmul(exponent, u) +
        tf.matmul(exponent - eye, tf.matmul(tf.linalg.inv(matrix), b)))
    expected = self.evaluate(tf.squeeze(expected))

    error_tolerance = 30 * time_step**(accuracy_order + 1)
    self.assertLess(np.max(np.abs(actual - expected)), error_tolerance) 

def testHomogeneousBackwards(self, scheme, accuracy_order):
    # Tests solving du/dt = At for a backward time step.
    # Compares with exact solution u(0) = exp(-At) u(t).
    time_step = 0.0001
    u = tf.constant([1, 2, -1, -2], dtype=tf.float64)
    matrix = tf.constant(
        [[1, -1, 0, 0], [3, 1, 2, 0], [0, -2, 1, 4], [0, 0, 3, 1]],
        dtype=tf.float64)

    tridiag_form = self._convert_to_tridiagonal_format(matrix)
    actual = self.evaluate(
        scheme(u, time_step, 0, lambda t: (tridiag_form, None)))

    expected = self.evaluate(
        tf.squeeze(
            tf.matmul(
                tf.linalg.expm(-matrix * time_step), tf.expand_dims(u, 1))))

    error_tolerance = 30 * time_step**(accuracy_order + 1)
    self.assertLess(np.max(np.abs(actual - expected)), error_tolerance) 

def testHomogeneous(self, scheme, accuracy_order):
    # Tests solving du/dt = At for a time step.
    # Compares with exact solution u(t) = exp(At) u(0).

    # Time step should be small enough to "resolve" different orders of accuracy
    time_step = 0.0001
    u = tf.constant([1, 2, -1, -2], dtype=tf.float64)
    matrix = tf.constant(
        [[1, -1, 0, 0], [3, 1, 2, 0], [0, -2, 1, 4], [0, 0, 3, 1]],
        dtype=tf.float64)

    tridiag_form = self._convert_to_tridiagonal_format(matrix)
    actual = self.evaluate(
        scheme(u, 0, time_step, lambda t: (tridiag_form, None)))
    expected = self.evaluate(
        tf.squeeze(
            tf.matmul(tf.linalg.expm(matrix * time_step), tf.expand_dims(u,
                                                                         1))))

    error_tolerance = 30 * time_step**(accuracy_order + 1)
    self.assertLess(np.max(np.abs(actual - expected)), error_tolerance) 

def _head(self, neck_outputs):
    logits = self._policy_logits(neck_outputs)
    if self._debug_writer:
      self._debug_writer.log_named_tensor(
          'action_logits', logits.numpy())
    value = tf.squeeze(self._baseline(neck_outputs), axis=-1)
    return common.AgentOutput(policy_logits=logits, baseline=value) 

def call(self, image_features, current_lstm_state):
    """Function call.

    Args:
      image_features: A tensor with shape[batch_size, num_views,
        feature_vector_length]
      current_lstm_state: A list of (state_c, state_h) tuple

    Returns:
      next_hidden_state: Hidden state vector [batch_size, lstm_space_size],
        current steps's LSTM output.
      next_lstm_state: Same shape as current_lstm_state.
    """
    # Attention-based visual-feature pooling. Pool the visual features of
    # shape [batch_size, num_views, feature_vector_length] to
    # [batch_size, attention_space_size].

    # LSTM state is a tuple (h, c) and `current_lstm_state` is a list of such
    # tuples. We use last LSTM layer's `h` to attention-pool current step's
    # image features.
    previous_step_lstm_output = current_lstm_state[-1][0]
    # [batch_size, 1, lstm_space_size]
    hidden_state = tf.expand_dims(previous_step_lstm_output, axis=1)
    # [batch_size, 1, attention_space_size]
    x = self._projection_hidden_layer(hidden_state)
    # [batch_size, num_view, attention_space_size]
    y = self._projection_image_feature(image_features)

    # v_t has shape[batch_size, 1, attention_space_size], representing the
    # current visual context.
    v_t = self.attention([x, y])


    v_t = tf.squeeze(v_t, axis=1)
    next_lstm_output, next_state = self.history_context_encoder(
        v_t, current_lstm_state)

    return (next_lstm_output, next_state) 

def _head(self, neck_outputs):
    # The shape of hidden_state is [batch_size * time, 1, hidden_size]
    hidden_state = neck_outputs['hidden_state']
    if self._scan_classifier is not None:
      scan_classifier_logits = self._scan_classifier(hidden_state)
    else:
      scan_classifier_logits = tf.zeros(shape=(tf.shape(hidden_state)[0], 61))
    image_features = tf.cast(neck_outputs[constants.PANO_ENC], tf.float32)

    # c_visual has shape [batch_size * time, 1, self._c_visual_attention_size]
    c_visual = self._visual_attention([
        self._visual_attention_project_ctext(neck_outputs['c_text']),
        self._visual_attention_project_feature(image_features),
    ])

    # Concatenating the h_t, c_text and c_visual as described in RCM paper.
    input_feature = tf.concat([hidden_state, neck_outputs['c_text'], c_visual],
                              axis=2)
    connection_encoding = neck_outputs[constants.CONN_ENC]
    connection_encoding = tf.cast(connection_encoding, tf.float32)
    logits = self._dot_product([
        self._project_feature(input_feature),
        self._project_action(connection_encoding)
    ])
    # The final shape of logits is [batch_size * time, num_connections]
    logits = tf.squeeze(logits, axis=1)
    # mask out invalid connections.
    valid_conn_mask = tf.cast(neck_outputs[constants.VALID_CONN_MASK],
                              tf.float32)
    logits += (1. - valid_conn_mask) * -_INFINITY
    value = self._value_network(tf.squeeze(neck_outputs['c_text'], axis=1))
    value = tf.squeeze(value, axis=1)
    return common.AgentOutput(
        policy_logits=logits,
        baseline={
            'value': value,
            'ins_classifier_logits': neck_outputs['ins_classifier_logits'],
            'scan_classifier_logits': scan_classifier_logits,
        }) 

def _head(self, neck_outputs):
    # The shape of hidden_state is [batch_size * time, 1, hidden_size]
    hidden_state = neck_outputs['hidden_state']
    image_features = tf.cast(neck_outputs[constants.PANO_ENC], tf.float32)

    # c_visual has shape [batch_size * time, 1, self._c_visual_attention_size]
    c_visual = self._visual_attention([
        self._visual_attention_project_ctext(neck_outputs['c_text']),
        self._visual_attention_project_feature(image_features),
    ])

    # Concatenating the h_t, c_text and c_visual as described in RCM paper.
    input_feature = tf.concat([hidden_state, neck_outputs['c_text'], c_visual],
                              axis=2)
    connection_encoding = neck_outputs[constants.CONN_ENC]
    connection_encoding = tf.cast(connection_encoding, tf.float32)
    logits = self._dot_product([
        self._project_feature(input_feature),
        self._project_action(connection_encoding)
    ])
    # The final shape of logits is [batch_size * time, num_connections]
    logits = tf.squeeze(logits, axis=1)
    # mask out invalid connections.
    valid_conn_mask = tf.cast(neck_outputs[constants.VALID_CONN_MASK],
                              tf.float32)
    logits += (1. - valid_conn_mask) * -_INFINITY
    value = self._value_network(tf.squeeze(neck_outputs['c_text'], axis=1))
    value = tf.squeeze(value, axis=1)
    return common.AgentOutput(policy_logits=logits, baseline=value) 

def testWrapperCall(self):
    wrapper = rnn_wrapper.RNNWrapper(
        tf.keras.layers.LSTM(3, return_state=True, return_sequences=True))

    batch_size = 2
    input_depth = 5
    inputs = np.random.rand(batch_size, input_depth).astype(np.float32)

    # Make sure wrapper call works when no time dimension is passed in.
    outputs, next_state = wrapper(inputs)

    inputs_time_dim = tf.expand_dims(inputs, axis=1)
    outputs_time_dim, next_state_time_dim = wrapper(inputs_time_dim)
    outputs_time_dim = tf.squeeze(outputs_time_dim, axis=1)

    outputs_manual_state, next_state_manual_state = wrapper(
        inputs, wrapper.get_initial_state(inputs))

    self.evaluate(tf.compat.v1.global_variables_initializer())
    for out_variant in (outputs, outputs_time_dim, outputs_manual_state):
      self.assertEqual(out_variant.shape, (batch_size, 3))
    for state_variant in (next_state, next_state_time_dim,
                          next_state_manual_state):
      self.assertLen(state_variant, 2)
      self.assertEqual(state_variant[0].shape, (batch_size, 3))
      self.assertEqual(state_variant[1].shape, (batch_size, 3))

    self.assertAllClose(outputs, outputs_time_dim)
    self.assertAllClose(outputs, outputs_manual_state)
    self.assertAllClose(next_state, next_state_time_dim)
    self.assertAllClose(next_state, next_state_manual_state) 

def squeeze(self, axis=None):
    """See tf.squeeze."""
    return self._apply_op(lambda t: tf.squeeze(t, axis)) 

def _update_variance(
    kappa, theta, epsilon, rho,
    current_vol, time_step, normals, psi_c=1.5):
  """Updates variance value."""
  del rho
  psi_c = tf.convert_to_tensor(psi_c, dtype=kappa.dtype)
  scaled_time = tf.exp(-kappa * time_step)
  epsilon_squared = epsilon**2
  m = theta + (current_vol - theta) * scaled_time
  s_squared = (
      current_vol * epsilon_squared * scaled_time / kappa
      * (1 - scaled_time) + theta * epsilon_squared / 2 / kappa
      * (1 - scaled_time)**2)
  psi = s_squared / m**2
  uniforms = 0.5 * (1 + tf.math.erf(normals[..., 0] / _SQRT_2))
  cond = psi < psi_c
  # Result where `cond` is true
  psi_inv = 2 / psi
  b_squared = psi_inv - 1 + tf.sqrt(psi_inv * (psi_inv - 1))

  a = m / (1 + b_squared)
  next_var_true = a * (tf.sqrt(b_squared) + tf.squeeze(normals[..., 1]))**2
  # Result where `cond` is false
  p = (psi - 1) / (psi + 1)
  beta = (1 - p) / m
  next_var_false = tf.where(uniforms > p,
                            tf.math.log(1 - p) - tf.math.log(1 - uniforms),
                            tf.zeros_like(uniforms)) / beta
  next_var = tf.where(cond, next_var_true, next_var_false)
  return next_var 

def test_univariate_xla_compatible(self):
    """Tests that univariate GBM sampling is XLA-compatible."""
    process = tff.models.GeometricBrownianMotion(0.05, 0.5, dtype=tf.float64)
    @tf.function
    def sample_fn():
      return process.sample_paths(
          times=[0.1, 0.5, 1.0], initial_state=2.0, num_samples=10000)
    samples = tf.xla.experimental.compile(sample_fn)[0]
    log_s = tf.math.log(samples)
    mean = tf.reduce_mean(log_s, axis=0)
    expected_mean = ((process._mu - process._sigma**2 / 2)
                     * 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 sparse_top_k_categorical_accuracy(y_true, y_pred, k=5):
  """TPU version of sparse_top_k_categorical_accuracy."""
  y_pred_rank = tf.convert_to_tensor(y_pred).get_shape().ndims
  y_true_rank = tf.convert_to_tensor(y_true).get_shape().ndims
  # If the shape of y_true is (num_samples, 1), squeeze to (num_samples,)
  if ((y_true_rank is not None) and
      (y_pred_rank is not None) and
      (len(tf.keras.backend.int_shape(y_true)) ==
       len(tf.keras.backend.int_shape(y_pred)))):
    y_true = tf.squeeze(y_true, [-1])

  y_true = tf.cast(y_true, 'int32')
  return tf.nn.in_top_k(y_true, y_pred, k) 

def _get_parameters(times, *params):
  """Gets parameter values at at specified `times`."""
  res = []
  for param in params:
    if callable(param):
      # Used only in drift and volatility computation.
      # Here `times` is of shape [1]
      t = tf.squeeze(times)
      # The result has to have shape [1] + param.shape
      param_value = tf.convert_to_tensor(param(t), dtype=times.dtype,
                                         name="param_value")
      res.append(tf.expand_dims(param_value, 0))
    else:
      res.append(param + tf.zeros(times.shape + param.shape, dtype=times.dtype))
  return res 

def _expected_exercise_fn(design, continuation_value, exercise_value):
  """Returns the expected continuation value for each path.

  Args:
    design: A real `Tensor` of shape `[basis_size, num_samples]`.
    continuation_value: A `Tensor` of shape `[num_samples, payoff_dim]` and of
      the same dtype as `design`. The optimal value of the option conditional on
      not exercising now or earlier, taking future information into account.
    exercise_value: A `Tensor` of the same shape and dtype as
      `continuation_value`. Value of the option if exercised immideately at
      the current time

  Returns:
    A `Tensor` of the same shape and dtype as `continuation_value` whose
    `(n, v)`-th entry represents the expected continuation value of sample path
    `n` under the `v`-th payoff scheme.
  """
  batch_design = tf.broadcast_to(
      design[..., None], design.shape + [continuation_value.shape[-1]])
  mask = tf.cast(exercise_value > 0, design.dtype)
  # Zero out contributions from samples we'd never exercise at this point (i.e.,
  # these extra observations do not change the regression coefficients).
  masked = tf.transpose(batch_design * mask, perm=(2, 1, 0))
  # For design matrix X and response y, the coefficients beta of the best linear
  # unbiased estimate are contained in the equation X'X beta = X'y. Here `lhs`
  # is X'X and `rhs` is X'y, or rather a tensor of such left hand and right hand
  # sides, one for each payoff dimension.
  lhs = tf.matmul(masked, masked, transpose_a=True)
  # Use pseudo inverse for the regression matrix to ensure stability of the
  # algorithm.
  lhs_pinv = tf.linalg.pinv(lhs)
  rhs = tf.matmul(
      masked,
      tf.expand_dims(tf.transpose(continuation_value), axis=-1),
      transpose_a=True)
  beta = tf.matmul(lhs_pinv, rhs)
  continuation = tf.matmul(tf.transpose(batch_design, perm=(2, 1, 0)), beta)
  return tf.maximum(tf.transpose(tf.squeeze(continuation, -1)), 0.0) 

def _put_valuer(sample_paths, time_index, strike_price, dtype=None, name=None):
  """Produces a callable from samples to payoff of a simple basket put option.

  Args:
    sample_paths: A `Tensor` of either `flaot32` or `float64` dtype and of
      shape `[num_samples, num_times, dim]`.
    time_index: An integer scalar `Tensor` that corresponds to the time
      coordinate at which the basis function is computed.
    strike_price: A `Tensor` of the same `dtype` as `sample_paths` and shape
      `[num_samples, num_strikes]`.
    dtype: Optional `dtype`. Either `tf.float32` or `tf.float64`. The `dtype`
      If supplied, represents the `dtype` for the 'strike_price' as well as
      for the input argument of the output payoff callable.
      Default value: `None`, which means that the `dtype` inferred by TensorFlow
      is used.
    name: Python `str` name prefixed to Ops created by the callable created
      by this function.
      Default value: `None` which is mapped to the default name 'put_valuer'

  Returns:
    A callable from `Tensor` of shape `[num_samples, num_exercise_times, dim]`
    and a scalar `Tensor` representing current time to a `Tensor` of shape
    `[num_samples, num_strikes]`.
  """
  name = name or "put_valuer"
  with tf.name_scope(name):
    strike_price = tf.convert_to_tensor(strike_price, dtype=dtype,
                                        name="strike_price")
    sample_paths = tf.convert_to_tensor(sample_paths, dtype=dtype,
                                        name="sample_paths")
    num_samples, _, dim = sample_paths.shape.as_list()

    slice_sample_paths = tf.slice(sample_paths, [0, time_index, 0],
                                  [num_samples, 1, dim])
    slice_sample_paths = tf.squeeze(slice_sample_paths, 1)
    average = tf.math.reduce_mean(slice_sample_paths, axis=-1, keepdims=True)
    return tf.nn.relu(strike_price - average) 

def _f_atm_second_derivative(self, s, cms_rates):
    """Computes second order derivative of _f_atm."""
    with tf.GradientTape() as g:
      g.watch(s)
      with tf.GradientTape() as gg:
        gg.watch(s)
        fx = self._f_atm(s, cms_rates)
      dfx = tf.squeeze(gg.gradient(fx, s))
    d2fx = tf.squeeze(g.gradient(dfx, s))
    return d2fx 

def _f_atm_first_derivative(self, s, cms_rates):
    """Computes first order derivative of _f_atm."""
    with tf.GradientTape() as g:
      g.watch(s)
      fx = self._f_atm(s, cms_rates)
    dfx = tf.squeeze(g.gradient(fx, s))
    return dfx 

def _adjust_convexity(self, valuation_date, market, model, pricing_context,
                        cms_rates, discount_factors):
    """Computes the convexity adjusted cms rate."""
    if model is None:
      return cms_rates
    elif model in (
        rc.InterestRateModelType.LOGNORMAL_SMILE_CONSISTENT_REPLICATION,
        rc.InterestRateModelType.NORMAL_SMILE_CONSISTENT_REPLICATION):
      return self._convexity_smile_replication(
          valuation_date, market, model, cms_rates, pricing_context)
    else:
      level = self._swap.annuity(valuation_date, market, None)
      expiry_time = dates.daycount_actual_365_fixed(
          start_date=valuation_date,
          end_date=self._coupon_start_dates,
          dtype=self._dtype)
      with tf.GradientTape() as g:
        g.watch(cms_rates)
        fx = self._fs(cms_rates)
      dfx = tf.squeeze(g.gradient(fx, cms_rates))
      swap_vol = tf.convert_to_tensor(pricing_context, dtype=self._dtype)
      if model == rc.InterestRateModelType.LOGNORMAL_RATE:
        cms_rates = cms_rates + dfx * level * (cms_rates**2) * (
            tf.math.exp(swap_vol**2 * expiry_time) - 1.0) / discount_factors
      else:
        cms_rates = cms_rates + dfx * level * (
            swap_vol**2 * expiry_time) / discount_factors
      return cms_rates 

def test_data_fitting(self):
    """Tests MLE estimation for a simple geometric GLM."""
    n, dim = 100, 3
    dtype = tf.float64
    np.random.seed(234095)
    x = np.random.choice([0, 1], size=[dim, n])
    s = 0.01 * np.sum(x, 0)
    p = 1. / (1 + np.exp(-s))
    y = np.random.geometric(p)
    x_data = tf.convert_to_tensor(value=x, dtype=dtype)
    y_data = tf.expand_dims(tf.convert_to_tensor(value=y, dtype=dtype), -1)

    def neg_log_likelihood(state):
      state_ext = tf.expand_dims(state, 0)
      linear_part = tf.matmul(state_ext, x_data)
      linear_part_ex = tf.stack([tf.zeros_like(linear_part), linear_part],
                                axis=0)
      term1 = tf.squeeze(
          tf.matmul(tf.reduce_logsumexp(linear_part_ex, axis=0), y_data), -1)
      term2 = (0.5 * tf.reduce_sum(state_ext * state_ext, axis=-1) -
               tf.reduce_sum(linear_part, axis=-1))
      return tf.squeeze(term1 + term2)

    self._check_algorithm(
        func=neg_log_likelihood,
        start_point=np.ones(shape=[dim]),
        expected_argmin=[-0.020460034354, 0.171708568111, 0.021200423717]) 

def _sample_paths(self, times, grid_step, keep_mask, num_requested_times,
                    num_samples, initial_state, random_type, seed, swap_memory):
    """Returns a sample of paths from the process."""
    dt = times[1:] - times[:-1]
    sqrt_dt = tf.sqrt(dt)
    current_state = initial_state + tf.zeros(
        [num_samples, self.dim()], dtype=initial_state.dtype)
    steps_num = tf.shape(dt)[-1]
    wiener_mean = tf.zeros((self.dim(), 1), dtype=self._dtype)

    cond_fn = lambda i, *args: i < steps_num

    def step_fn(i, written_count, current_state, result):
      """Performs one step of Euler scheme."""
      current_time = times[i + 1]
      dw = random_ops.mv_normal_sample((num_samples,),
                                       mean=wiener_mean,
                                       random_type=random_type,
                                       seed=seed)
      dw = dw * sqrt_dt[i]
      dt_inc = dt[i] * self.drift_fn()(current_time, current_state)  # pylint: disable=not-callable
      dw_inc = tf.squeeze(
          tf.matmul(self.volatility_fn()(current_time, current_state), dw), -1)  # pylint: disable=not-callable
      next_state = current_state + dt_inc + dw_inc

      def write_next_state_to_result():
        # Replace result[:, written_count, :] with next_state.
        one_hot = tf.one_hot(written_count, depth=num_requested_times)
        mask = tf.expand_dims(one_hot > 0, axis=-1)
        return tf.where(mask, tf.expand_dims(next_state, axis=1), result)

      # Keep only states for times requested by user.
      result = tf.cond(keep_mask[i + 1],
                       write_next_state_to_result,
                       lambda: result)
      written_count += tf.cast(keep_mask[i + 1], dtype=tf.int32)
      return i + 1, written_count, next_state, result

    # Maximum number iterations is passed to the while loop below. It improves
    # performance of the while loop on a GPU and is needed for XLA-compilation
    # comptatiblity
    maximum_iterations = (
        tf.cast(1. / grid_step, dtype=tf.int32) + tf.size(times))
    result = tf.zeros((num_samples, num_requested_times, self.dim()))
    _, _, _, result = tf.compat.v1.while_loop(
        cond_fn,
        step_fn, (0, 0, current_state, result),
        maximum_iterations=maximum_iterations,
        swap_memory=swap_memory)

    return result 

def call(self, inputs, initial_state=None, mask=None, training=False):
    """Perform the computation.

    Args:
      inputs: A tuple containing tensors in batch-major format, each shaped
        `[batch_size, n, ...]`.
      initial_state: (Optional) An initial state for the wrapped layer. If not
        provided, `get_initial_state()` is used instead.
      mask: The mask to pass down to the wrapped layer.
      training: Whether the output is being used for training.

    Returns:
      A 2-tuple `(outputs, final_state)` where:

       - `outputs` contains the outputs for all time steps of the unroll; this
         is typically a tensor shaped `[batch_size, n, ...]`.
       - `final_state` contains the final state.
    """
    inputs_flat = [
        tf.convert_to_tensor(x, name='input', dtype_hint=self.dtype)
        for x in tf.nest.flatten(inputs)
    ]
    has_time_axis = all(
        [x.shape.ndims is None or x.shape.ndims > 2 for x in inputs_flat])
    if not has_time_axis:
      inputs_flat = [tf.expand_dims(t, axis=1) for t in inputs_flat]
    inputs = tf.nest.pack_sequence_as(inputs, inputs_flat)

    # TODO(b/158804957): tf.function changes "if tensor:" to tensor bool expr.
    # pylint: disable=literal-comparison
    if initial_state is None or initial_state is () or initial_state is []:
      initial_state = self._layer.get_initial_state(inputs)
    # pylint: enable=literal-comparison

    outputs = self._layer(
        inputs, initial_state=initial_state, mask=mask, training=training)

    output, new_state = outputs[0], outputs[1:]

    # Keras RNN's outputs[1:] does not match the nest structure of its cells'
    # state_size property.  Restructure the output state to match.
    new_state = tf.nest.pack_sequence_as(
        self.state_size, tf.nest.flatten(new_state))

    if not has_time_axis:
      output = tf.nest.map_structure(lambda t: tf.squeeze(t, axis=1), output)

    # Outputs are in output, and state is in outputs[1:]
    return output, new_state


# Register with Keras so we can do type(layer).from_config(layer.get_config()) 

def call(self,
           image_embed,
           instructions,
           instruction_lengths,
           training=False):
    assert self.num_channels == image_embed.shape[3]

    text_embed = self.text_embedder(instructions)
    text_embed = self.rnn(text_embed, instruction_lengths, training)
    text_embed_1, text_embed_2 = tf.split(text_embed, 2, axis=-1)
    batch_size = text_embed.shape[0]

    # Compute 1x1 convolution weights
    kern1 = self.dense_k1(text_embed_1)
    kern2 = self.dense_k2(text_embed_2)
    kern1 = tf.reshape(kern1, (
        batch_size, 1, 1, self.num_channels, self.num_channels))
    kern2 = tf.reshape(kern2, (
        batch_size, 1, 1, self.num_channels, self.num_channels))

    f0 = image_embed
    f1 = self.conv1(f0)
    f2 = self.conv2(f1)

    # Filter encoded image features to produce language-conditioned features
    #

    g1 = utils.parallel_conv2d(f1, kern1, 1, "SAME")
    g2 = utils.parallel_conv2d(f2, kern2, 1, "SAME")

    h2 = self.deconv2(g2)
    h2_g1 = tf.concat([h2, g1], axis=3)  # Assuming NHWC

    h1 = self.deconv1(h2_g1)

    out1 = self.dense1(h1)
    out2 = self.dense2(out1)
    out = tf.squeeze(self.out_dense(out2), -1)

    out_flat = tf.reshape(out, [batch_size, -1])
    # So that the output forms a prob distribution.
    out_flat = tf.nn.softmax(out_flat)
    return out_flat 

def _piecewise_constant_integrate(x1, x2, jump_locations, values, batch_rank):
  """Integrates piecewise constant function between `x1` and `x2`."""
  # Initializer already verified that `jump_locations` and `values` have the
  # same shape.
  # Expand batch size to one if there is no batch shape.
  if x1.shape.as_list()[:batch_rank]:
    no_batch_shape = False
  else:
    no_batch_shape = True
    x1 = tf.expand_dims(x1, 0)
    x2 = tf.expand_dims(x2, 0)
  if not jump_locations.shape.as_list()[:-1]:
    jump_locations = tf.expand_dims(jump_locations, 0)
    values = tf.expand_dims(values, 0)
    batch_rank += 1

  # Compute the index matrix that is later used for `tf.gather_nd`.
  index_matrix = _prepare_index_matrix(
      x1.shape.as_list()[:-1], x1.shape.as_list()[-1], tf.int32)
  # Compute integral values at the jump locations starting from the first jump
  # location.
  event_shape = values.shape[(batch_rank+1):]
  num_data_points = values.shape.as_list()[batch_rank]
  diff = jump_locations[..., 1:] - jump_locations[..., :-1]
  # Broadcast `diff` to the shape of
  # `batch_shape + [num_data_points - 2] + [1] * sample_rank`.
  for _ in event_shape:
    diff = tf.expand_dims(diff, -1)
  slice_indices = batch_rank * [slice(None)]
  slice_indices += [slice(1, num_data_points - 1)]
  integrals = tf.cumsum(values[slice_indices] * diff, batch_rank)
  # Pad integrals with zero values on left and right.
  batch_shape = integrals.shape.as_list()[:batch_rank]
  zeros = tf.zeros(batch_shape + [1] + event_shape, dtype=integrals.dtype)
  integrals = tf.concat([zeros, integrals, zeros], axis=batch_rank)
  # Get jump locations and values and the integration end points
  value1, jump_location1, indices_nd1 = _get_indices_and_values(
      x1, index_matrix, jump_locations, values, 'left', batch_rank)
  value2, jump_location2, indices_nd2 = _get_indices_and_values(
      x2, index_matrix, jump_locations, values, 'right', batch_rank)
  integrals1 = tf.gather_nd(integrals, indices_nd1)
  integrals2 = tf.gather_nd(integrals, indices_nd2)
  # Broadcast `x1`, `x2`, `jump_location1`, `jump_location2` to the shape
  # `batch_shape + [num_points] + [1] * sample_rank`.
  for _ in event_shape:
    x1 = tf.expand_dims(x1, -1)
    x2 = tf.expand_dims(x2, -1)
    jump_location1 = tf.expand_dims(jump_location1, -1)
    jump_location2 = tf.expand_dims(jump_location2, -1)
  # Compute the value of the integral.
  res = ((jump_location1 - x1) * value1
         + (x2 - jump_location2) * value2
         + integrals2 - integrals1)
  if no_batch_shape:
    return tf.squeeze(res, 0)
  else:
    return res 

def get_discriminator_loss(learner_agent_output, env_output, actor_agent_output,
                           actor_action, reward_clipping, discounting,
                           baseline_cost, entropy_cost, num_steps):
  """Discriminator loss."""
  del actor_agent_output
  del actor_action
  del reward_clipping
  del discounting
  del baseline_cost
  del entropy_cost

  first_true = utils.get_first_true_column(env_output.observation['disc_mask'])
  output_logits = learner_agent_output.policy_logits
  output_logits = tf.squeeze(output_logits, axis=1)
  output_logits = tf.boolean_mask(output_logits, first_true)
  output_affine_a, output_affine_b = learner_agent_output.baseline

  # Get the first true.
  labels = tf.cast(env_output.observation['label'], tf.float32)
  labels = tf.boolean_mask(labels, first_true)

  positive_label = tf.equal(labels, tf.constant(1.0))
  positive_logits = tf.boolean_mask(output_logits, positive_label)
  tf.summary.histogram(
      'distribution/sigmoid_positive_logits',
      tf.sigmoid(positive_logits),
      step=num_steps)
  tf.summary.histogram(
      'distribution/positive_logits', positive_logits, step=num_steps)

  negative_label = tf.equal(labels, tf.constant(0.0))
  negative_logits = tf.boolean_mask(output_logits, negative_label)
  tf.summary.histogram(
      'distribution/sigmoid_negative_logits',
      tf.sigmoid(negative_logits),
      step=num_steps)
  tf.summary.histogram(
      'distribution/negative_logits', negative_logits, step=num_steps)

  tf.summary.scalar(
      'labels/positive_label',
      tf.reduce_mean(tf.cast(positive_label, tf.float32)),
      step=num_steps)

  tf.summary.scalar('labels/labels', tf.reduce_mean(labels), step=num_steps)
  tf.summary.scalar(
      'affine_transform/a', tf.reduce_mean(output_affine_a), step=num_steps)
  tf.summary.scalar(
      'affine_transform/b', tf.reduce_mean(output_affine_b), step=num_steps)

  cross_entropy = tf.nn.weighted_cross_entropy_with_logits(
      labels=labels, logits=output_logits, pos_weight=5)
  return cross_entropy