Method Compile

tensorflow/compiler/tf2xla/kernels/xla_svd_op.cc:61–87

    OP_REQUIRES_OK(ctx, ctx->GetAttr("full_matrices", &full_matrices_));
  }
  void Compile(XlaOpKernelContext* ctx) override {
    const TensorShape input_shape = ctx->InputShape("input");
    int m = input_shape.dim_size(input_shape.dims() - 2);
    int n = input_shape.dim_size(input_shape.dims() - 1);
    // This is based on heuristics that approx log(n) sweep updates are needed.
    // Note: the heuristics provides no theoretical guarantee, max_iter=100 and
    // epsilon should be used to determine exit condition.
    int max_iter = 2 * tensorflow::Log2Ceiling(std::max(m, n));
    auto result = xla::SVD(ctx->Input(0), max_iter, 1e-6);
    ctx->SetOutput(0, result.d);
    if (compute_uv_) {
      int p = std::min(m, n);
      if (!full_matrices_) {
        if (p < m) {
          result.u = xla::SliceInMinorDims(result.u, {0, 0}, {m, p});
        }
        if (p < n) {
          result.v = xla::SliceInMinorDims(result.v, {0, 0}, {n, p});
        }
      }
      ctx->SetOutput(1, result.u);
      ctx->SetOutput(2, result.v);
    } else {
      ctx->SetOutput(1, xla::ScalarLike(ctx->Input(0), 0.0));
      ctx->SetOutput(2, xla::ScalarLike(ctx->Input(0), 0.0));
    }
  }

 private:
  bool compute_uv_;