def test_batch_jacobian_fixed_shape(self):
   x = random_ops.random_uniform([2, 3, 5])
   y = x * x
   batch_jacobian_pfor = gradients.batch_jacobian(y, x, use_pfor=True)
   batch_jacobian_while = gradients.batch_jacobian(y, x, use_pfor=False)
   two_x = 2 * x
   answer = array_ops.stack(
       [array_ops.diag(two_x[0]),
        array_ops.diag(two_x[1])])
   self.run_and_assert_equal(answer, batch_jacobian_pfor)
   self.run_and_assert_equal(answer, batch_jacobian_while)

 def _jacobian(self, experimental_use_pfor):
   persistent = context.executing_eagerly and not experimental_use_pfor
   with backprop.GradientTape(persistent=persistent) as g:
     x = constant_op.constant([1., 2.])
     y = constant_op.constant([3., 4.])
     g.watch(x)
     g.watch(y)
     z = x * x * y
   jacobian = g.jacobian(z, [x, y],
                         experimental_use_pfor=experimental_use_pfor)
   answer = [array_ops.diag(2 * x * y), array_ops.diag(x * x)]
   return jacobian, answer

 def _batch_jacobian(self, experimental_use_pfor):
   persistent = context.executing_eagerly and not experimental_use_pfor
   with backprop.GradientTape(persistent=persistent) as g:
     x = constant_op.constant([[1., 2.], [3., 4.]])
     y = constant_op.constant([[3., 4.], [5., 6.]])
     g.watch(x)
     z = x * x * y
   batch_jacobian = g.batch_jacobian(
       z, x, experimental_use_pfor=experimental_use_pfor)
   answer = array_ops.stack([array_ops.diag(2 * x[0] * y[0]),
                             array_ops.diag(2 * x[1] * y[1])])
   return batch_jacobian, answer

 def test_batch_jacobian_unknown_shape(self):
   with self.test_session() as sess:
     x = array_ops.placeholder(dtypes.float32)
     y = x * x
     batch_jacobian_pfor = gradients.batch_jacobian(y, x, use_pfor=True)
     batch_jacobian_while = gradients.batch_jacobian(y, x, use_pfor=False)
     two_x = 2 * x
     answer = array_ops.stack(
         [array_ops.diag(two_x[0]),
          array_ops.diag(two_x[1])])
     ans, pfor_value, while_value = sess.run(
         [answer, batch_jacobian_pfor, batch_jacobian_while],
         feed_dict={x: [[1, 2], [3, 4]]})
     self.assertAllClose(ans, pfor_value)
     self.assertAllClose(ans, while_value)

