Add cholesky_ex to list of supported functions in compiler (#1570)

Summary:
Pull Request resolved: #1570

Xitong's latest version of the GEP model uses a different version of the Cholesky operator that was not yet recognized by the compiler. The `cholesky_ex` pytorch API returns a tuple consisting of the result and a tensor containing the indices of the inputs which were invalid inputs.

We have no easy way to replicate this behavior in BMG, so instead we'll have the compiler emulate the pytorch behavior, but assume that there are no errors. It returns a tuple of the graph node and a zero tensor.

Reviewed By: gafter

Differential Revision: D38017818

fbshipit-source-id: e7a010d1e9f493ddb61c4a6c162841b81507117e

Author: ericlippert

src/beanmachine/ppl/compiler/special_function_caller.py

@@ -444,6 +444,7 @@ def __init__(self, bmg: BMGraphBuilder) -> None:
             torch.bitwise_right_shift: self._torch_bitwise_right_shift,
             torch.Tensor.cholesky: self._torch_cholesky,
             torch.linalg.cholesky: self._torch_cholesky,
+            torch.linalg.cholesky_ex: self._torch_cholesky_ex,
             torch.Tensor.div: self._torch_div,
             torch.div: self._torch_div,
             torch.Tensor.divide: self._torch_div,
@@ -930,6 +931,23 @@ def _torch_cholesky(
         # TODO: What to do with upper?
         return self._bmg.add_cholesky(input)
 
+    def _torch_cholesky_ex(
+        self,
+        input: BMGNode,
+        upper: Optional[BMGNode] = None,
+        check_errors: Optional[BMGNode] = None,
+        out: Any = None,
+    ) -> BMGNode:
+        # TODO: What to do with upper and check_errors?
+        # cholesky_ex returns a named tuple (L, info) where
+        # L is the result matrix and info is a tensor containing
+        # an index saying which input element was not
+        # positive-definite. We pretend that this operation always
+        # succeeds and return a graph node and a zero error index.
+        return torch.return_types.linalg_cholesky_ex(  # pyre-ignore
+            (self._bmg.add_cholesky(input), torch.tensor(0))
+        )
+
     def _torch_transpose(
         self,
         input: BMGNode,

src/beanmachine/ppl/compiler/tests/cholesky_test.py

@@ -56,6 +56,15 @@ def cholesky4():
     return t.cholesky()
 
 
+@bm.functional
+def cholesky5():
+    n0 = norm(0) * norm(0)
+    n1 = norm(1) * norm(1)
+    t = tensor([[n0, 0.0], [0.0, n1]])
+    L, _ = torch.linalg.cholesky_ex(t)
+    return L
+
+
 # TODO: Test with a non-square matrix, should give an error.
 
 
@@ -101,6 +110,8 @@ def test_cholesky(self) -> None:
         self.assertEqual(expected.strip(), observed.strip())
         observed = BMGInference().to_dot([cholesky3()], {})
         self.assertEqual(expected.strip(), observed.strip())
+        observed = BMGInference().to_dot([cholesky5()], {})
+        self.assertEqual(expected.strip(), observed.strip())
 
         expected = """
 digraph "graph" {