Python tensorflow.matrix_diag_part() Examples

def effective_steps(masks, num_edge_residues, name=None):
    """ Returns the effective number of steps, i.e. number of residues that are non-missing and are not just
        padding, given a masking matrix.

    Args:
        masks: A batch of square masking matrices (batch is last dimension)
        [MAX_SEQ_LENGTH, MAX_SEQ_LENGTH, BATCH_SIZE]

    Returns:
        A vector with the effective number of steps
        [BATCH_SIZE]

    """

    with tf.name_scope(name, 'effective_steps', [masks]) as scope:
        masks = tf.convert_to_tensor(masks, name='masks')
        
        traces = tf.matrix_diag_part(tf.transpose(masks, [2, 0, 1]))
        eff_stepss = tf.add(tf.reduce_sum(traces, [1]), num_edge_residues, name=scope) # NUM_EDGE_RESIDUES shouldn't be here, but I'm keeping it for 
                                                                                       # legacy reasons. Just be clear that it's _always_ wrong to have
                                                                                       # it here, even when it's not equal to 0.

        return eff_stepss 

def matrix_mean_wo_diagonal(matrix, num_row, num_col=None, name='mu_wo_diag'):
    """ This function calculates the mean of the matrix elements not in the diagonal

    2018.4.9 - replace tf.diag_part with tf.matrix_diag_part
    tf.matrix_diag_part can be used for rectangle matrix while tf.diag_part can only be used for square matrix

    :param matrix:
    :param num_row:
    :type num_row: float
    :param num_col:
    :type num_col: float
    :param name:
    :return:
    """
    with tf.name_scope(name):
        if num_col is None:
            mu = (tf.reduce_sum(matrix) - tf.reduce_sum(tf.matrix_diag_part(matrix))) / (num_row * (num_row - 1.0))
        else:
            mu = (tf.reduce_sum(matrix) - tf.reduce_sum(tf.matrix_diag_part(matrix))) \
                 / (num_row * num_col - tf.minimum(num_col, num_row))

    return mu


######################################################################## 

def fit(self, x=None, y=None):
    # p(coeffs | x, y) = Normal(coeffs |
    #   mean = (1/noise_variance) (1/noise_variance x^T x + I)^{-1} x^T y,
    #   covariance = (1/noise_variance x^T x + I)^{-1})
    # TODO(trandustin): We newly fit the data at each call. Extend to do
    # Bayesian updating.
    kernel_matrix = tf.matmul(x, x, transpose_a=True) / self.noise_variance
    coeffs_precision = tf.matrix_set_diag(
        kernel_matrix, tf.matrix_diag_part(kernel_matrix) + 1.)
    coeffs_precision_tril = tf.linalg.cholesky(coeffs_precision)
    self.coeffs_precision_tril_op = tf.linalg.LinearOperatorLowerTriangular(
        coeffs_precision_tril)
    self.coeffs_mean = self.coeffs_precision_tril_op.solvevec(
        self.coeffs_precision_tril_op.solvevec(tf.einsum('nm,n->m', x, y)),
        adjoint=True) / self.noise_variance
    # TODO(trandustin): To be fully Keras-compatible, return History object.
    return 

def get_matrix_tree(r, A):
    L = tf.reduce_sum(A, 1)
    L = tf.matrix_diag(L)
    L = L - A

    r_diag = tf.matrix_diag(r)
    LL = L + r_diag

    LL_inv = tf.matrix_inverse(LL)  #batch_l, doc_l, doc_l
    LL_inv_diag_ = tf.matrix_diag_part(LL_inv)

    d0 = tf.multiply(r, LL_inv_diag_)

    LL_inv_diag = tf.expand_dims(LL_inv_diag_, 2)

    tmp1 = tf.multiply(A, tf.matrix_transpose(LL_inv_diag))
    tmp2 = tf.multiply(A, tf.matrix_transpose(LL_inv))

    d = tmp1 - tmp2
    d = tf.concat([tf.expand_dims(d0,[1]), d], 1)
    return d 

