Added alternative slogdet to return sign and logdet of det op

Author: tanish1729

pytensor/tensor/nlinalg.py

@@ -266,7 +266,10 @@ def __str__(self):
         return "SLogDet"
 
 
-slogdet = Blockwise(SLogDet())
+# slogdet = Blockwise(SLogDet())
+def slogdet(x):
+    det_val = det(x)
+    return ptm.sign(det_val), ptm.log(ptm.abs(det_val))
 
 
 class Eig(Op):

pytensor/tensor/rewriting/linalg.py

@@ -34,7 +34,6 @@
     inv,
     kron,
     pinv,
-    slogdet,
     svd,
 )
 from pytensor.tensor.rewriting.basic import (
@@ -785,43 +784,43 @@ def rewrite_det_blockdiag(fgraph, node):
     return [prod(det_sub_matrices)]
 
 
-@register_canonicalize
-@register_stabilize
-@node_rewriter([slogdet])
-def rewrite_slogdet_blockdiag(fgraph, node):
-    """
-    This rewrite simplifies the slogdet of a blockdiagonal matrix by extracting the individual sub matrices and returning the sign and logdet values computed using those
-
-    slogdet(block_diag(a,b,c,....)) = prod(sign(a), sign(b), sign(c),...), sum(logdet(a), logdet(b), logdet(c),....)
-
-    Parameters
-    ----------
-    fgraph: FunctionGraph
-        Function graph being optimized
-    node: Apply
-        Node of the function graph to be optimized
-
-    Returns
-    -------
-    list of Variable, optional
-        List of optimized variables, or None if no optimization was performed
-    """
-    # Check for inner block_diag operation
-    potential_block_diag = node.inputs[0].owner
-    if not (
-        potential_block_diag
-        and isinstance(potential_block_diag.op, Blockwise)
-        and isinstance(potential_block_diag.op.core_op, BlockDiagonal)
-    ):
-        return None
-
-    # Find the composing sub_matrices
-    sub_matrices = potential_block_diag.inputs
-    sign_sub_matrices, logdet_sub_matrices = zip(
-        *[slogdet(sub_matrices[i]) for i in range(len(sub_matrices))]
-    )
-
-    return [prod(sign_sub_matrices), sum(logdet_sub_matrices)]
+# @register_canonicalize
+# @register_stabilize
+# @node_rewriter([slogdet])
+# def rewrite_slogdet_blockdiag(fgraph, node):
+#     """
+#     This rewrite simplifies the slogdet of a blockdiagonal matrix by extracting the individual sub matrices and returning the sign and logdet values computed using those
+
+#     slogdet(block_diag(a,b,c,....)) = prod(sign(a), sign(b), sign(c),...), sum(logdet(a), logdet(b), logdet(c),....)
+
+#     Parameters
+#     ----------
+#     fgraph: FunctionGraph
+#         Function graph being optimized
+#     node: Apply
+#         Node of the function graph to be optimized
+
+#     Returns
+#     -------
+#     list of Variable, optional
+#         List of optimized variables, or None if no optimization was performed
+#     """
+#     # Check for inner block_diag operation
+#     potential_block_diag = node.inputs[0].owner
+#     if not (
+#         potential_block_diag
+#         and isinstance(potential_block_diag.op, Blockwise)
+#         and isinstance(potential_block_diag.op.core_op, BlockDiagonal)
+#     ):
+#         return None
+
+#     # Find the composing sub_matrices
+#     sub_matrices = potential_block_diag.inputs
+#     sign_sub_matrices, logdet_sub_matrices = zip(
+#         *[slogdet(sub_matrices[i]) for i in range(len(sub_matrices))]
+#     )
+
+#     return [prod(sign_sub_matrices), sum(logdet_sub_matrices)]
 
 
 @register_canonicalize
@@ -858,39 +857,39 @@ def rewrite_diag_kronecker(fgraph, node):
     return [outer_prod_as_vector]
 
 
-@register_canonicalize
-@register_stabilize
-@node_rewriter([slogdet])
-def rewrite_slogdet_kronecker(fgraph, node):
-    """
-    This rewrite simplifies the slogdet of a kronecker-structured matrix by extracting the individual sub matrices and returning the sign and logdet values computed using those
-
-    Parameters
-    ----------
-    fgraph: FunctionGraph
-        Function graph being optimized
-    node: Apply
-        Node of the function graph to be optimized
-
-    Returns
-    -------
-    list of Variable, optional
-        List of optimized variables, or None if no optimization was performed
-    """
-    # Check for inner kron operation
-    potential_kron = node.inputs[0].owner
-    if not (potential_kron and isinstance(potential_kron.op, KroneckerProduct)):
-        return None
-
-    # Find the matrices
-    a, b = potential_kron.inputs
-    signs, logdets = zip(*[slogdet(a), slogdet(b)])
-    sizes = [a.shape[-1], b.shape[-1]]
-    prod_sizes = prod(sizes, no_zeros_in_input=True)
-    signs_final = [signs[i] ** (prod_sizes / sizes[i]) for i in range(2)]
-    logdet_final = [logdets[i] * prod_sizes / sizes[i] for i in range(2)]
-
-    return [prod(signs_final, no_zeros_in_input=True), sum(logdet_final)]
+# @register_canonicalize
+# @register_stabilize
+# @node_rewriter([slogdet])
+# def rewrite_slogdet_kronecker(fgraph, node):
+#     """
+#     This rewrite simplifies the slogdet of a kronecker-structured matrix by extracting the individual sub matrices and returning the sign and logdet values computed using those
+
+#     Parameters
+#     ----------
+#     fgraph: FunctionGraph
+#         Function graph being optimized
+#     node: Apply
+#         Node of the function graph to be optimized
+
+#     Returns
+#     -------
+#     list of Variable, optional
+#         List of optimized variables, or None if no optimization was performed
+#     """
+#     # Check for inner kron operation
+#     potential_kron = node.inputs[0].owner
+#     if not (potential_kron and isinstance(potential_kron.op, KroneckerProduct)):
+#         return None
+
+#     # Find the matrices
+#     a, b = potential_kron.inputs
+#     signs, logdets = zip(*[slogdet(a), slogdet(b)])
+#     sizes = [a.shape[-1], b.shape[-1]]
+#     prod_sizes = prod(sizes, no_zeros_in_input=True)
+#     signs_final = [signs[i] ** (prod_sizes / sizes[i]) for i in range(2)]
+#     logdet_final = [logdets[i] * prod_sizes / sizes[i] for i in range(2)]
+
+#     return [prod(signs_final, no_zeros_in_input=True), sum(logdet_final)]
 
 
 @register_canonicalize