def _symmetric_matrix_square_root(mat, eps=1e-10):
  """Compute square root of a symmetric matrix.

  Note that this is different from an elementwise square root. We want to
  compute M' where M' = sqrt(mat) such that M' * M' = mat.

  Also note that this method **only** works for symmetric matrices.

  Args:
    mat: Matrix to take the square root of.
    eps: Small epsilon such that any element less than eps will not be square
      rooted to guard against numerical instability.

  Returns:
    Matrix square root of mat.
  """
  # Unlike numpy, tensorflow's return order is (s, u, v)
  s, u, v = linalg_ops.svd(mat)
  # sqrt is unstable around 0, just use 0 in such case
  si = array_ops.where(math_ops.less(s, eps), s, math_ops.sqrt(s))
  # Note that the v returned by Tensorflow is v = V
  # (when referencing the equation A = U S V^T)
  # This is unlike Numpy which returns v = V^T
  return math_ops.matmul(
      math_ops.matmul(u, array_ops.diag(si)), v, transpose_b=True)

 def test_noise_decreasing(self):
   for dtype in [dtypes.float32, dtypes.float64]:
     with variable_scope.variable_scope(dtype.name):
       random_model = RandomStateSpaceModel(
           state_dimension=5, state_noise_dimension=4,
           configuration=state_space_model.StateSpaceModelConfiguration(
               dtype=dtype, num_features=1))
       random_model.initialize_graph()
       original_covariance = array_ops.diag(
           array_ops.ones(shape=[5], dtype=dtype))
       _, new_covariance, _ = random_model._exogenous_noise_decreasing(
           current_times=[[1]],
           exogenous_values=constant_op.constant([[-2.]], dtype=dtype),
           state=[
               -array_ops.ones(shape=[1, 5], dtype=dtype),
               original_covariance[None], [0]
           ])
       with self.cached_session() as session:
         variables.global_variables_initializer().run()
         evaled_new_covariance, evaled_original_covariance = session.run(
             [new_covariance[0], original_covariance])
         new_variances = numpy.diag(evaled_new_covariance)
         original_variances = numpy.diag(evaled_original_covariance)
         for i in range(5):
           self.assertLess(new_variances[i], original_variances[i])

 def _diagOp(self, diag, dtype, expected_ans, use_gpu):
   with self.cached_session(use_gpu=use_gpu):
     tf_ans = array_ops.diag(ops.convert_to_tensor(diag.astype(dtype)))
     out = self.evaluate(tf_ans)
     tf_ans_inv = array_ops.diag_part(expected_ans)
     inv_out = self.evaluate(tf_ans_inv)
   self.assertAllClose(out, expected_ans)
   self.assertAllClose(inv_out, diag)
   self.assertShapeEqual(expected_ans, tf_ans)
   self.assertShapeEqual(diag, tf_ans_inv)

 def diagOp(self, diag, dtype, expected_ans, use_gpu=False):
   with self.test_session(use_gpu=use_gpu):
     tf_ans = array_ops.diag(ops.convert_to_tensor(diag.astype(dtype)))
     out = tf_ans.eval()
     tf_ans_inv = array_ops.diag_part(expected_ans)
     inv_out = tf_ans_inv.eval()
   self.assertAllClose(out, expected_ans)
   self.assertAllClose(inv_out, diag)
   self.assertShapeEqual(expected_ans, tf_ans)
   self.assertShapeEqual(diag, tf_ans_inv)

  def __init__(self,
               data,
               num_classes,
               initial_means=None,
               params='wmc',
               covariance_type=FULL_COVARIANCE,
               random_seed=0):
    """Constructor.

    Args:
      data: a list of Tensors with data, each row is a new example.
      num_classes: number of clusters.
      initial_means: a Tensor with a matrix of means. If None, means are
        computed by sampling randomly.
      params: Controls which parameters are updated in the training
        process. Can contain any combination of "w" for weights, "m" for
        means, and "c" for covariances.
      covariance_type: one of "full", "diag".
      random_seed: Seed for PRNG used to initialize seeds.

    Raises:
      Exception if covariance type is unknown.
    """
    self._params = params
    self._random_seed = random_seed
    self._covariance_type = covariance_type
    if self._covariance_type not in [DIAG_COVARIANCE, FULL_COVARIANCE]:
      raise Exception(  # pylint: disable=g-doc-exception
          'programmer error: Invalid covariance type: %s' %
          self._covariance_type)
    # Create sharded variables for multiple shards. The following
    # lists are indexed by shard.
    # Probability per example in a class.
    num_shards = len(data)
    self._probs = [None] * num_shards
    # Prior probability.
    self._prior_probs = [None] * num_shards
    # Membership weights w_{ik} where "i" is the i-th example and "k"
    # is the k-th mixture.
    self._w = [None] * num_shards
    # Number of examples in a class.
    self._points_in_k = [None] * num_shards
    first_shard = data[0]
    self._dimensions = array_ops.shape(first_shard)[1]
    self._num_classes = num_classes
    # Small value to guarantee that covariances are invertible.
    self._min_var = array_ops.diag(
        array_ops.ones(array_ops.stack([self._dimensions]))) * 1e-3
    self._create_variables()
    self._initialize_variables(data, initial_means)
    # Operations of partial statistics for the computation of the means.
    self._w_mul_x = []
    # Operations of partial statistics for the computation of the covariances.
    self._w_mul_x2 = []
    self._define_graph(data)

  def testComputePiTracenorm(self):
    with ops.Graph().as_default(), self.test_session() as sess:
      random_seed.set_random_seed(200)
      left_factor = array_ops.diag([1., 2., 0., 1.])
      right_factor = array_ops.ones([2., 2.])

      # pi is the sqrt of the left trace norm divided by the right trace norm
      pi = fb.compute_pi_tracenorm(left_factor, right_factor)

      pi_val = sess.run(pi)
      self.assertEqual(1., pi_val)

 def testShuffle2d(self):
   with self.cached_session() as sess:
     with self.test_scope():
       x = array_ops.diag(math_ops.range(20))
       shuffle = random_ops.random_shuffle(x)
     result = sess.run(shuffle)
     expected = np.diag(range(20)).flatten()
     # Compare sets to avoid randomness behavior changes but make sure still
     # have all the values.
     self.assertAllEqual(len(result.flatten()), len(expected))
     self.assertAllEqual(set(result.flatten()), set(expected))