def _update_memory(self, old_memory, w_samples, new_z_mean, new_z_var):
    """Setting new_z_var=0 for sample based update."""
    old_mean, old_cov = old_memory
    wR = self.get_w_to_z_mean(w_samples, old_memory.M_mean)
    wU, wUw = self._read_cov(w_samples, old_memory)
    sigma_z = wUw + new_z_var + self._obs_noise_stddev**2  # [S, B]
    delta = new_z_mean - wR  # [S, B, C]
    c_z = wU / tf.expand_dims(sigma_z, -1)  # [S, B, M]
    posterior_mean = old_mean + tf.einsum('sbm,sbc->bmc', c_z, delta)
    posterior_cov = old_cov - tf.einsum('sbm,sbn->bmn', c_z, wU)
    # Clip diagonal elements for numerical stability
    posterior_cov = tf.matrix_set_diag(
        posterior_cov,
        tf.clip_by_value(tf.matrix_diag_part(posterior_cov), EPSILON, 1e10))
    new_memory = MemoryState(M_mean=posterior_mean, M_cov=posterior_cov)
    return new_memory 

def fit(self, x=None, y=None):
    # p(coeffs | x, y) = Normal(coeffs |
    #   mean = (1/noise_variance) (1/noise_variance x^T x + I)^{-1} x^T y,
    #   covariance = (1/noise_variance x^T x + I)^{-1})
    # TODO(trandustin): We newly fit the data at each call. Extend to do
    # Bayesian updating.
    kernel_matrix = tf.matmul(x, x, transpose_a=True) / self.noise_variance
    coeffs_precision = tf.matrix_set_diag(
        kernel_matrix, tf.matrix_diag_part(kernel_matrix) + 1.)
    coeffs_precision_tril = tf.linalg.cholesky(coeffs_precision)
    self.coeffs_precision_tril_op = tf.linalg.LinearOperatorLowerTriangular(
        coeffs_precision_tril)
    self.coeffs_mean = self.coeffs_precision_tril_op.solvevec(
        self.coeffs_precision_tril_op.solvevec(tf.einsum('nm,n->m', x, y)),
        adjoint=True) / self.noise_variance
    # TODO(trandustin): To be fully Keras-compatible, return History object.
    return 

def _svd(A, name=None):
        """Tensorflow svd with gradient.

        Parameters
        ----------
        A : Tensor
            The matrix for which to compute singular values.
        name : string, optional

        Returns
        -------
        s : Tensor
            The singular values of A.

        """
        S0, U0, V0 = map(tf.stop_gradient,
                         tf.svd(A, full_matrices=True, name=name))
        # A = U * S * V.T
        # S = inv(U) * A * inv(V.T) = U.T * A * V  (orthogonal matrices)
        S = tf.matmul(U0, tf.matmul(A, V0),
                      transpose_a=True)
        return tf.matrix_diag_part(S) 

def _define_full_covariance_probs(self, shard_id, shard):
    """Defines the full covariance probabilties per example in a class.

    Updates a matrix with dimension num_examples X num_classes.

    Args:
      shard_id: id of the current shard.
      shard: current data shard, 1 X num_examples X dimensions.
    """
    diff = shard - self._means
    cholesky = tf.cholesky(self._covs + self._min_var)
    log_det_covs = 2.0 * tf.reduce_sum(tf.log(tf.matrix_diag_part(cholesky)), 1)
    x_mu_cov = tf.square(
        tf.matrix_triangular_solve(
            cholesky, tf.transpose(
                diff, perm=[0, 2, 1]), lower=True))
    diag_m = tf.transpose(tf.reduce_sum(x_mu_cov, 1))
    self._probs[shard_id] = -0.5 * (
        diag_m + tf.to_float(self._dimensions) * tf.log(2 * np.pi) +
        log_det_covs) 

def testGrad(self):
    shapes = ((3, 3), (2, 3), (3, 2), (5, 3, 3))
    with self.test_session(use_gpu=self._use_gpu):
      for shape in shapes:
        x = tf.constant(np.random.rand(*shape), dtype=np.float32)
        y = tf.matrix_diag_part(x)
        error = tf.test.compute_gradient_error(x, x.get_shape().as_list(),
                                               y, y.get_shape().as_list())
        self.assertLess(error, 1e-4) 

