add total size hints

Author: ferrine

pymc3/data.py

@@ -144,7 +144,7 @@ class Minibatch(tt.TensorVariable):
     >>> with pm.Model() as model:
     ...     mu = pm.Flat('mu')
     ...     sd = pm.HalfNormal('sd')
-    ...     lik = pm.Normal('lik', mu, sd, observed=x)
+    ...     lik = pm.Normal('lik', mu, sd, observed=x, total_size=(100, 100))
 
     Then you can perform regular Variational Inference out of the box
     >>> with model:
@@ -188,18 +188,21 @@ class Minibatch(tt.TensorVariable):
     For more complex slices some more code is needed that can seem not so clear
     They are
     >>> moredata = np.random.rand(10, 20, 30, 40, 50)
+    
+    default total_size is then (10, 20, 30, 40, 50) but 
+    can be less verbose in sove cases
 
-    1) Advanced indexing
+    1) Advanced indexing, `total_size = (10, Ellipsis, 50)`
     >>> x = Minibatch(moredata, [2, Ellipsis, 10])
 
     We take slice only for the first and last dimension
     >>> assert x.eval().shape == (2, 20, 30, 40, 10)
 
-    2) skipping particular dimension
+    2) skipping particular dimension, total_size = (10, None, 30) 
     >>> x = Minibatch(moredata, [2, None, 20])
     >>> assert x.eval().shape == (2, 20, 20, 40, 50)
 
-    3) mixing that all
+    3) mixing that all, total_size = (10, None, 30, Ellipsis, 50)
     >>> x = Minibatch(moredata, [2, None, 20, Ellipsis, 10])
     >>> assert x.eval().shape == (2, 20, 20, 40, 10)
 

pymc3/tests/test_minibatches.py

@@ -223,6 +223,18 @@ def test_common_errors(self):
                 Normal('n', observed=[[1]], total_size=[Ellipsis, Ellipsis])
             assert 'Double Ellipsis' in str(e.value)
 
+    def test_mixed1(self):
+        with pm.Model():
+            data = np.random.rand(10, 20, 30, 40, 50)
+            mb = pm.Minibatch(data, [2, None, 20, Ellipsis, 10])
+            Normal('n', observed=mb, total_size=(10, None, 30, Ellipsis, 50))
+
+    def test_mixed2(self):
+        with pm.Model():
+            data = np.random.rand(10, 20, 30, 40, 50)
+            mb = pm.Minibatch(data, [2, None, 20])
+            Normal('n', observed=mb, total_size=(10, None, 30))
+
     def test_free_rv(self):
         with pm.Model() as model4:
             Normal('n', observed=[[1, 1],