def _test1(op, grad_e, grad_v):
    """Gradient for SelfAdjointEigV2 derived with Joan with no adjustment for subspace"""
    e = op.outputs[0]
    v = op.outputs[1]
    #dim = v.get_shape()
    with ops.control_dependencies([grad_e.op, grad_v.op]):
        if grad_v is not None:  
            E = array_ops.diag(e)
            v_proj = array_ops.slice(v, [0,0], [20,2])
            grad_grassman = grad_v - math_ops.batch_matmul(math_ops.batch_matmul(v_proj, array_ops.transpose(v_proj)), grad_v)
            grad_a = math_ops.batch_matmul(grad_grassman, math_ops.batch_matmul(E, array_ops.transpose(grad_v)))+math_ops.batch_matmul(grad_v, math_ops.batch_matmul(E, array_ops.transpose(grad_grassman)))
    return grad_a

 def testBatchGradientUnknownSize(self):
   with self.test_session():
     batch_size = constant_op.constant(3)
     matrix_size = constant_op.constant(4)
     batch_identity = array_ops.tile(
         array_ops.expand_dims(
             array_ops.diag(array_ops.ones([matrix_size])), 0),
         [batch_size, 1, 1])
     determinants = linalg_ops.matrix_determinant(batch_identity)
     reduced = math_ops.reduce_sum(determinants)
     sum_grad = gradients_impl.gradients(reduced, batch_identity)[0]
     self.assertAllClose(batch_identity.eval(), sum_grad.eval())

 def testDiagGrad(self):
   np.random.seed(0)
   shapes = ((3,), (3, 3), (3, 3, 3))
   dtypes = (dtypes_lib.float32, dtypes_lib.float64)
   with self.test_session(use_gpu=False):
     errors = []
     for shape in shapes:
       for dtype in dtypes:
         x1 = constant_op.constant(np.random.rand(*shape), dtype=dtype)
         y = array_ops.diag(x1)
         error = gradient_checker.compute_gradient_error(
             x1, x1.get_shape().as_list(), y, y.get_shape().as_list())
         tf_logging.info("error = %f", error)
         self.assertLess(error, 1e-4)

def pairwise_distance(feature, squared=False):
  """Computes the pairwise distance matrix with numerical stability.

  output[i, j] = || feature[i, :] - feature[j, :] ||_2

  Args:
    feature: 2-D Tensor of size [number of data, feature dimension].
    squared: Boolean, whether or not to square the pairwise distances.

  Returns:
    pairwise_distances: 2-D Tensor of size [number of data, number of data].
  """
  pairwise_distances_squared = math_ops.add(
      math_ops.reduce_sum(
          math_ops.square(feature),
          axis=[1],
          keepdims=True),
      math_ops.reduce_sum(
          math_ops.square(
              array_ops.transpose(feature)),
          axis=[0],
          keepdims=True)) - 2.0 * math_ops.matmul(
              feature, array_ops.transpose(feature))

  # Deal with numerical inaccuracies. Set small negatives to zero.
  pairwise_distances_squared = math_ops.maximum(pairwise_distances_squared, 0.0)
  # Get the mask where the zero distances are at.
  error_mask = math_ops.less_equal(pairwise_distances_squared, 0.0)

  # Optionally take the sqrt.
  if squared:
    pairwise_distances = pairwise_distances_squared
  else:
    pairwise_distances = math_ops.sqrt(
        pairwise_distances_squared + math_ops.to_float(error_mask) * 1e-16)

  # Undo conditionally adding 1e-16.
  pairwise_distances = math_ops.multiply(
      pairwise_distances, math_ops.to_float(math_ops.logical_not(error_mask)))

  num_data = array_ops.shape(feature)[0]
  # Explicitly set diagonals to zero.
  mask_offdiagonals = array_ops.ones_like(pairwise_distances) - array_ops.diag(
      array_ops.ones([num_data]))
  pairwise_distances = math_ops.multiply(pairwise_distances, mask_offdiagonals)
  return pairwise_distances

