Replace addr_decomp with the one in PyTorch Core, issue #833 (#836)

* Replace addr_decomp with the one in PyTorch Core, issue #833

* only allow OpOverload to be passed to _register_jit_decomposition_bypass_script

* change variable name to vmap_decompositions_lib

* change function name to _register_python_decomposition_vmap

Author: yushangdi

functorch/_src/eager_transforms.py

@@ -1305,6 +1305,20 @@ def get_function_def(sig):
     torch.jit._register_decomposition(decomp, graph)
 
 
+# use an alternate way to register an operator into the decomposition table
+# _register_jit_decomposition doesn't work for some operators, e.g. addr,
+#  because the Tensor types generated cannot be unioned by torchscript
+# decomp should be type OpOverload
+vmap_decompositions_lib = torch.library.Library("aten", "IMPL", "FuncTorchBatched")
+
+
+def _register_python_decomposition_vmap(decomp):
+    if decomp in decomposition_table:
+        vmap_decompositions_lib.impl(decomp, decomposition_table[decomp])
+    else:
+        raise RuntimeError(f"could not find decomposition for {decomp}")
+
+
 _register_jit_decomposition(torch.ops.aten.trace.default)
 _register_jit_decomposition(torch.ops.aten.nll_loss_backward.default)
 _register_jit_decomposition(torch.ops.aten.nll_loss2d_backward.default)
@@ -1316,3 +1330,4 @@ def get_function_def(sig):
 _register_jit_decomposition(torch.ops.aten.binary_cross_entropy_backward.default)
 _register_jit_decomposition(torch.ops.aten.binary_cross_entropy.default)
 _register_jit_decomposition(torch.ops.aten.native_layer_norm_backward.default)
+_register_python_decomposition_vmap(torch.ops.aten.addr.default)

functorch/csrc/BatchRulesBinaryOps.cpp

@@ -264,14 +264,6 @@ std::tuple<Tensor,optional<int64_t>> masked_select_backward_batch_rule(
   return std::make_tuple(result, 0);
 }
 
-Tensor addr_decomposition(
-    const Tensor& self, const Tensor& vec1, const Tensor& vec2,
-    const Scalar& beta, const Scalar& alpha) {
-
-  auto outer = alpha * vec1.unsqueeze(-1) * vec2.unsqueeze(-2);
-  return self * beta + outer;
-}
-
 std::tuple<Tensor,optional<int64_t>> cdist_backward_batch_rule(
     const Tensor& grad, optional<int64_t> grad_bdim,
     const Tensor& x1, optional<int64_t> x1_bdim,
@@ -369,7 +361,6 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   BINARY_SCALAR_2(add, Tensor, Scalar);
   POINTWISE_BOXED(addcdiv);
   POINTWISE_BOXED(addcmul);
-  m.impl("addr", addr_decomposition);
   BINARY_POINTWISE(atan2);
   BINARY_SCALAR_2(bitwise_and, Tensor, Scalar);
   BINARY_POINTWISE2(bitwise_or, Tensor);