Added addr op decomposition

functorch/csrc/BatchRulesBinaryOps.cpp

@@ -155,6 +155,14 @@ std::tuple<Tensor,optional<int64_t>> masked_select_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;
+}
+
 TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
 #define BINARY_POINTWISE2(op, overload) \
   VMAP_SUPPORT(#op"."#overload, BINARY_POINTWISE_BATCH_RULE(ATEN_FN2(op, overload)));
@@ -193,6 +201,7 @@ 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);

test/test_ops.py

@@ -570,7 +570,6 @@ def test_vmapjvp(self, device, dtype, op):
     @skipOps('TestOperators', 'test_vmapvjp_has_batch_rule', vmapvjp_fail.union({
         xfail('view_as_complex'),
         xfail('__getitem__'),
-        xfail('addr'),
         xfail('cdist'),
         xfail('cholesky'),
         xfail('clamp'),

test/test_vmap.py

@@ -3099,7 +3099,6 @@ def test_vmap_exhaustive(self, device, dtype, op):
 
     @ops(functorch_lagging_op_db + additional_op_db, allowed_dtypes=(torch.float,))
     @skipOps('TestVmapOperatorsOpInfo', 'test_op_has_batch_rule', vmap_fail.union({
-        xfail('addr'),
         xfail('cdist'),
         xfail('complex'),
         xfail('copysign'),