def row_mean_wo_diagonal(matrix, num_col, name='mu_wo_diag'):
    """ This function calculates the mean of each row of the matrix elements excluding the diagonal
    
    :param matrix:
    :param num_col:
    :type num_col: float
    :param name:
    :return: 
    """
    with tf.name_scope(name):
        return (tf.reduce_sum(matrix, axis=1) - tf.matrix_diag_part(matrix)) / (num_col - 1.0)


######################################################################### 

def get_dkl_total(self, memory_state):
    """Compute the KL-divergence between a memory distribution and its prior."""
    R, U = memory_state
    B, K, _ = R.get_shape().as_list()
    U.get_shape().assert_is_compatible_with([B, K, K])
    R_prior, U_prior = self.get_prior_state(B)
    p_diag = tf.matrix_diag_part(U_prior)
    q_diag = tf.matrix_diag_part(U)  # B, K
    t1 = self._code_size * tf.reduce_sum(q_diag / p_diag, -1)
    t2 = tf.reduce_sum((R - R_prior)**2 / tf.expand_dims(
        p_diag, -1), [-2, -1])
    t3 = -self._code_size * self._memory_size
    t4 = self._code_size * tf.reduce_sum(tf.log(p_diag) - tf.log(q_diag), -1)
    return t1 + t2 + t3 + t4 

def KL(self):
        """
        The KL divergence from the variational distribution to the prior

        :return: KL divergence from N(q_mu, q_sqrt) to N(0, I), independently for each GP
        """
        # if self.white:
        #     return gauss_kl(self.q_mu, self.q_sqrt)
        # else:
        #     return gauss_kl(self.q_mu, self.q_sqrt, self.Ku)

        self.build_cholesky_if_needed()

        KL = -0.5 * self.num_outputs * self.num_inducing
        KL -= 0.5 * tf.reduce_sum(tf.log(tf.matrix_diag_part(self.q_sqrt) ** 2))

        if not self.white:
            KL += tf.reduce_sum(tf.log(tf.matrix_diag_part(self.Lu))) * self.num_outputs
            KL += 0.5 * tf.reduce_sum(tf.square(tf.matrix_triangular_solve(self.Lu_tiled, self.q_sqrt, lower=True)))
            Kinv_m = tf.cholesky_solve(self.Lu, self.q_mu)
            KL += 0.5 * tf.reduce_sum(self.q_mu * Kinv_m)
        else:
            KL += 0.5 * tf.reduce_sum(tf.square(self.q_sqrt))
            KL += 0.5 * tf.reduce_sum(self.q_mu**2)

        return KL 

def _objective(xx, obj):
  """Objective function as custom op so that we can overload gradients."""
  with tf.name_scope('objective'):
    chol = tf.cholesky(xx)
    choli = tf.linalg.inv(chol)

    rq = tf.matmul(choli, tf.matmul(obj, choli, transpose_b=True))
    eigval = tf.matrix_diag_part(rq)
    loss = tf.trace(rq)
    grad = functools.partial(_objective_grad, xx, obj)
  return (loss, eigval, chol), grad 

def _objective_grad(xx, obj, grad_loss, grad_eigval, grad_chol):
  """Symbolic form of the gradient of the objective with stop_gradients."""
  del grad_eigval
  del grad_chol
  with tf.name_scope('objective_grad'):
    chol = tf.cholesky(xx)
    choli = tf.linalg.inv(chol)
    rq = tf.matmul(choli, tf.matmul(obj, choli, transpose_b=True))

    dl = tf.diag(tf.matrix_diag_part(choli))
    triu = tf.matrix_band_part(tf.matmul(rq, dl), 0, -1)
    gxx = -1.0 * tf.matmul(choli, triu, transpose_a=True)
    gobj = tf.matmul(choli, dl, transpose_a=True)

    return grad_loss * gxx, grad_loss * gobj 

