Add MXNet backend

pyhf/tensor/mxnet_backend.py

@@ -0,0 +1,183 @@
+import mxnet as mx
+from mxnet import nd
+import logging
+log = logging.getLogger(__name__)
+
+
+class mxnet_backend(object):
+    """Backend for MXNet"""
+
+    def __init__(self, **kwargs):
+        self.session = kwargs.get('session')
+
+    def tolist(self, tensor_in):
+        """
+        Convert a tensor to a list
+
+        Args:
+            tensor_in: MXNet tensor
+
+        Returns:
+            The possibly nested list of tensor elements.
+        """
+        tensor_in = self.astensor(tensor_in)
+        return tensor_in.asnumpy().tolist()
+
+    def outer(self, tensor_in_1, tensor_in_2):
+        """
+        The outer product of two tensors: u.v^T
+
+        Args:
+            tensor_in_1: tensor object
+            tensor_in_2: tensor object
+
+        Returns:
+            MXNet NDArray: The outer product
+        """
+        tensor_in_1 = self.astensor(tensor_in_1)
+        tensor_in_2 = self.astensor(tensor_in_2)
+        pass
+
+    def astensor(self, tensor_in):
+        """
+        Convert a tensor to an MXNet NDArray
+
+        Args:
+            tensor_in: tensor object
+
+        Returns:
+            MXNet NDArray: a multi-dimensional, fixed-size homogenous array.
+        """
+        return nd.array(tensor_in)
+
+    def sum(self, tensor_in, axis=None):
+        """
+        Compute the sum of array elements over given axes.
+
+        Args:
+            tensor_in: tensor object
+            axis: the axes over which to sum
+
+        Returns:
+            MXNet NDArray: ndarray of the sum over the axes
+        """
+        tensor_in = self.astensor(tensor_in)
+        if axis is None or tensor_in.shape == nd.Size([]):
+            return nd.sum(tensor_in)
+        else:
+            return nd.sum(tensor_in, axis)
+
+    def product(self, tensor_in, axis=None):
+        pass
+
+    def ones(self, shape):
+        """
+        A new array filled with all ones, with the given shape.
+
+        Args:
+            shape: the shape of the array
+
+        Returns:
+            MXNet NDArray: ndarray of 1's with given shape
+        """
+        return nd.ones(shape)
+
+    def zeros(self, shape):
+        """
+        A new array filled with all zeros, with the given shape.
+
+        Args:
+            shape: the shape of the array
+
+        Returns:
+            MXNet NDArray: ndarray of 0's with given shape
+        """
+        return nd.zeros(shape)
+
+    def power(self, tensor_in_1, tensor_in_2):
+        """
+        Result of first array elements raised to powers from second array,
+        element-wise with broadcasting.
+
+        Args:
+            tensor_in_1: tensor object
+            tensor_in_2: tensor object
+
+        Returns:
+            MXNet NDArray: first array elements raised to powers from second array
+        """
+        tensor_in_1 = self.astensor(tensor_in_1)
+        tensor_in_2 = self.astensor(tensor_in_2)
+        return nd.power(tensor_in_1, tensor_in_2)
+
+    def sqrt(self, tensor_in):
+        """
+        Element-wise square-root value of the input.
+
+        Args:
+            tensor_in: tensor object
+
+        Returns:
+            MXNet NDArray: element-wise square-root value
+        """
+        tensor_in = self.astensor(tensor_in)
+        return nd.sqrt(tensor_in)
+
+    def divide(self, tensor_in_1, tensor_in_2):
+        """
+        Element-wise division of the input arrays with broadcasting.
+
+        Args:
+            tensor_in_1: tensor object
+            tensor_in_2: tensor object
+
+        Returns:
+            MXNet NDArray: element-wise division of the input arrays
+        """
+        tensor_in_1 = self.astensor(tensor_in_1)
+        tensor_in_2 = self.astensor(tensor_in_2)
+        return nd.divide(tensor_in_1, tensor_in_2)
+
+    def log(self, tensor_in):
+        """
+        Element-wise Natural logarithmic value of the input.
+
+        Args:
+            tensor_in: tensor object
+
+        Returns:
+            MXNet NDArray: element-wise Natural logarithmic value
+        """
+        tensor_in = self.astensor(tensor_in)
+        return nd.log(tensor_in)
+
+    def exp(self, tensor_in):
+        """
+        Element-wise exponential value of the input.
+
+        Args:
+            tensor_in: tensor object
+
+        Returns:
+            MXNet NDArray: element-wise exponential value
+        """
+        tensor_in = self.astensor(tensor_in)
+        return nd.exp(tensor_in)
+
+    def stack(self, sequence, axis=0):
+        pass
+
+    def where(self, mask, tensor_in_1, tensor_in_2):
+        pass
+
+    def concatenate(self, sequence):
+        pass
+
+    def simple_broadcast(self, *args):
+        pass
+
+    def poisson(self, n, lam):
+        pass
+
+    def normal(self, x, mu, sigma):
+        pass

