Much better api

Author: ferrine

pymc3/data.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 import pymc3 as pm
+from pymc3 import theanof
 import theano.tensor as tt
 import theano
 
@@ -149,7 +150,7 @@ def __next__(self):
     next = __next__
 
 
-class Minibatch(object):
+class Minibatch(tt.TensorVariable):
     """Multidimensional minibatch
 
     Parameters
@@ -173,18 +174,31 @@ class Minibatch(object):
     minibatch : minibatch tensor
         Used for training
     """
-    def __init__(self, data, batch_size=128, in_memory_size=None, random_seed=42, update_shared_f=None):
+    @theanof.change_flags(compute_test_value='raise')
+    def __init__(self, data, batch_size=128, in_memory_size=None,
+                 random_seed=42, update_shared_f=None,
+                 broadcastable=None, name='Minibatch'):
         data = pm.smartfloatX(np.asarray(data))
         self._random_seed = random_seed
         in_memory_slc = self._to_slices(in_memory_size)
-        self.data = data
         self.batch_size = batch_size
         self.shared = theano.shared(data[in_memory_slc])
         self.update_shared_f = update_shared_f
         self.random_slc = self._to_random_slices(self.shared.shape, batch_size)
-        self.minibatch = self.shared[self.random_slc]
+        minibatch = self.shared[self.random_slc]
+        if broadcastable is None:
+            broadcastable = (False, ) * minibatch.ndim
+        minibatch = tt.patternbroadcast(minibatch, broadcastable)
+        self.minibatch = minibatch
+        super(Minibatch, self).__init__(
+            self.minibatch.type, None, None, name=name)
+        theano.Apply(
+            theano.compile.view_op,
+            inputs=[self.minibatch], outputs=[self])
+        self.tag.test_value = copy(self.minibatch.tag.test_value)
 
-    def rslice(self, total, size, seed):
+    @staticmethod
+    def rslice(total, size, seed):
         if size is None:
             return slice(None)
         elif isinstance(size, int):
@@ -284,13 +298,8 @@ def update_shared(self):
     def set_value(self, value):
         self.shared.set_value(np.asarray(value, self.dtype))
 
-    @property
-    def dtype(self):
-        return self.shared.dtype
-
-    @property
-    def type(self):
-        return self.shared.type
-
-    def __repr__(self):
-        return '<Minibatch of %s>' % self.batch_size
+    def clone(self):
+        ret = self.type()
+        ret.name = self.name
+        ret.tag = copy(self.tag)
+        return ret

pymc3/tests/test_minibatches.py

@@ -246,22 +246,30 @@ class TestMinibatch(object):
 
     def test_1d(self):
         mb = pm.Minibatch(self.data, 20)
-        assert mb.minibatch.eval().shape == (20, 10, 40, 10, 50)
+        assert mb.eval().shape == (20, 10, 40, 10, 50)
 
     def test_2d(self):
         with pytest.raises(TypeError):
             pm.Minibatch(self.data, (10, 5))
         mb = pm.Minibatch(self.data, [(10, 42), (4, 42)])
-        assert mb.minibatch.eval().shape == (10, 4, 40, 10, 50)
+        assert mb.eval().shape == (10, 4, 40, 10, 50)
 
     def test_special1(self):
         mb = pm.Minibatch(self.data, [(10, 42), None, (4, 42)])
-        assert mb.minibatch.eval().shape == (10, 10, 4, 10, 50)
+        assert mb.eval().shape == (10, 10, 4, 10, 50)
 
     def test_special2(self):
         mb = pm.Minibatch(self.data, [(10, 42), Ellipsis, (4, 42)])
-        assert mb.minibatch.eval().shape == (10, 10, 40, 10, 4)
+        assert mb.eval().shape == (10, 10, 40, 10, 4)
 
     def test_special3(self):
         mb = pm.Minibatch(self.data, [(10, 42), None, Ellipsis, (4, 42)])
-        assert mb.minibatch.eval().shape == (10, 10, 40, 10, 4)
+        assert mb.eval().shape == (10, 10, 40, 10, 4)
+
+    def test_cloning_available(self):
+        gop = pm.Minibatch(np.arange(100), 1)
+        res = gop ** 2
+        shared = theano.shared(np.array([10]))
+        res1 = theano.clone(res, {gop: shared})
+        f = theano.function([], res1)
+        assert f() == np.array([100])