def matrix_diag_part(self, a):
        return tf.matrix_diag_part(a) 

def logdet(self, A, **kwargs):
        A = (A + self.matrix_transpose(A)) / 2.
        term = tf.log(tf.matrix_diag_part(self.cholesky(A, **kwargs)))
        return 2 * tf.reduce_sum(term, -1) 

def build_variational(hps, kernel, z_pos, x, n_particles):
    bn = zs.BayesianNet()
    z_mean = tf.get_variable(
        'z/mean', [hps.n_z], hps.dtype, tf.zeros_initializer())
    z_cov_raw = tf.get_variable(
        'z/cov_raw', initializer=tf.eye(hps.n_z, dtype=hps.dtype))
    z_cov_tril = tf.matrix_set_diag(
        tf.matrix_band_part(z_cov_raw, -1, 0),
        tf.nn.softplus(tf.matrix_diag_part(z_cov_raw)))
    fz = bn.multivariate_normal_cholesky(
        'fz', z_mean, z_cov_tril, n_samples=n_particles)
    bn.stochastic('fx', gp_conditional(z_pos, fz, x, False, kernel))
    return bn 

def testInvalidShape(self):
    with self.assertRaisesRegexp(ValueError, "must be at least rank 2"):
      tf.matrix_diag_part(0) 

def testRectangularBatch(self):
    with self.test_session(use_gpu=self._use_gpu):
      v_batch = np.array([[1.0, 2.0],
                          [4.0, 5.0]])
      mat_batch = np.array(
          [[[1.0, 0.0, 0.0],
            [0.0, 2.0, 0.0]],
           [[4.0, 0.0, 0.0],
            [0.0, 5.0, 0.0]]])
      self.assertEqual(mat_batch.shape, (2, 2, 3))
      mat_batch_diag = tf.matrix_diag_part(mat_batch)
      self.assertEqual((2, 2), mat_batch_diag.get_shape())
      self.assertAllEqual(mat_batch_diag.eval(), v_batch) 

def testSquareBatch(self):
    with self.test_session(use_gpu=self._use_gpu):
      v_batch = np.array([[1.0, 2.0, 3.0],
                          [4.0, 5.0, 6.0]])
      mat_batch = np.array(
          [[[1.0, 0.0, 0.0],
            [0.0, 2.0, 0.0],
            [0.0, 0.0, 3.0]],
           [[4.0, 0.0, 0.0],
            [0.0, 5.0, 0.0],
            [0.0, 0.0, 6.0]]])
      self.assertEqual(mat_batch.shape, (2, 3, 3))
      mat_batch_diag = tf.matrix_diag_part(mat_batch)
      self.assertEqual((2, 3), mat_batch_diag.get_shape())
      self.assertAllEqual(mat_batch_diag.eval(), v_batch) 

def testRectangular(self):
    with self.test_session(use_gpu=self._use_gpu):
      mat = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
      mat_diag = tf.matrix_diag_part(mat)
      self.assertAllEqual(mat_diag.eval(), np.array([1.0, 5.0]))
      mat = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
      mat_diag = tf.matrix_diag_part(mat)
      self.assertAllEqual(mat_diag.eval(), np.array([1.0, 4.0])) 

def testSquare(self):
    with self.test_session(use_gpu=self._use_gpu):
      v = np.array([1.0, 2.0, 3.0])
      mat = np.diag(v)
      mat_diag = tf.matrix_diag_part(mat)
      self.assertEqual((3,), mat_diag.get_shape())
      self.assertAllEqual(mat_diag.eval(), v) 

def test_MatrixDiagPart(self):
        if td._tf_version[:2] >= (0, 12):
            t = tf.matrix_diag_part(self.random(3, 4, 4, 5))
            self.check(t) 

def _chollogdet(L):
    """Log det of a cholesky, where L is (..., D, D)."""
    ldiag = tf.maximum(tf.abs(tf.matrix_diag_part(L)), JIT)  # keep > 0
    logdet = 2. * tf.reduce_sum(tf.log(ldiag))
    return logdet 

