Function House

tensorflow/compiler/xla/client/lib/qr.cc:75–117

Computes a Householder reflection of the form: H = I - tau v v.T. such that H . ( x1 ) = ( x1 ) ( x2 ) = ( x2 ) ( ... ) = ( ... ) ( xk ) = ( beta ) ( ... ) ( 0 ) ( ... ) ( 0 ) Unlike the usual formulation, we allow the caller to supply 'k' rather than only providing the relevant part of 'x' to maintain XLA's static shape invariant. In addition, the implementation supports batchin

// TODO(phawkins): LAPACK's xLARFG implementation has code for handling
// overflows in the norm/beta calculations. Perhaps do the same here.
Status House(XlaOp x, XlaOp k, absl::Span<const int64> batch_dims,
             const int64 m, XlaOp* v, XlaOp* tau, XlaOp* beta) {
  XlaBuilder* const builder = x.builder();
  TF_ASSIGN_OR_RETURN(Shape x_shape, builder->GetShape(x));
  const PrimitiveType type = x_shape.element_type();

  std::vector<int64> batch_dim_ids(batch_dims.size());
  std::iota(batch_dim_ids.begin(), batch_dim_ids.end(), 0);
  const int64 minor_dim = batch_dims.size();

  XlaOp zero = ScalarLike(x, 0.0);
  XlaOp one = ScalarLike(x, 1.0);

  // alpha = x[k]
  XlaOp alpha = Reshape(DynamicSliceInMinorDims(x, {k}, {1}), batch_dims);

  // Compute x[k+1:] (padded with zeros in elements 0..k)
  XlaOp iota = Iota(builder, S32, m);
  XlaOp x_after_k = Mul(x, ConvertElementType(Gt(iota, k), type),
                        /*broadcast_dimensions=*/{minor_dim});

  // sigma = np.dot(x[k+1:], x[k+1:])
  auto sigma = Reduce(x_after_k * x_after_k, zero,
                      CreateScalarAddComputation(type, builder), {minor_dim});
  // mu = np.sqrt(x[k]*x[k] + sigma)
  auto mu = Sqrt(Square(alpha) + sigma);

  auto sigma_is_zero = Eq(sigma, zero);

  *beta = Select(sigma_is_zero, alpha, Select(Lt(alpha, zero), one, -one) * mu);
  *tau = Select(sigma_is_zero, Broadcast(zero, batch_dims),
                (*beta - alpha) / *beta);
  auto divisor =
      Select(sigma_is_zero, Broadcast(one, batch_dims), alpha - *beta);

  auto e_k = Broadcast(ConvertElementType(Eq(iota, k), type),
                       std::vector<int64>(batch_dims.size(), 1));

  // Form v as [0, 0, ..., 1] ++ x[k+1:] / divisor
  // If sigma is zero, x[k+1:] is zero, so use any non-zero divisor.
  *v = e_k + Div(x_after_k, divisor, /*broadcast_dimensions=*/batch_dim_ids);
  return Status::OK();
}

// Householder QR decomposition. Algorithm 5.2.1 from Golub and Van
// Loan "Matrix Computations", 4th Edition. This is an unblocked implementation