Update infer_shape signatures

Author: brandonwillard

pymc3/distributions/multivariate.py

@@ -740,7 +740,7 @@ def perform(self, node, inputs, outputs):
             pm._log.exception("Failed to check if %s positive definite", x)
             raise
 
-    def infer_shape(self, node, shapes):
+    def infer_shape(self, fgraph, node, shapes):
         return [[]]
 
     def grad(self, inp, grads):

pymc3/math.py

@@ -359,7 +359,7 @@ def grad(self, inputs, gout):
         idx = tt.arange(gz.shape[-1])
         return [gz[..., idx, idx]]
 
-    def infer_shape(self, nodes, shapes):
+    def infer_shape(self, fgraph, nodes, shapes):
         return [(shapes[0][0],) + (shapes[0][1],) * 2]
 
 
@@ -418,7 +418,7 @@ def grad(self, inputs, gout):
         ]
         return [gout[0][slc] for slc in slices]
 
-    def infer_shape(self, nodes, shapes):
+    def infer_shape(self, fgraph, nodes, shapes):
         first, second = zip(*shapes)
         return [(tt.add(*first), tt.add(*second))]
 

pymc3/ode/ode.py

@@ -210,7 +210,7 @@ def perform(self, node, inputs_storage, output_storage):
         # simulate states and sensitivities in one forward pass
         output_storage[0][0], output_storage[1][0] = self._simulate(y0, theta)
 
-    def infer_shape(self, node, input_shapes):
+    def infer_shape(self, fgraph, node, input_shapes):
         s_y0, s_theta = input_shapes
         output_shapes = [(self.n_times, self.n_states), (self.n_times, self.n_states, self.n_p)]
         return output_shapes