model.py: changes to the length 1 fix

pymc3/model.py

@@ -401,6 +401,8 @@ class ValueGradFunction:
     """
     def __init__(self, cost, grad_vars, extra_vars=None, dtype=None,
                  casting='no', **kwargs):
+        from .distributions import TensorType
+
         if extra_vars is None:
             extra_vars = []
 
@@ -437,12 +439,12 @@ def __init__(self, cost, grad_vars, extra_vars=None, dtype=None,
         self._extra_vars_shared = {}
         for var in extra_vars:
             shared = theano.shared(var.tag.test_value, var.name + '_shared__')
-            # manually set the TensorType for length 1 arrays due to a
-            # theano conversion problem
-            if isinstance(var.tag.test_value, np.ndarray):
-                if var.tag.test_value.ndim == 1:
-                    if len(var.tag.test_value) == 1:
-                        shared.type = theano.tensor.TensorType(var.dtype, (True,))
+            # test TensorType compatibility
+            if hasattr(var.tag.test_value, 'shape'):
+                testtype = TensorType(var.dtype, var.tag.test_value.shape)
+
+                if testtype != shared.type:
+                    shared.type = testtype
             self._extra_vars_shared[var.name] = shared
             givens.append((var, shared))