Function HouseCol

tensorflow/compiler/xla/client/lib/svd.cc:174–232

Householder reflection on the trailing elements of a col of a matrix. After applying it on the matrix, all elements in [(i+1):, j] become zeros, i.e., H = I - beta * [1; v] * [1; v]', then, H * A[i:, j] = [xnorm, 0, 0, ..., 0]

// H * A[i:, j] = [xnorm, 0, 0, ..., 0]
//
StatusOr<HouseHolderResult> HouseCol(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 m = ShapeUtil::GetDimension(a_shape, -2);
  XlaOp zero = ScalarLike(i, 0);
  XlaOp x = DynamicSliceInMinorDims(a, {zero, j}, {m, 1});

  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 - 2);
  auto zeros = ZerosLike(x);
  auto v = Select(Gt(idx, i), 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 - 2}));

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

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

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

  v = Select(
      BroadcastInDim(Lt(sigma, eps), x_shape.dimensions(), broadcast_dims), v,
      v / v_0i);
  v = Select(Eq(idx, i), 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(v, false, BatchDot(v, true, a, false, precision),
                            false, precision)));

  return result;