Added batch rule for LU op (#326)

vfdev-5 · web-flow · commit 03a9fcd8fa0c · 2021-12-13T17:23:03.000-05:00
Description: - Added batch rule for LU op - Updated tests Related to #240
diff --git a/functorch/csrc/BatchRulesLinearAlgebra.cpp b/functorch/csrc/BatchRulesLinearAlgebra.cpp
@@ -184,6 +184,7 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   VARIADIC_BDIMS_BOXED(triangular_solve);
 
   VARIADIC_BDIMS_BOXED(_det_lu_based_helper);
+  VARIADIC_BDIMS_BOXED(_lu_with_info);
 }
 }}
 
diff --git a/test/test_vmap.py b/test/test_vmap.py
@@ -3188,7 +3188,6 @@ def test_vmap_exhaustive(self, device, dtype, op):
         xfail('linalg.solve'),
         xfail('linalg.svdvals'),
         xfail('linalg.tensorinv'),
-        xfail('lu'),
         xfail('lu_solve'),
         xfail('lu_unpack'),
         xfail('masked_fill'),

Original file line number	Diff line number	Diff line change
`@@ -184,6 +184,7 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {`
`184`	`184`	`VARIADIC_BDIMS_BOXED(triangular_solve);`
`185`	`185`
`186`	`186`	`VARIADIC_BDIMS_BOXED(_det_lu_based_helper);`
	`187`	`+ VARIADIC_BDIMS_BOXED(_lu_with_info);`
`187`	`188`	`}`
`188`	`189`	`}}`
`189`	`190`