Proposal: Dist shape refactor

pymc3/distributions/distribution.py

@@ -94,7 +94,7 @@ def dist(cls, *args, **kwargs):
         return dist
 
     def __init__(self, shape_supp, shape_ind, shape_reps, bcast, dtype,
-                 testval=None, defaults=None, transform=None):
+                 testval=None, defaults=None, transform=None, *args, **kwargs):
         r"""
         Distributions are specified in terms of the shape of their support, the shape
         of the space of independent instances and the shape of the space of replications.
@@ -175,6 +175,9 @@ def __init__(self, shape_supp, shape_ind, shape_reps, bcast, dtype,
         self.shape_reps = _as_tensor_shape_variable(shape_reps)
         self.ndim_reps = tt.get_vector_length(self.shape_reps)
 
+        self.bcast = bcast
+        self.dtype = dtype
+
         ndim_sum = self.ndim_supp + self.ndim_ind + self.ndim_reps
         if ndim_sum == 0:
             self.shape = tt.constant([], dtype='int64')
@@ -197,7 +200,7 @@ def __init__(self, shape_supp, shape_ind, shape_reps, bcast, dtype,
             testval = self.get_test_value(defaults=self.defaults)
 
         self.testval = testval
-        self.type = tt.TensorType(str(dtype), bcast)
+        self.type = tt.TensorType(str(dtype), self.bcast)
 
     def default(self):
         return self.get_test_value(self.testval, self.defaults)

pymc3/distributions/transforms.py

@@ -46,13 +46,14 @@ class TransformedDistribution(Distribution):
     """A distribution that has been transformed from one space into another."""
 
     def __init__(self, dist, transform, *args, **kwargs):
-        """
+        r"""
         Parameters
         ----------
         dist : Distribution
+            TODO
         transform : Transform
-        args, kwargs
-            arguments to Distribution"""
+            TODO
+        """
         forward = transform.forward
         testval = forward(dist.default())
 
@@ -61,9 +62,18 @@ def __init__(self, dist, transform, *args, **kwargs):
         v = forward(FreeRV(name='v', distribution=dist))
         self.type = v.type
 
+        # We can get the transformed support shape from a single dummy var in
+        # only the support (i.e. without the independent or replication dimensions).
+        shape_supp = forward(tt.alloc(1, *dist.shape_supp)).shape
+
+        # XXX: We assume these two shapes don't change under a transform.
+        shape_ind = dist.shape_ind
+        shape_reps = dist.shape_reps
+
         super(TransformedDistribution, self).__init__(
-            v.shape.tag.test_value, v.dtype,
-            testval, dist.defaults, *args, **kwargs)
+              shape_supp, shape_ind, shape_reps,
+              v.broadcastable, v.dtype,
+              testval.tag.test_value, dist.defaults, *args, **kwargs)
 
         if transform.name == 'stickbreaking':
             b = np.hstack(((np.atleast_1d(self.shape) == 1)[:-1], False))
@@ -193,7 +203,7 @@ class StickBreaking(Transform):
     Parameters
     ----------
     eps : float, positive value
-        A small value for numerical stability in invlogit. 
+        A small value for numerical stability in invlogit.
     """
 
     name = "stickbreaking"
@@ -250,7 +260,7 @@ def backward(self, y):
 
     def forward(self, x):
         return x
-        
+
     def jacobian_det(self, x):
         return 0
 

pymc3/variational/advi.py

@@ -208,11 +208,11 @@ def logp_(input):
         r = MRG_RandomStreams(seed=random_seed)
 
     if n_mcsamples == 1:
-        n = r.normal(size=inarray.tag.test_value.shape)
+        n = r.normal(size=np.shape(inarray.tag.test_value))
         q = n * tt.exp(w) + u
         elbo = logp_(q) + tt.sum(w) + 0.5 * l * (1 + tt.log(2.0 * np.pi))
     else:
-        n = r.normal(size=(n_mcsamples, u.tag.test_value.shape[0]))
+        n = r.normal(size=(n_mcsamples, np.shape(u.tag.test_value)[0]))
         qs = n * tt.exp(w) + u
         logps, _ = theano.scan(fn=lambda q: logp_(q),
                                outputs_info=None,
@@ -255,7 +255,7 @@ def optimizer(loss, param):
             i_int = i.astype('int64')
             value = param_.get_value(borrow=True)
             accu = theano.shared(
-                np.zeros(value.shape + (n_win,), dtype=value.dtype))
+                np.zeros(np.shape(value) + (n_win,), dtype=value.dtype))
             grad = tt.grad(loss, param_)
 
             # Append squared gradient vector to accu_new
@@ -324,17 +324,17 @@ def rvs(x):
     for v in global_RVs:
         u = theano.shared(vparams['means'][str(v)]).ravel()
         w = theano.shared(vparams['stds'][str(v)]).ravel()
-        n = r.normal(size=u.tag.test_value.shape)
-        updates.update({v: (n * w + u).reshape(v.tag.test_value.shape)})
+        n = r.normal(size=np.shape(u.tag.test_value))
+        updates.update({v: (n * w + u).reshape(np.shape(v.tag.test_value))})
 
     if local_RVs is not None:
         for v_, (uw, _) in local_RVs.items():
             v = get_transformed(v_)
             u = uw[0].ravel()
             w = uw[1].ravel()
-            n = r.normal(size=u.tag.test_value.shape)
+            n = r.normal(size=np.shape(u.tag.test_value))
             updates.update(
-                {v: (n * tt.exp(w) + u).reshape(v.tag.test_value.shape)})
+                {v: (n * tt.exp(w) + u).reshape(np.shape(v.tag.test_value))})
 
     # Replace some nodes of the graph with variational distributions
     vars = model.free_RVs