setup.py

@@ -19,6 +19,9 @@
     'torch': [
       'torch'
     ],
+    'mxnet':[
+      'mxnet',
+    ],
     'develop': [
        'pyflakes',
        'pytest>=3.2.0',
@@ -29,7 +32,9 @@
        'uproot',
        'papermill',
        'torch',
-       'tensorflow'
+       'tensorflow',
+       'mxnet>=1.0.0',
+       'graphviz'
     ]
   },
   entry_points = {

tests/test_tensor.py

@@ -1,33 +1,39 @@
 from pyhf.tensor.pytorch_backend import pytorch_backend
 from pyhf.tensor.numpy_backend import numpy_backend
 from pyhf.tensor.tensorflow_backend import tensorflow_backend
+from pyhf.tensor.mxnet_backend import mxnet_backend
 from pyhf.simplemodels import hepdata_like
 import tensorflow as tf
 
 
 def test_common_tensor_backends():
     tf_sess = tf.Session()
-    for tb in [numpy_backend(), pytorch_backend(), tensorflow_backend(session = tf_sess)]:
-        assert tb.tolist(tb.astensor([1,2,3])) == [1,2,3]
-        assert tb.tolist(tb.ones((2,3))) == [[1,1,1],[1,1,1]]
-        assert tb.tolist(tb.sum([[1,2,3],[4,5,6]], axis = 0)) == [5,7,9]
-        assert tb.tolist(tb.product([[1,2,3],[4,5,6]], axis = 0)) == [4,10,18]
-        assert tb.tolist(tb.power([1,2,3],[1,2,3])) == [1,4,27]
-        assert tb.tolist(tb.divide([4,9,16],[2,3,4])) == [2,3,4]
-        assert tb.tolist(tb.outer([1,2,3],[4,5,6])) == [[4,5,6],[8,10,12],[12,15,18]]
-        assert tb.tolist(tb.sqrt([4,9,16])) == [2,3,4]
-        assert tb.tolist(tb.stack([tb.astensor([1,2,3]),tb.astensor([4,5,6])])) == [[1,2,3],[4,5,6]]
-        assert tb.tolist(tb.concatenate([tb.astensor([1,2,3]),tb.astensor([4,5,6])])) == [1,2,3,4,5,6]
-        assert tb.tolist(tb.log(tb.exp([2,3,4]))) == [2,3,4]
+    for tb in [numpy_backend(), pytorch_backend(),
+               tensorflow_backend(session=tf_sess), mxnet_backend()]:
+        assert tb.tolist(tb.astensor([1, 2, 3])) == [1, 2, 3]
+        assert tb.tolist(tb.ones((2, 3))) == [[1, 1, 1], [1, 1, 1]]
+        assert tb.tolist(tb.sum([[1, 2, 3], [4, 5, 6]], axis=0)) == [5, 7, 9]
+        assert tb.tolist(
+            tb.product([[1, 2, 3], [4, 5, 6]], axis=0)) == [4, 10, 18]
+        assert tb.tolist(tb.power([1, 2, 3], [1, 2, 3])) == [1, 4, 27]
+        assert tb.tolist(tb.divide([4, 9, 16], [2, 3, 4])) == [2, 3, 4]
+        assert tb.tolist(
+            tb.outer([1, 2, 3], [4, 5, 6])) == [[4, 5, 6], [8, 10, 12], [12, 15, 18]]
+        assert tb.tolist(tb.sqrt([4, 9, 16])) == [2, 3, 4]
+        assert tb.tolist(tb.stack(
+            [tb.astensor([1, 2, 3]), tb.astensor([4, 5, 6])])) == [[1, 2, 3], [4, 5, 6]]
+        assert tb.tolist(tb.concatenate(
+            [tb.astensor([1, 2, 3]), tb.astensor([4, 5, 6])])) == [1, 2, 3, 4, 5, 6]
+        assert tb.tolist(tb.log(tb.exp([2, 3, 4]))) == [2, 3, 4]
         assert tb.tolist(tb.where(
-            tb.astensor([1,0,1]),
-            tb.astensor([1,1,1]),
-            tb.astensor([2,2,2]))) == [1,2,1]
+            tb.astensor([1, 0, 1]),
+            tb.astensor([1, 1, 1]),
+            tb.astensor([2, 2, 2]))) == [1, 2, 1]
 