def _matrix_square_root(mat, eps=1e-10):
  """Compute symmetric square root of matrix.

  Equivalent to matrix square root when matrix is invertible; note that this is
  different from an elementwise square root. We want to compute M' where M' =
  sqrt(mat) such that M' * M' = mat.

  Args:
    mat: Matrix to take the square root of.
    eps: Small epsilon such that any element less than eps will not be square
      rooted to guard against numerical instability.

  Returns:
    Matrix square root of mat.
  """
  s, u, v = linalg_ops.svd(mat)
  # sqrt is unstable around 0, just use 0 in such case
  si = array_ops.where(math_ops.less(s, eps), s, math_ops.sqrt(s))
  return math_ops.matmul(
      math_ops.matmul(u, array_ops.diag(si)), v, transpose_b=True)

def log_noninformative_covariance_prior(covariance):
  """Compute a relatively uninformative prior for noise parameters.

  Helpful for avoiding noise over-estimation, where noise otherwise decreases
  very slowly during optimization.

  See:
    Villegas, C. On the A Priori Distribution of the Covariance Matrix.
    Ann. Math. Statist. 40 (1969), no. 3, 1098--1099.

  Args:
    covariance: A covariance matrix.
  Returns:
    For a [p x p] matrix:
      log(det(covariance)^(-(p + 1) / 2))
  """
  # Avoid zero/negative determinants due to numerical errors
  covariance += array_ops.diag(1e-8 * array_ops.ones(
      shape=[array_ops.shape(covariance)[0]], dtype=covariance.dtype))
  power = -(math_ops.cast(array_ops.shape(covariance)[0] + 1,
                          covariance.dtype) / 2.)
  return power * math_ops.log(linalg_ops.matrix_determinant(covariance))

  def _define_distance_to_clusters(self, data):
    """Defines the Mahalanobis distance to the assigned Gaussian."""
    # TODO(xavigonzalvo): reuse (input - mean) * cov^-1 * (input -
    # mean) from log probability function.
    self._all_scores = []
    for shard in data:
      all_scores = []
      shard = array_ops.expand_dims(shard, 0)
      for c in xrange(self._num_classes):
        if self._covariance_type == FULL_COVARIANCE:
          cov = self._covs[c, :, :]
        elif self._covariance_type == DIAG_COVARIANCE:
          cov = array_ops.diag(self._covs[c, :])
        inverse = linalg_ops.matrix_inverse(cov + self._min_var)
        inv_cov = array_ops.tile(
            array_ops.expand_dims(inverse, 0),
            array_ops.stack([self._num_examples, 1, 1]))
        diff = array_ops.transpose(shard - self._means[c, :, :], perm=[1, 0, 2])
        m_left = math_ops.matmul(diff, inv_cov)
        all_scores.append(
            math_ops.sqrt(
                math_ops.matmul(
                    m_left, array_ops.transpose(
                        diff, perm=[0, 2, 1]))))
      self._all_scores.append(
          array_ops.reshape(
              array_ops.concat(all_scores, 1),
              array_ops.stack([self._num_examples, self._num_classes])))

    # Distance to the associated class.
    self._all_scores = array_ops.concat(self._all_scores, 0)
    assignments = array_ops.concat(self.assignments(), 0)
    rows = math_ops.to_int64(math_ops.range(0, self._num_examples))
    indices = array_ops.concat(
        [array_ops.expand_dims(rows, 1), array_ops.expand_dims(assignments, 1)],
        1)
    self._scores = array_ops.gather_nd(self._all_scores, indices)

