[batch-rule] householder_product

Original code written by @kshitij12345 in
#322,
this PR is just a rebase onto main

Author: zou3519

codegen/gen_vmap_plumbing.py

@@ -274,6 +274,7 @@ def gen_vmap_plumbing(native_function: NativeFunction) -> str:
     'gcd',
     'igamma',
     'igammac',
+    'linalg_householder_product',
     'logaddexp',
     'logaddexp2',
     'lcm',

functorch/csrc/BatchRulesDecompositions.cpp

@@ -217,6 +217,7 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   OP_DECOMPOSE2(vsplit, int);
   OP_DECOMPOSE2(vsplit, array);
   OP_DECOMPOSE(vstack);
+  OP_DECOMPOSE(orgqr);
   OP_DECOMPOSE2(unflatten, int);
   OP_DECOMPOSE(_convolution_double_backward);
   OP_DECOMPOSE(conv_transpose1d);

functorch/csrc/BatchRulesLinearAlgebra.cpp

@@ -162,6 +162,20 @@ void _linalg_check_errors_batch_rule(const Tensor& info, optional<int64_t> info_
   at::_linalg_check_errors(info_, api_name, false);
 }
 
+std::tuple<Tensor, c10::optional<int64_t>>
+householder_product_batch_rule(const Tensor &input, c10::optional<int64_t> input_bdim,
+                               const Tensor &tau, c10::optional<int64_t> tau_bdim)
+{
+  auto input_ = moveBatchDimToFront(input, input_bdim);
+  auto tau_ = moveBatchDimToFront(tau, tau_bdim);
+
+  auto batch_size = get_bdim_size2(input, input_bdim, tau, tau_bdim);
+
+  input_ = ensure_has_bdim(input_, input_bdim.has_value(), batch_size);
+  tau_ = ensure_has_bdim(tau_, tau_bdim.has_value(), batch_size);
+  return std::make_tuple(at::linalg_householder_product(input_, tau_), 0);
+}
+
 TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   VMAP_SUPPORT(bmm, bmm_batch_rule);
   m.impl("addmv", addmv_decomp);
@@ -172,6 +186,7 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   VMAP_SUPPORT(mv, mv_batch_rule);
   VMAP_SUPPORT(mm, mm_batch_rule);
   m.impl("linear", linear_decomp);
+  VMAP_SUPPORT(linalg_householder_product, householder_product_batch_rule);
 
   VMAP_SUPPORT(_linalg_check_errors, _linalg_check_errors_batch_rule);
 

functorch/csrc/VmapGeneratedPlumbing.h

@@ -5431,6 +5431,24 @@ at::Tensor logdet_generated_plumbing(const at::Tensor & self) {
   return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
 }
 template <typename batch_rule_t, batch_rule_t batch_rule>
+at::Tensor linalg_householder_product_generated_plumbing(const at::Tensor & input, const at::Tensor & tau) {
+  c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
+  auto maybe_layer = maybeCurrentDynamicLayer();
+  TORCH_INTERNAL_ASSERT(maybe_layer.has_value());
+  int64_t cur_level = maybe_layer->layerId();
+  if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(tau, cur_level)) {
+    return at::_ops::linalg_householder_product::call(input, tau);
+  }
+  Tensor input_value;
+  optional<int64_t> input_bdim;
+  std::tie(input_value, input_bdim) = unwrapTensorAtLevel(input, cur_level);
+  Tensor tau_value;
+  optional<int64_t> tau_bdim;
+  std::tie(tau_value, tau_bdim) = unwrapTensorAtLevel(tau, cur_level);
+  auto results = batch_rule(input_value, input_bdim, tau_value, tau_bdim);
+  return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
+}
+template <typename batch_rule_t, batch_rule_t batch_rule>
 at::Tensor linalg_pinv_generated_plumbing(const at::Tensor & self, double rcond, bool hermitian) {
   c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
   auto maybe_layer = maybeCurrentDynamicLayer();

test/test_vmap.py

@@ -3199,7 +3199,6 @@ def test_vmap_exhaustive(self, device, dtype, op):
         xfail('linalg.cholesky'),
         xfail('linalg.eigvals'),
         xfail('linalg.eigvalsh'),
-        xfail('linalg.householder_product'),
         xfail('linalg.inv'),
         xfail('linalg.lstsq'),
         xfail('linalg.lstsq', 'grad_oriented'),