Function HouseRow

tensorflow/compiler/xla/client/lib/svd.cc:109–165

Householder reflection on the trailing elements of a vector. H = I - beta * [1, v]' * [1, v] H * x = [..., xnorm, 0, ..., 0] ..., j, j + 1, ..., n def house(x, j, eps): sigma = np.linalg.norm(x[(j + 1):]) v = np.zeros_like(x) v[(j + 1):] = x[(j + 1):] if sigma < eps: beta = 0 else: mu = sigma * np.sqrt((x[j]/sigma)**2 + 1) if x[j] <= 0: v[j] = x[j] - mu else: v[j] = -sigma / (x[j] + mu) * sigma

// A[i, j:] * H = [sigma, 0, 0, ..., 0]
//
StatusOr<HouseHolderResult> HouseRow(XlaOp a, XlaOp i, XlaOp j, XlaOp eps,
                                     PrecisionConfig::Precision precision) {
  XlaBuilder* builder = a.builder();
  TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a));
  const int64 num_dims = a_shape.rank();
  const int64 n = ShapeUtil::GetDimension(a_shape, -1);
  XlaOp zero = ScalarLike(i, 0);
  XlaOp x = DynamicSliceInMinorDims(a, {i, zero}, {1, n});

  const int64 num_batch_dims = num_dims - 2;
  std::vector<int64> batch_dims(num_batch_dims);
  for (int k = 0; k < num_batch_dims; ++k) {
    batch_dims[k] = ShapeUtil::GetDimension(a_shape, k);
  }

  TF_ASSIGN_OR_RETURN(Shape x_shape, builder->GetShape(x));
  auto idx = Iota(builder, ShapeUtil::MakeShape(S32, x_shape.dimensions()),
                  num_dims - 1);
  auto zeros = ZerosLike(x);
  auto v = Select(Gt(idx, j), x, zeros);

  auto one = ScalarLike(v, 1.0);

  auto sigma =
      Sqrt(Reduce(Square(v), ScalarLike(v, 0.0),
                  CreateScalarAddComputation(x_shape.element_type(), builder),
                  {num_dims - 1}));

  std::vector<int64> broadcast_dims(num_dims - 1);
  std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0);
  auto x_0j = DynamicSliceInMinorDims(x, {zero, j}, {1, 1});
  auto mu = Mul(sigma, Sqrt(Square(Div(x_0j, sigma, broadcast_dims)) + one),
                broadcast_dims);

  auto v_0j = Select(
      Le(x_0j, ScalarLike(x_0j, 0.0)), Sub(x_0j, mu),
      -Mul(sigma, Div(sigma, Add(x_0j, mu), broadcast_dims), broadcast_dims));

  auto beta = Div(ScalarLike(v_0j, 2.0),
                  (Square(Div(sigma, v_0j, broadcast_dims)) + one));

  v = Select(
      BroadcastInDim(Lt(sigma, eps), x_shape.dimensions(), broadcast_dims), v,
      v / v_0j);
  v = Select(Eq(idx, j), zeros + one, v);

  beta = Select(Lt(Add(sigma, ZerosLike(beta), broadcast_dims), eps),
                ZerosLike(beta), beta);

  HouseHolderResult result;
  result.v = v;
  result.beta = beta;
  result.a = Sub(a, Mul(beta, BatchDot(BatchDot(a, false, v, true, precision),
                                       v, precision)));

  return result;
}