def covariance_initializer(shape, dtype, partition_info=None):  # pylint: disable=unused-argument
  if INIT_COVARIANCES_AT_ZERO:
    return array_ops.diag(array_ops.zeros(shape[0], dtype))
  return array_ops.diag(array_ops.ones(shape[0], dtype))

def batch_matrix_pow(matrices, powers):
  """Compute powers of matrices, e.g. A^3 = matmul(matmul(A, A), A).

  Uses exponentiation by squaring, with O(log(p)) matrix multiplications to
  compute A^p.

  Args:
    matrices: [batch size x N x N]
    powers: Which integer power to raise each matrix to [batch size]
  Returns:
    The matrices raised to their respective powers, same dimensions as the
    "matrices" argument.
  """

  def terminate_when_all_zero(current_argument, residual_powers, accumulator):
    del current_argument, accumulator  # not used for condition
    do_exit = math_ops.reduce_any(
        math_ops.greater(residual_powers, array_ops.ones_like(residual_powers)))
    return do_exit

  def do_iteration(current_argument, residual_powers, accumulator):
    """Compute one step of iterative exponentiation by squaring.

    The recursive form is:
      power(A, p) = { power(matmul(A, A), p / 2) for even p
                    { matmul(A, power(matmul(A, A), (p - 1) / 2)) for odd p
      power(A, 0) = I

    The power(A, 0) = I case is handeled by starting with accumulator set to the
    identity matrix; matrices with zero residual powers are passed through
    unchanged.

    Args:
      current_argument: On this step, what is the first argument (A^2..^2) to
          the (unrolled) recursive function? [batch size x N x N]
      residual_powers: On this step, what is the second argument (residual p)?
          [batch_size]
      accumulator: Accumulates the exterior multiplications from the odd
          powers (initially the identity matrix). [batch_size x N x N]
    Returns:
      Updated versions of each argument for one step of the unrolled
      computation. Does not change parts of the batch which have a residual
      power of zero.
    """
    is_even = math_ops.equal(residual_powers % 2,
                             array_ops.zeros(
                                 array_ops.shape(residual_powers),
                                 dtype=dtypes.int32))
    new_accumulator = array_ops.where(is_even, accumulator,
                                      math_ops.matmul(accumulator,
                                                      current_argument))
    new_argument = math_ops.matmul(current_argument, current_argument)
    do_update = math_ops.greater(residual_powers, 1)
    new_residual_powers = residual_powers - residual_powers % 2
    new_residual_powers //= 2
    # Stop updating if we've reached our base case; some batch elements may
    # finish sooner than others
    accumulator = array_ops.where(do_update, new_accumulator, accumulator)
    current_argument = array_ops.where(do_update, new_argument,
                                       current_argument)
    residual_powers = array_ops.where(do_update, new_residual_powers,
                                      residual_powers)
    return (current_argument, residual_powers, accumulator)

  matrices = ops.convert_to_tensor(matrices)
  powers = math_ops.cast(powers, dtype=dtypes.int32)
  ident = array_ops.expand_dims(
      array_ops.diag(
          array_ops.ones([array_ops.shape(matrices)[1]], dtype=matrices.dtype)),
      0)
  ident_tiled = array_ops.tile(ident, [array_ops.shape(matrices)[0], 1, 1])
  (final_argument,
   final_residual_power, final_accumulator) = control_flow_ops.while_loop(
       terminate_when_all_zero, do_iteration, [matrices, powers, ident_tiled])
  return array_ops.where(
      math_ops.equal(final_residual_power,
                     array_ops.zeros_like(
                         final_residual_power, dtype=dtypes.int32)),
      ident_tiled, math_ops.matmul(final_argument, final_accumulator))