Numpy det and slogdet operators (#15861)

ckt624 · reminisce · commit 295fc145510a · 2019-10-07T10:48:52.000-07:00
* Add alias.

Add tests.

Add slogdet tests.

Add docs

Change shapes

Change tests.

Change slogdet tests

Change style.

* Fix

* Fix
diff --git a/python/mxnet/_numpy_op_doc.py b/python/mxnet/_numpy_op_doc.py
@@ -20,6 +20,128 @@
 """Doc placeholder for numpy ops with prefix _np."""
 
 
+def _np__linalg_det(a):
+    """
+    det(a)
+
+    Compute the determinant of an array.
+
+    Parameters
+    ----------
+    a : (..., M, M) ndarray
+        Input array to compute determinants for.
+
+    Returns
+    -------
+    det : (...) ndarray
+        Determinant of `a`.
+
+    See Also
+    --------
+    slogdet : Another way to represent the determinant, more suitable
+    for large matrices where underflow/overflow may occur.
+
+    Notes
+    -----
+
+    Broadcasting rules apply, see the `numpy.linalg` documentation for
+    details.
+
+    The determinant is computed via LU factorization using the LAPACK
+    routine z/dgetrf.
+
+    Examples
+    --------
+    The determinant of a 2-D array [[a, b], [c, d]] is ad - bc:
+
+    >>> a = np.array([[1, 2], [3, 4]])
+    >>> np.linalg.det(a)
+    -2.0
+
+    Computing determinants for a stack of matrices:
+
+    >>> a = np.array([ [[1, 2], [3, 4]], [[1, 2], [2, 1]], [[1, 3], [3, 1]] ])
+    >>> a.shape
+    (3, 2, 2)
+    >>> np.linalg.det(a)
+    array([-2., -3., -8.])
+    """
+    pass
+
+
+def _np__linalg_slogdet(a):
+    """
+    slogdet(a)
+
+    Compute the sign and (natural) logarithm of the determinant of an array.
+
+    If an array has a very small or very large determinant, then a call to
+    `det` may overflow or underflow. This routine is more robust against such
+    issues, because it computes the logarithm of the determinant rather than
+    the determinant itself.
+
+    Parameters
+    ----------
+    a : (..., M, M) ndarray
+        Input array, has to be a square 2-D array.
+
+    Returns
+    -------
+    sign : (...) ndarray
+        A number representing the sign of the determinant. For a real matrix,
+        this is 1, 0, or -1.
+    logdet : (...) array_like
+        The natural log of the absolute value of the determinant.
+
+    If the determinant is zero, then `sign` will be 0 and `logdet` will be
+    -Inf. In all cases, the determinant is equal to ``sign * np.exp(logdet)``.
+
+    See Also
+    --------
+    det
+
+    Notes
+    -----
+
+    Broadcasting rules apply, see the `numpy.linalg` documentation for
+    details.
+
+    The determinant is computed via LU factorization using the LAPACK
+    routine z/dgetrf.
+
+
+    Examples
+    --------
+    The determinant of a 2-D array ``[[a, b], [c, d]]`` is ``ad - bc``:
+
+    >>> a = np.array([[1, 2], [3, 4]])
+    >>> (sign, logdet) = np.linalg.slogdet(a)
+    >>> (sign, logdet)
+    (-1., 0.69314718055994529)
+    >>> sign * np.exp(logdet)
+    -2.0
+
+    Computing log-determinants for a stack of matrices:
+
+    >>> a = np.array([ [[1, 2], [3, 4]], [[1, 2], [2, 1]], [[1, 3], [3, 1]] ])
+    >>> a.shape
+    (3, 2, 2)
+    >>> sign, logdet = np.linalg.slogdet(a)
+    >>> (sign, logdet)
+    (array([-1., -1., -1.]), array([ 0.69314718,  1.09861229,  2.07944154]))
+    >>> sign * np.exp(logdet)
+    array([-2., -3., -8.])
+
+    This routine succeeds where ordinary `det` does not:
+
+    >>> np.linalg.det(np.eye(500) * 0.1)
+    0.0
+    >>> np.linalg.slogdet(np.eye(500) * 0.1)
+    (1., -1151.2925464970228)
+    """
+    pass
+
+
 def _np_ones_like(a):
     """
     Return an array of ones with the same shape and type as a given array.
diff --git a/src/operator/tensor/la_op.cc b/src/operator/tensor/la_op.cc
@@ -941,6 +941,7 @@ NNVM_REGISTER_OP(_backward_linalg_inverse)
 
 NNVM_REGISTER_OP(_linalg_det)
 .add_alias("linalg_det")
+.add_alias("_np__linalg_det")
 .describe(R"code(Compute the determinant of a matrix.
 Input is a tensor *A* of dimension *n >= 2*.
 
@@ -991,6 +992,7 @@ NNVM_REGISTER_OP(_backward_linalg_det)
 .set_attr<FCompute>("FCompute<cpu>", LaOpDetBackward<cpu, 1, det_backward>);
 
 NNVM_REGISTER_OP(_linalg_slogdet)
+.add_alias("_np__linalg_slogdet")
 .add_alias("linalg_slogdet")
 .describe(R"code(Compute the sign and log of the determinant of a matrix.
 Input is a tensor *A* of dimension *n >= 2*.
diff --git a/src/operator/tensor/la_op.cu b/src/operator/tensor/la_op.cu
@@ -100,12 +100,14 @@ NNVM_REGISTER_OP(_backward_linalg_inverse)
 .set_attr<FCompute>("FCompute<gpu>", LaOpBackward<gpu, 2, 2, 2, 1, inverse_backward>);
 
 NNVM_REGISTER_OP(_linalg_det)
+.add_alias("_np__linalg_det")
 .set_attr<FCompute>("FCompute<gpu>", LaOpDetForward<gpu, 1, det>);
 
 NNVM_REGISTER_OP(_backward_linalg_det)
 .set_attr<FCompute>("FCompute<gpu>", LaOpDetBackward<gpu, 1, det_backward>);
 
 NNVM_REGISTER_OP(_linalg_slogdet)
+.add_alias("_np__linalg_slogdet")
 .set_attr<FCompute>("FCompute<gpu>", LaOpDetForward<gpu, 2, slogdet>);
 
 NNVM_REGISTER_OP(_backward_linalg_slogdet)
diff --git a/src/operator/tensor/la_op.h b/src/operator/tensor/la_op.h
@@ -26,6 +26,7 @@
 #define MXNET_OPERATOR_TENSOR_LA_OP_H_
 
 #include <mxnet/operator_util.h>
+#include <mxnet/imperative.h>
 #include <vector>
 #include <algorithm>
 #include "../mshadow_op.h"
@@ -428,7 +429,11 @@ inline bool DetShape(const nnvm::NodeAttrs& attrs,
   CHECK_EQ(in[ndim-2], in[ndim-1]) << "Input A's last two dimension must be equal";
   mxnet::TShape out;
   if (ndim == 2) {
-    out = mxnet::TShape(1, 1);
+    if (Imperative::Get()->is_np_shape()) {
+      out = mxnet::TShape(0, 1);
+    } else {
+      out = mxnet::TShape(1, 1);
+    }
   } else {
     out = mxnet::TShape(in.begin(), in.end() - 2);
   }
diff --git a/tests/python/unittest/test_numpy_op.py b/tests/python/unittest/test_numpy_op.py