def testSample(self):
    with self.test_session():
      scale = make_pd(1., 2)
      df = 4

      chol_w = distributions.WishartCholesky(
          df, chol(scale), cholesky_input_output_matrices=False)

      x = chol_w.sample(1, seed=42).eval()
      chol_x = [chol(x[0])]

      full_w = distributions.WishartFull(
          df, scale, cholesky_input_output_matrices=False)
      self.assertAllClose(x, full_w.sample(1, seed=42).eval())

      chol_w_chol = distributions.WishartCholesky(
          df, chol(scale), cholesky_input_output_matrices=True)
      self.assertAllClose(chol_x, chol_w_chol.sample(1, seed=42).eval())
      eigen_values = tf.matrix_diag_part(chol_w_chol.sample(1000, seed=42))
      np.testing.assert_array_less(0., eigen_values.eval())

      full_w_chol = distributions.WishartFull(
          df, scale, cholesky_input_output_matrices=True)
      self.assertAllClose(chol_x, full_w_chol.sample(1, seed=42).eval())
      eigen_values = tf.matrix_diag_part(full_w_chol.sample(1000, seed=42))
      np.testing.assert_array_less(0., eigen_values.eval())

      # Check first and second moments.
      df = 4.
      chol_w = distributions.WishartCholesky(
          df=df,
          scale=chol(make_pd(1., 3)),
          cholesky_input_output_matrices=False)
      x = chol_w.sample(10000, seed=42)
      self.assertAllEqual((10000, 3, 3), x.get_shape())

      moment1_estimate = tf.reduce_mean(x, reduction_indices=[0]).eval()
      self.assertAllClose(chol_w.mean().eval(),
                          moment1_estimate,
                          rtol=0.05)

      # The Variance estimate uses the squares rather than outer-products
      # because Wishart.Variance is the diagonal of the Wishart covariance
      # matrix.
      variance_estimate = (
          tf.reduce_mean(tf.square(x), reduction_indices=[0]) -
          tf.square(moment1_estimate)).eval()
      self.assertAllClose(chol_w.variance().eval(),
                          variance_estimate,
                          rtol=0.05)

  # Test that sampling with the same seed twice gives the same results. 

def build(self):

    dd_q_input = Input((self.config.nb_supervised_doc, self.config.doc_topk_term, 1), name='dd_q_input')
    dd_d_input = Input((self.config.nb_supervised_doc, self.config.doc_topk_term,
                        self.config.hist_size), name='dd_d_input')

    dd_q_w = Dense(1, kernel_initializer=self.initializer_gate, use_bias=False, name='dd_q_gate')(dd_q_input)
    dd_q_w = Lambda(lambda x: softmax(x, axis=2), output_shape=(
                  self.config.nb_supervised_doc, self.config.doc_topk_term,), name='dd_q_softmax')(dd_q_w)

    z = dd_d_input
    for i in range(self.config.nb_layers):
      z = Dense(self.config.hidden_size[i], activation='tanh',
                kernel_initializer=self.initializer_fc, name='hidden')(z)
    z = Dense(self.config.out_size, kernel_initializer=self.initializer_fc, name='dd_d_gate')(z)
    z = Reshape((self.config.nb_supervised_doc, self.config.doc_topk_term,))(z)
    dd_q_w = Reshape((self.config.nb_supervised_doc, self.config.doc_topk_term,))(dd_q_w)
    # out = Dot(axes=[2, 2], name='dd_pseudo_out')([z, dd_q_w])

    out = Lambda(lambda x: K.batch_dot(x[0], x[1], axes=[2, 2]), name='dd_pseudo_out')([z, dd_q_w])
    dd_init_out = Lambda(lambda x: tf.matrix_diag_part(x), output_shape=(self.config.nb_supervised_doc,), name='dd_init_out')(out)
    '''
    dd_init_out = Lambda(lambda x: tf.reduce_sum(x, axis=2), output_shape=(self.config.nb_supervised_doc,))(z)
    '''
    #dd_out = Reshape((self.config.nb_supervised_doc,))(dd_out)

    # dd out gating
    dd_gate = Input((self.config.nb_supervised_doc, 1), name='baseline_doc_score')
    dd_w = Dense(1, kernel_initializer=self.initializer_gate, use_bias=False, name='dd_gate')(dd_gate)
    # dd_w = Lambda(lambda x: softmax(x, axis=1), output_shape=(self.config.nb_supervised_doc,), name='dd_softmax')(dd_w)

    # dd_out = Dot(axes=[1, 1], name='dd_out')([dd_init_out, dd_w])
    dd_w = Reshape((self.config.nb_supervised_doc,))(dd_w)
    dd_init_out = Reshape((self.config.nb_supervised_doc,))(dd_init_out)


    if self.config.method in [1, 3]: # no doc gating, with dense layer
      z = dd_init_out
    elif self.config.method == 2:
      logging.info("Apply doc gating")
      z = Multiply(name='dd_out')([dd_init_out, dd_w])
    else:
      raise ValueError("Method not initialized, please check config file")

    if self.config.method in [1, 2]:
      logging.info("Dense layer on top")
      z = Dense(self.config.merge_hidden, activation='tanh', name='merge_hidden')(z)
      out = Dense(self.config.merge_out, name='score')(z)
    else:
      logging.info("Apply doc gating, No dense layer on top, sum up scores")
      out = Dot(axes=[1, 1], name='score')([z, dd_w])

    model = Model(inputs=[dd_q_input, dd_d_input, dd_gate], outputs=[out])
    print(model.summary())

    return model 

