Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

remove flatten.using_ints, linalg_*, linear, log_softmax.int, logdet, special_* from xfail list #110985

Closed
wants to merge 4 commits into from
Closed

remove flatten.using_ints, linalg_*, linear, log_softmax.int, logdet, special_* from xfail list #110985

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a922fdb
remove flatten.using_ints, linalg_*, linear, log_softmax.int, logdet,…
guilhermeleobas Oct 10, 2023
b288bd8
Update on "remove flatten.using_ints, linalg_*, linear, log_softmax.i…
guilhermeleobas Oct 10, 2023
f9724a8
Update on "remove flatten.using_ints, linalg_*, linear, log_softmax.i…
guilhermeleobas Oct 11, 2023
26634de
Update on "remove flatten.using_ints, linalg_*, linear, log_softmax.i…
guilhermeleobas Oct 20, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion aten/src/ATen/functorch/BatchRulesBinaryOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -434,7 +434,6 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchBatched, m) {
BINARY_SCALAR_2(rsub, Tensor, Scalar);

BINARY_SCALAR_3_Tensor(special_xlog1py, other_scalar, self_scalar);
BINARY_SCALAR_3_Tensor(special_xlogy, other_scalar, self_scalar);
BINARY_SCALAR_3_Tensor(special_zeta, other_scalar, self_scalar);

VMAP_SUPPORT2(where, self, where_self_batch_rule);
Expand Down
10 changes: 10 additions & 0 deletions aten/src/ATen/functorch/BatchRulesDecompositions.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,7 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchBatchedDecomposition, m) {
OP_DECOMPOSE(fix);
OP_DECOMPOSE(fliplr);
OP_DECOMPOSE(flipud);
OP_DECOMPOSE2(flatten, using_ints);
OP_DECOMPOSE2(float_power, Tensor_Tensor);
OP_DECOMPOSE2(float_power, Tensor_Scalar);
OP_DECOMPOSE2(float_power, Scalar);
Expand Down Expand Up @@ -155,12 +156,14 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchBatchedDecomposition, m) {
OP_DECOMPOSE2(ldexp, Tensor);
OP_DECOMPOSE2(less_equal, Tensor );
OP_DECOMPOSE2(less, Tensor );
OP_DECOMPOSE(linear);
OP_DECOMPOSE(linalg_cond);
OP_DECOMPOSE(linalg_cholesky);
OP_DECOMPOSE(linalg_det);
OP_DECOMPOSE(linalg_eigvalsh);
OP_DECOMPOSE(linalg_eigvals);
OP_DECOMPOSE(linalg_inv);
OP_DECOMPOSE(linalg_lu_factor);
OP_DECOMPOSE(linalg_matmul);
OP_DECOMPOSE(linalg_matrix_norm);
OP_DECOMPOSE2(linalg_matrix_norm, str_ord);
Expand All @@ -172,12 +175,16 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchBatchedDecomposition, m) {
OP_DECOMPOSE(linalg_solve_ex);
OP_DECOMPOSE(linalg_svd);
OP_DECOMPOSE(linalg_svdvals);
OP_DECOMPOSE(linalg_pinv);
OP_DECOMPOSE(linalg_tensorinv);
OP_DECOMPOSE2(linalg_pinv, atol_rtol_float);
m.impl("linalg_vander", native::linalg_vander_symint);
OP_DECOMPOSE(cumprod_backward);
OP_DECOMPOSE(linalg_matrix_power);
OP_DECOMPOSE(linalg_vecdot);
OP_DECOMPOSE(log_sigmoid);
OP_DECOMPOSE(logdet);
OP_DECOMPOSE2(log_softmax, int);
OP_DECOMPOSE(_lu_with_info);
OP_DECOMPOSE(matmul);
OP_DECOMPOSE(matrix_H);
Expand Down Expand Up @@ -248,6 +255,9 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchBatchedDecomposition, m) {
OP_DECOMPOSE(special_psi);
OP_DECOMPOSE(special_round);
OP_DECOMPOSE(special_sinc);
OP_DECOMPOSE(special_xlogy);
OP_DECOMPOSE2(special_xlogy, other_scalar);
OP_DECOMPOSE2(special_xlogy, self_scalar);


m.impl("split.sizes", native::split_symint);
Expand Down
28 changes: 0 additions & 28 deletions aten/src/ATen/functorch/BatchRulesLinearAlgebra.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,21 +146,6 @@ Tensor baddbmm_decomp(
return out;
}

Tensor linear_decomp(
const Tensor& input, const Tensor& weight,
const c10::optional<Tensor>& bias_opt) {
auto result = input.matmul(weight.t());
if (bias_opt) {
// NB: It's too much work to figure out how to actually fuse the bias so
// we're not going to.
// TODO: if the result isn't batched but bias is, then we need to do the following.
// Otherwise, it can just be in-place. We should write a more nuanced
// decomposition rule
return result.add(*bias_opt);
}
return result;
}

Tensor addmm_decomp(const Tensor& self, const Tensor& mat1, const Tensor& mat2, const Scalar& beta, const Scalar& alpha) {
// Decomposition that is probably not very fast...
return at::add(self * beta, at::mm(mat1, mat2), alpha);
Expand Down Expand Up @@ -345,14 +330,6 @@ threeOutputs linalg_lu_factor_ex_batch_rule(
return std::make_tuple(std::get<0>(res), 0, std::get<1>(res), 0, std::get<2>(res), 0);
}

twoOutputs linalg_lu_factor_batch_rule(
const Tensor& A, c10::optional<int64_t> A_bdim, bool pivot) {
TORCH_CHECK(rankWithoutBatchDim(A, A_bdim) >= 2, "torch.lu_factor: Expected tensor with 2 or more dimensions. Got size: ", A.sizes(), " instead");
const auto A_ = moveBatchDimToFront(A, A_bdim);
const auto res = at::linalg_lu_factor(A_, pivot);
return std::make_tuple(std::get<0>(res), 0, std::get<1>(res), 0);
}

oneOutput matrix_exp_batch_rule(const Tensor& self, c10::optional<int64_t> self_bdim) {
TORCH_CHECK(rankWithoutBatchDim(self, self_bdim) >= 2, "torch.matrix_exp: The input tensor A must have at least 2 dimensions.");
const auto self_ = moveBatchDimToFront(self, self_bdim).contiguous(); // seems to be a bug
Expand Down Expand Up @@ -582,14 +559,11 @@ LINALG_CHECK_MATRIX_UNARY_TWO_OUT(linalg_cholesky_ex, linalg.cholesky);
LINALG_CHECK_MATRIX_UNARY_TWO_OUT(linalg_eig, linalg.eig);
LINALG_CHECK_MATRIX_UNARY_TWO_OUT(linalg_inv_ex, linalg.inv_ex);
LINALG_CHECK_MATRIX_UNARY_THREE_OUT(linalg_ldl_factor_ex, torch.linalg.ldl_factor_ex);
LINALG_CHECK_MATRIX_UNARY_ONE_OUT(linalg_pinv, linalg.pinv);
LINALG_CHECK_MATRIX_UNARY_ONE_OUT2(linalg_pinv, atol_rtol_float, linalg.pinv);
LINALG_CHECK_MATRIX_UNARY_TWO_OUT(linalg_qr, linalg.qr);
LINALG_CHECK_MATRIX_UNARY_TWO_OUT(linalg_slogdet, linalg.slogdet);
LINALG_CHECK_MATRIX_BINARY_ONE_OUT(linalg_solve_triangular, linalg.solve_triangular);

LINALG_CHECK_MATRIX_UNARY_TWO_OUT(geqrf, geqrf);
LINALG_CHECK_MATRIX_UNARY_ONE_OUT(logdet, logdet);
LINALG_CHECK_MATRIX_BINARY_TWO_OUT(triangular_solve, triangular_solve);
LINALG_CHECK_MATRIX_UNARY_THREE_OUT(_linalg_det, linalg.det);
LINALG_CHECK_MATRIX_UNARY_TWO_OUT(_linalg_eigh, linalg.eigh);
Expand All @@ -605,13 +579,11 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchBatched, m) {
VMAP_SUPPORT(dot, dot_batch_rule);
VMAP_SUPPORT(mv, mv_batch_rule);
VMAP_SUPPORT(mm, mm_batch_rule);
m.impl("linear", linear_decomp);
guilhermeleobas marked this conversation as resolved.
Show resolved Hide resolved
VMAP_SUPPORT(linalg_lu_solve, linalg_lu_solve_batch_rule);
VMAP_SUPPORT(linalg_householder_product, householder_product_batch_rule);
VMAP_SUPPORT(cholesky_solve, cholesky_solve_batch_rule); // custom dim error
VMAP_SUPPORT(linalg_lstsq, linalg_lstsq_batch_rule); // custom errors and sometimes empty return
VMAP_SUPPORT(linalg_lu_factor_ex, linalg_lu_factor_ex_batch_rule);
VMAP_SUPPORT(linalg_lu_factor, linalg_lu_factor_batch_rule);
guilhermeleobas marked this conversation as resolved.
Show resolved Hide resolved
VMAP_SUPPORT(linalg_matrix_exp, matrix_exp_batch_rule);
VMAP_SUPPORT(_linalg_solve_ex, solve_ex_batch_rule);
VMAP_SUPPORT(linalg_cross, cross_batch_rule);
Expand Down
1 change: 0 additions & 1 deletion aten/src/ATen/functorch/BatchRulesReduceOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -472,7 +472,6 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchBatched, m) {
m.impl("dist", dist_decomp);
REDUCTION_BOXED_ARGS(kthvalue, 2, KEEPDIM_CASE_VARIABLE, 3);
REDUCTION_BOXED_ARGS(linalg_vector_norm, 2, KEEPDIM_CASE_VARIABLE, 3);
REDUCTION_NO_KEEPDIM_ARG(log_softmax.int);
REDUCTION_NO_KEEPDIM_ARG(logcumsumexp);
REDUCTION_WITH_KEEPDIM_ARG(logsumexp);
m.impl("max", max_decomp);
Expand Down
1 change: 0 additions & 1 deletion aten/src/ATen/functorch/BatchRulesViews.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -556,7 +556,6 @@ std::tuple<Tensor,optional<int64_t>> triu_batch_rule(
}

TORCH_LIBRARY_IMPL(aten, FuncTorchBatched, m) {
m.impl("flatten.using_ints", static_cast<decltype(&ATEN_FN2(flatten, using_ints))>(native::flatten));
VMAP_SUPPORT(flip, flip_batch_rule);
m.impl("trace", trace_decomp);
VMAP_SUPPORT(tril, tril_batch_rule);
Expand Down
20 changes: 0 additions & 20 deletions test/functorch/test_vmap_registrations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,22 +15,12 @@
)

xfail_functorch_batched = {
"aten::flatten.using_ints",
"aten::is_nonzero",
"aten::item",
"aten::linalg_pinv",
"aten::linalg_pinv.atol_rtol_float",
"aten::linalg_slogdet",
"aten::linalg_lu_factor",
"aten::linear",
"aten::log_softmax.int",
"aten::logdet",
"aten::masked_select_backward",
"aten::one_hot",
"aten::silu_backward",
"aten::special_xlogy",
"aten::special_xlogy.other_scalar",
"aten::special_xlogy.self_scalar",
"aten::where",
}

Expand Down Expand Up @@ -85,7 +75,6 @@
"aten::fill_diagonal_",
"aten::fix_",
"aten::flatten.named_out_dim",
"aten::flatten.using_ints",
"aten::flatten.using_names",
"aten::flatten_dense_tensors",
"aten::float_power_.Scalar",
Expand Down Expand Up @@ -116,20 +105,14 @@
"aten::less_equal_.Tensor",
"aten::linalg_cond.p_str",
"aten::linalg_eigh.eigvals",
"aten::linalg_lu_factor",
"aten::linalg_matrix_rank",
"aten::linalg_matrix_rank.out_tol_tensor",
"aten::linalg_matrix_rank.tol_tensor",
"aten::linalg_pinv",
"aten::linalg_pinv.atol_rtol_float",
"aten::linalg_pinv.out_rcond_tensor",
"aten::linalg_pinv.rcond_tensor",
"aten::linalg_slogdet",
"aten::linalg_svd.U",
"aten::linalg_tensorsolve",
"aten::linear",
"aten::log_softmax.int",
"aten::logdet",
"aten::logsumexp.names",
"aten::lstm.data",
"aten::lstm.input",
Expand Down Expand Up @@ -213,9 +196,6 @@
"aten::special_shifted_chebyshev_polynomial_v.x_scalar",
"aten::special_shifted_chebyshev_polynomial_w.n_scalar",
"aten::special_shifted_chebyshev_polynomial_w.x_scalar",
"aten::special_xlogy",
"aten::special_xlogy.other_scalar",
"aten::special_xlogy.self_scalar",
"aten::square_",
"aten::sspaddmm",
"aten::std.correction_names",
Expand Down
Loading