-        assert list(map(tb.tolist,tb.simple_broadcast(
-            tb.astensor([1,1,1]),
+        assert list(map(tb.tolist, tb.simple_broadcast(
+            tb.astensor([1, 1, 1]),
             tb.astensor([2]),
-            tb.astensor([3,3,3])))) == [[1,1,1],[2,2,2],[3,3,3]]
+            tb.astensor([3, 3, 3])))) == [[1, 1, 1], [2, 2, 2], [3, 3, 3]]
 
 
 def test_pdf_eval():
@@ -36,38 +42,40 @@ def test_pdf_eval():
     oldlib = pyhf.tensorlib
 
     tf_sess = tf.Session()
-    backends = [numpy_backend(poisson_from_normal = True), pytorch_backend(), tensorflow_backend(session = tf_sess)]
+    backends = [numpy_backend(poisson_from_normal=True),
+                pytorch_backend(),
+                tensorflow_backend(session=tf_sess),
+                mxnet_backend()]
 
     values = []
     for b in backends:
-
         pyhf.tensorlib = b
 
         source = {
-          "binning": [2,-0.5,1.5],
-          "bindata": {
-            "data":    [120.0, 180.0],
-            "bkg":     [100.0, 150.0],
-            "bkgsys_up":  [102, 190],
-            "bkgsys_dn":  [98, 100],
-            "sig":     [30.0, 95.0]
-          }
+            "binning": [2, -0.5, 1.5],
+            "bindata": {
+                "data":    [120.0, 180.0],
+                "bkg":     [100.0, 150.0],
+                "bkgsys_up":  [102, 190],
+                "bkgsys_dn":  [98, 100],
+                "sig":     [30.0, 95.0]
+            }
         }
         spec = {
             'singlechannel': {
                 'signal': {
                     'data': source['bindata']['sig'],
-                    'mods': [{'name': 'mu','type': 'normfactor','data': None}]
+                    'mods': [{'name': 'mu', 'type': 'normfactor', 'data': None}]
                 },
                 'background': {
                     'data': source['bindata']['bkg'],
-                    'mods': [{'name': 'bkg_norm','type': 'histosys','data': {
+                    'mods': [{'name': 'bkg_norm', 'type': 'histosys', 'data': {
                         'lo_hist': source['bindata']['bkgsys_dn'], 'hi_hist': source['bindata']['bkgsys_up'],
                     }}]
                 }
             }
         }
-        pdf  = pyhf.hfpdf(spec)
+        pdf = pyhf.hfpdf(spec)
         data = source['bindata']['data'] + pdf.config.auxdata
 
         v1 = pdf.logpdf(pdf.config.suggested_init(), data)
@@ -84,24 +92,27 @@ def test_pdf_eval_2():
     oldlib = pyhf.tensorlib
 
     tf_sess = tf.Session()
-    backends = [numpy_backend(poisson_from_normal = True), pytorch_backend(), tensorflow_backend(session = tf_sess)]
+    backends = [numpy_backend(poisson_from_normal=True),
+                pytorch_backend(),
+                tensorflow_backend(session=tf_sess),
+                mxnet_backend()]
 
     values = []
     for b in backends:
-
         pyhf.tensorlib = b
 
         source = {
-          "binning": [2,-0.5,1.5],
-          "bindata": {
-            "data":    [120.0, 180.0],
-            "bkg":     [100.0, 150.0],
-            "bkgerr":     [10.0, 10.0],
-            "sig":     [30.0, 95.0]
-          }
+            "binning": [2, -0.5, 1.5],
+            "bindata": {
+                "data":    [120.0, 180.0],
+                "bkg":     [100.0, 150.0],
+                "bkgerr":     [10.0, 10.0],
+                "sig":     [30.0, 95.0]
+            }
         }
 
-        pdf  = hepdata_like(source['bindata']['sig'], source['bindata']['bkg'], source['bindata']['bkgerr'])
+        pdf = hepdata_like(source['bindata']['sig'], source['bindata'][
+                           'bkg'], source['bindata']['bkgerr'])
         data = source['bindata']['data'] + pdf.config.auxdata
 
         v1 = pdf.logpdf(pdf.config.suggested_init(), data)