def layer_xyzq(matrix_rt, scope='pose', name='xyzq'):
    # Rotation Matrix to quaternion + xyz
    with tf.variable_scope(scope):
        qw = tf.sqrt(tf.reduce_sum(tf.matrix_diag_part(matrix_rt), axis=-1)) / 2.0
        qx = (matrix_rt[:, 2, 1] - matrix_rt[:, 1, 2]) / (4 * qw)
        qy = (matrix_rt[:, 0, 2] - matrix_rt[:, 2, 0]) / (4 * qw)
        qz = (matrix_rt[:, 1, 0] - matrix_rt[:, 0, 1]) / (4 * qw)

        '''
        # See : http://www.euclideanspace.com/maths/geometry/rotations/conversions/matrixToQuaternion/
        trace = tf.reduce_sum(tf.matrix_diag_part(matrix_rt), axis=-1) - 1.0

        if trace > 0:
            S = tf.sqrt(trace + 1.0) * 2
            qw = 0.25 * S
            qx = (matrix_rt[:, 2, 1] - matrix_rt[:, 1, 2]) / S
            qy = (matrix_rt[:, 0, 2] - matrix_rt[:, 2, 0]) / S
            qz = (matrix_rt[:, 1, 0] - matrix_rt[:, 0, 1]) / S
        elif matrix_rt[:, 0, 0] > matrix_rt[:, 1, 1] and matrix_rt[:, 0, 0] > matrix_rt[:, 2, 2]:
            S = tf.sqrt(1 + matrix_rt[:, 0, 0] - matrix_rt[:, 1, 1] - matrix_rt[:, 2, 2]) * 2
            qw = (matrix_rt[:, 2, 1] - matrix_rt[:, 1, 2]) / S
            qx = 0.25 * S
            qy = (matrix_rt[:, 0, 1] + matrix_rt[:, 1, 0]) / S
            qz = (matrix_rt[:, 0, 2] + matrix_rt[:, 2, 0]) / S
        elif matrix_rt[:, 1, 1] > matrix_rt[:, 2, 2]:
            S = tf.sqrt(1 + matrix_rt[:, 1, 1] - matrix_rt[:, 0, 0] - matrix_rt[:, 2, 2]) * 2
            qw = (matrix_rt[:, 0, 2] - matrix_rt[:, 2, 0]) / S
            qx = (matrix_rt[:, 0, 1] + matrix_rt[:, 1, 0]) / S
            qy = 0.25 * S
            qz = (matrix_rt[:, 1, 2] + matrix_rt[:, 2, 1]) / S
        else:
            S = tf.sqrt(1 + matrix_rt[:, 2, 2] - matrix_rt[:, 0, 0] - matrix_rt[:, 1, 1]) * 2
            qw = (matrix_rt[:, 1, 0] - matrix_rt[:, 0, 1]) / S
            qx = (matrix_rt[:, 0, 2] + matrix_rt[:, 2, 0]) / S
            qy = (matrix_rt[:, 1, 2] + matrix_rt[:, 2, 1]) / S
            qz = 0.25 * S
        '''

        x = matrix_rt[:, 0, 3]
        y = matrix_rt[:, 1, 3]
        z = matrix_rt[:, 2, 3]

        xyzq = tf.stack([qw, qx, qy, qz, x, y, z], axis=-1, name=name)
    return xyzq 

def call(self, inputs):
    if self.conditional_inputs is None and self.conditional_outputs is None:
      covariance_matrix = self.covariance_fn(inputs, inputs)
      # Tile locations so output has shape [units, batch_size]. Covariance will
      # broadcast to [units, batch_size, batch_size], and we perform
      # shape manipulations to get a random variable over [batch_size, units].
      loc = self.mean_fn(inputs)
      loc = tf.tile(loc[tf.newaxis], [self.units] + [1] * len(loc.shape))
    else:
      knn = self.covariance_fn(inputs, inputs)
      knm = self.covariance_fn(inputs, self.conditional_inputs)
      kmm = self.covariance_fn(self.conditional_inputs, self.conditional_inputs)
      kmm = tf.matrix_set_diag(
          kmm, tf.matrix_diag_part(kmm) + tf.keras.backend.epsilon())
      kmm_tril = tf.linalg.cholesky(kmm)
      kmm_tril_operator = tf.linalg.LinearOperatorLowerTriangular(kmm_tril)
      knm_operator = tf.linalg.LinearOperatorFullMatrix(knm)

      # TODO(trandustin): Vectorize linear algebra for multiple outputs. For
      # now, we do each separately and stack to obtain a locations Tensor of
      # shape [units, batch_size].
      loc = []
      for conditional_outputs_unit in tf.unstack(self.conditional_outputs,
                                                 axis=-1):
        center = conditional_outputs_unit - self.mean_fn(
            self.conditional_inputs)
        loc_unit = knm_operator.matvec(
            kmm_tril_operator.solvevec(kmm_tril_operator.solvevec(center),
                                       adjoint=True))
        loc.append(loc_unit)
      loc = tf.stack(loc) + self.mean_fn(inputs)[tf.newaxis]

      covariance_matrix = knn
      covariance_matrix -= knm_operator.matmul(
          kmm_tril_operator.solve(
              kmm_tril_operator.solve(knm, adjoint_arg=True), adjoint=True))

    covariance_matrix = tf.matrix_set_diag(
        covariance_matrix,
        tf.matrix_diag_part(covariance_matrix) + tf.keras.backend.epsilon())

    # Form a multivariate normal random variable with batch_shape units and
    # event_shape batch_size. Then make it be independent across the units
    # dimension. Then transpose its dimensions so it is [batch_size, units].
    random_variable = ed.MultivariateNormalFullCovariance(
        loc=loc, covariance_matrix=covariance_matrix)
    random_variable = ed.Independent(random_variable.distribution,
                                     reinterpreted_batch_ndims=1)
    bijector = tfp.bijectors.Inline(
        forward_fn=lambda x: tf.transpose(x, [1, 0]),
        inverse_fn=lambda y: tf.transpose(y, [1, 0]),
        forward_event_shape_fn=lambda input_shape: input_shape[::-1],
        forward_event_shape_tensor_fn=lambda input_shape: input_shape[::-1],
        inverse_log_det_jacobian_fn=lambda y: tf.cast(0, y.dtype),
        forward_min_event_ndims=2)
    random_variable = ed.TransformedDistribution(random_variable.distribution,
                                                 bijector=bijector)
    return random_variable