Cryptic shape issue with 1x1 matrix

[jessegrabowski]

Hello,

I am trying to hunt down the cause of a shape issue in my code and I am really stumped. Here's a minimal example that reproduces the error:

import numpy as np
import aesara.tensor as at
from aesara.tensor.nlinalg import matrix_dot

def step_func(y, P, Z, H):
    nan_mask = at.isnan(y)
    W = at.set_subtensor(at.eye(y.shape[0])[nan_mask, nan_mask], 0.0)
    Z_masked = W.dot(Z)
    H_masked = W.dot(H)
    
    F = matrix_dot(Z_masked, P, Z_masked.T) + H_masked
    F_inv = at.linalg.solve(F, at.eye(F.shape[0]))
    
    K = matrix_dot(P, Z_masked.T, F_inv)
    z = matrix_dot(K, H_masked, K.T)
    
    return z

data = np.arange(10).astype(float)[:, None]

k_obs = 1
k_states = 2
k_stochastic = 2

Z = at.zeros((k_obs, k_states))
H = at.zeros((k_obs, k_obs))
P = at.zeros((k_stochastic, k_stochastic))

Z = at.set_subtensor(Z[0, 0], 1.0)
H = at.set_subtensor(H[0, 0], 0.3)
P = at.set_subtensor(P[[0, 1], [0, 1]], 0.7)

z = step_func(data[0], P, Z, H)

Long error traceback

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Input In [201], in <cell line: 1>()
----> 1 z.eval()

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\graph\basic.py:566, in Variable.eval(self, inputs_to_values)
    564 inputs = tuple(sorted(inputs_to_values.keys(), key=id))
    565 if inputs not in self._fn_cache:
--> 566     self._fn_cache[inputs] = function(inputs, self)
    567 args = [inputs_to_values[param] for param in inputs]
    569 rval = self._fn_cache[inputs](*args)

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\compile\function\__init__.py:337, in function(inputs, outputs, mode, updates, givens, no_default_updates, accept_inplace, name, rebuild_strict, allow_input_downcast, profile, on_unused_input)
    331     fn = orig_function(
    332         inputs, outputs, mode=mode, accept_inplace=accept_inplace, name=name
    333     )
    334 else:
    335     # note: pfunc will also call orig_function -- orig_function is
    336     #      a choke point that all compilation must pass through
--> 337     fn = pfunc(
    338         params=inputs,
    339         outputs=outputs,
    340         mode=mode,
    341         updates=updates,
    342         givens=givens,
    343         no_default_updates=no_default_updates,
    344         accept_inplace=accept_inplace,
    345         name=name,
    346         rebuild_strict=rebuild_strict,
    347         allow_input_downcast=allow_input_downcast,
    348         on_unused_input=on_unused_input,
    349         profile=profile,
    350         output_keys=output_keys,
    351     )
    352 return fn

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\compile\function\pfunc.py:363, in pfunc(params, outputs, mode, updates, givens, no_default_updates, accept_inplace, name, rebuild_strict, allow_input_downcast, profile, on_unused_input, output_keys)
    350     profile = ProfileStats(message=profile)
    352 inputs, cloned_outputs = construct_pfunc_ins_and_outs(
    353     params,
    354     outputs,
   (...)
    360     allow_input_downcast,
    361 )
--> 363 return orig_function(
    364     inputs,
    365     cloned_outputs,
    366     mode,
    367     accept_inplace=accept_inplace,
    368     name=name,
    369     profile=profile,
    370     on_unused_input=on_unused_input,
    371     output_keys=output_keys,
    372 )

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\compile\function\types.py:1732, in orig_function(inputs, outputs, mode, accept_inplace, name, profile, on_unused_input, output_keys)
   1730 try:
   1731     Maker = getattr(mode, "function_maker", FunctionMaker)
-> 1732     m = Maker(
   1733         inputs,
   1734         outputs,
   1735         mode,
   1736         accept_inplace=accept_inplace,
   1737         profile=profile,
   1738         on_unused_input=on_unused_input,
   1739         output_keys=output_keys,
   1740         name=name,
   1741     )
   1742     with config.change_flags(compute_test_value="off"):
   1743         fn = m.create(defaults)

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\compile\function\types.py:1471, in FunctionMaker.__init__(self, inputs, outputs, mode, accept_inplace, function_builder, profile, on_unused_input, fgraph, output_keys, name)
   1465 opt_time = None
   1467 with config.change_flags(
   1468     compute_test_value=config.compute_test_value_opt,
   1469     traceback__limit=config.traceback__compile_limit,
   1470 ):
-> 1471     optimizer_profile = optimizer(fgraph)
   1473     end_optimizer = time.time()
   1474     opt_time = end_optimizer - start_optimizer

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\graph\opt.py:112, in GlobalOptimizer.__call__(self, fgraph)
    106 def __call__(self, fgraph):
    107     """Optimize a `FunctionGraph`.
    108 
    109     This is the same as ``self.optimize(fgraph)``.
    110 
    111     """
--> 112     return self.optimize(fgraph)

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\graph\opt.py:103, in GlobalOptimizer.optimize(self, fgraph, *args, **kwargs)
     94 """
     95 
     96 This is meant as a shortcut for the following::
   (...)
    100 
    101 """
    102 self.add_requirements(fgraph)
--> 103 ret = self.apply(fgraph, *args, **kwargs)
    104 return ret

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\graph\opt.py:280, in SeqOptimizer.apply(self, fgraph)
    278 nb_nodes_before = len(fgraph.apply_nodes)
    279 t0 = time.time()
--> 280 sub_prof = optimizer.optimize(fgraph)
    281 l.append(float(time.time() - t0))
    282 sub_profs.append(sub_prof)

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\graph\opt.py:103, in GlobalOptimizer.optimize(self, fgraph, *args, **kwargs)
     94 """
     95 
     96 This is meant as a shortcut for the following::
   (...)
    100 
    101 """
    102 self.add_requirements(fgraph)
--> 103 ret = self.apply(fgraph, *args, **kwargs)
    104 return ret

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\graph\opt.py:280, in SeqOptimizer.apply(self, fgraph)
    278 nb_nodes_before = len(fgraph.apply_nodes)
    279 t0 = time.time()
--> 280 sub_prof = optimizer.optimize(fgraph)
    281 l.append(float(time.time() - t0))
    282 sub_profs.append(sub_prof)

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\graph\opt.py:103, in GlobalOptimizer.optimize(self, fgraph, *args, **kwargs)
     94 """
     95 
     96 This is meant as a shortcut for the following::
   (...)
    100 
    101 """
    102 self.add_requirements(fgraph)
--> 103 ret = self.apply(fgraph, *args, **kwargs)
    104 return ret

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\graph\opt.py:2329, in EquilibriumOptimizer.apply(self, fgraph, start_from)
   2327 nb = change_tracker.nb_imported
   2328 t_opt = time.time()
-> 2329 lopt_change = self.process_node(fgraph, node, lopt)
   2330 time_opts[lopt] += time.time() - t_opt
   2331 if not lopt_change:

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\graph\opt.py:1850, in NavigatorOptimizer.process_node(self, fgraph, node, lopt)
   1848 lopt = lopt or self.local_opt
   1849 try:
-> 1850     replacements = lopt.transform(fgraph, node)
   1851 except Exception as e:
   1852     if self.failure_callback is not None:

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\graph\opt.py:1055, in FromFunctionLocalOptimizer.transform(self, fgraph, node)
   1050     if not (
   1051         node.op in self._tracks or isinstance(node.op, self._tracked_types)
   1052     ):
   1053         return False
-> 1055 return self.fn(fgraph, node)

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\tensor\blas.py:1796, in local_dot22_to_ger_or_gemv(fgraph, node)
   1793 elif not xb[0] and not xb[1] and yb[1]:
   1794     # x is matrix, y is vector, try gemv
   1795     yv = y.dimshuffle(0)
-> 1796     zeros = at.AllocEmpty(x.dtype)(x.shape[0])
   1797     rval = gemv_no_inplace(zeros, one, x, yv, zero)
   1798     new_out = [rval.dimshuffle(0, "x")]

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\graph\op.py:294, in Op.__call__(self, *inputs, **kwargs)
    252 r"""Construct an `Apply` node using :meth:`Op.make_node` and return its outputs.
    253 
    254 This method is just a wrapper around :meth:`Op.make_node`.
   (...)
    291 
    292 """
    293 return_list = kwargs.pop("return_list", False)
--> 294 node = self.make_node(*inputs, **kwargs)
    296 if config.compute_test_value != "off":
    297     compute_test_value(node)

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\tensor\basic.py:4176, in AllocEmpty.make_node(self, *_shape)
   4175 def make_node(self, *_shape):
-> 4176     _shape, bcast = infer_broadcastable(_shape)
   4177     otype = TensorType(dtype=self.dtype, shape=bcast)
   4178     output = otype()

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\tensor\basic.py:1447, in infer_broadcastable(shape)
   1443     raise TypeError(f"Shapes must be scalar integers; got {s_as_str}")
   1445 sh = [check_type(as_tensor_variable(s, ndim=0)) for s in shape]
-> 1447 shape_fg = FunctionGraph(
   1448     outputs=sh,
   1449     features=[ShapeFeature()],
   1450     clone=True,
   1451 )
   1452 folded_shape = optimize_graph(shape_fg, custom_opt=topo_constant_folding).outputs
   1454 bcast = tuple(getattr(s, "data", s) == 1 for s in folded_shape)

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\graph\fg.py:155, in FunctionGraph.__init__(self, inputs, outputs, features, clone, update_mapping, memo, copy_inputs, copy_orphans)
    152     self.add_input(in_var, check=False)
    154 for output in outputs:
--> 155     self.import_var(output, reason="init")
    156 for i, output in enumerate(outputs):
    157     self.clients[output].append(("output", i))

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\graph\fg.py:296, in FunctionGraph.import_var(self, var, reason, import_missing)
    294 # Imports the owners of the variables
    295 if var.owner and var.owner not in self.apply_nodes:
--> 296     self.import_node(var.owner, reason=reason, import_missing=import_missing)
    297 elif (
    298     var.owner is None
    299     and not isinstance(var, Constant)
    300     and var not in self.inputs
    301 ):
    302     from aesara.graph.null_type import NullType

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\graph\fg.py:377, in FunctionGraph.import_node(self, apply_node, check, reason, import_missing)
    375         self.variables.add(input)
    376     self.add_client(input, (node, i))
--> 377 self.execute_callbacks("on_import", node, reason)

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\graph\fg.py:579, in FunctionGraph.execute_callbacks(self, name, *args, **kwargs)
    577         continue
    578     tf0 = time.time()
--> 579     fn(self, *args, **kwargs)
    580     self.execute_callbacks_times[feature] += time.time() - tf0
    581 self.execute_callbacks_time += time.time() - t0

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\tensor\basic_opt.py:1313, in ShapeFeature.on_import(self, fgraph, node, reason)
   1310         o_shapes[sh_idx] = tuple(new_shape)
   1312 for r, s in zip(node.outputs, o_shapes):
-> 1313     self.set_shape(r, s)

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\tensor\basic_opt.py:1086, in ShapeFeature.set_shape(self, r, s, override)
   1084         shape_vars.append(constant(r.type.shape[i], dtype="int64"))
   1085     else:
-> 1086         shape_vars.append(self.unpack(s[i], r))
   1087 assert all(
   1088     not hasattr(r.type, "broadcastable") or not r.type.broadcastable[i] or
   1089     # The two following comparison are a speed optimization
   (...)
   1093     for i in range(r.type.ndim)
   1094 )
   1095 self.shape_of[r] = tuple(shape_vars)

File ~\miniconda3\envs\pymc_dev\lib\site-packages\aesara\tensor\basic_opt.py:986, in ShapeFeature.unpack(self, s_i, var)
    977 """Return a symbolic integer scalar for the shape element s_i.
    978 
    979 The s_i argument was produced by the infer_shape() of an Op subclass.
   (...)
    983 
    984 """
    985 # unpack the s_i that the Op returned
--> 986 assert s_i is not None
    987 if s_i == 1:
    988     # don't make the optimizer merge a zillion ones together
    989     # by always returning the same object to represent 1
    990     return self.lscalar_one

The error seems to be a confluence of several things, because I found several ways to make it go away. None of them, however, give me a general solution to my problem, so here I am.

The first is that if I just set up everything with symbolic objects rather than using at.zeros, it works fine. For example:

data = at.matrix()
P = at.matrix()
Z = at.matrix()
H = at.matrix()

z = step_func(data[0], P, Z, H)
z.eval({data:np.arange(10).astype(float)[:, None],
        P: np.eye(2),
        Z: np.array([[1.0, 0.0]]),
        H: np.array([[1.0]])})

Does not throw an error. As an aside, I have run into a lot of shape issues trying to use at.zeros and at.zeros_like in general. Should these be avoided? I am using them here because of that's how I've written my API. These tend to be sparse matrices, and I wanted users to be able to assign only the relevant elements when they set up their problem.

Anyway, given the zeros/set_subtensor setup, the proximate cause of the error is the final matrix_dot inside the function. Removing this, I can get back all the intermediate computations and confirm that shapes conform. In particular, confirm that H_masked isn't being cast to a scalar. This is important because the ultimate cause seems to have something to do with the matrix H being 1 x 1. Changing the input variables so that H is 2 x 2 fixes the problem, for example:

data = np.arange(10).astype(float)[:, None]
data = data.repeat(2, axis=1)

k_obs = 2
k_states = 2
k_stochastic = 2

Z = at.zeros((k_obs, k_states))
H = at.zeros((k_obs, k_obs))
P = at.zeros((k_stochastic, k_stochastic))

Z = at.set_subtensor(Z[[0, 1], [0, 1]], 1.0)
H = at.set_subtensor(H[[0, 1], [0, 1]], 0.3)
P = at.set_subtensor(P[[0, 1], [0, 1]], 0.7)
z = step_func(data[0], P, Z, H)
z.eval()

Does not throw an error. H being 1x1 is an important special case to my problem (one observed time series), so I really want to figure this out.

In addition, removing the first four lines that zero out columns of H and Z associated with missing data also fixes the problem. That is, this step_func works fine:

def step_func(y, P, Z, H):
    F = matrix_dot(Z, P, Z.T) + H
    F_inv = at.linalg.solve(F, at.eye(F.shape[0]))
    
    K = matrix_dot(P, Z.T, F_inv)
    z = matrix_dot(K, H, K.T)
    
    return z

But then I lose the ability to interpolate missing data.

My favorite fix, however, is changing F = matrix_dot(Z_masked, P, Z_masked.T) + H_masked to F = matrix_dot(Z_masked, P, Z_masked.T). I have no idea why that particular addition (pun intended) would cause a shape error in a matrix multiplication down-stream.

I hope I'm missing something obvious as usual. Also as usual, your help and time are greatly appreciated.

[brandonwillard]

This looks like a genuine bug.

[jessegrabowski]

Any thoughts on how I could help get to the bottom of it?

Running the code snippet I posted above in debug mode, it looks like it fails during the application of an EquilibriumOptimizer (I assume this is associated with the SVD that happens in linalg.solve?), but I think the problem is actually happening before that. H_masked is definitely getting converted to a scalar during optimization, because the assertion error comes from asking for the shape of the inputs to dot22, and it's getting back nothing. What's odd is that this scalar conversion only happens if I do the dot with the W matrix first. Of note is that W is also a 1x1 matrix. Would dot22 (or some other optimization) convert outputs to scalars if all inputs are 1x1?

[brandonwillard]

It looks like a problem with Elemwise.infer_shape; it shouldn't be returning a tuple with Nones. I think I found an MRE last night using only at.add and at.eye(1) inputs, but it still needs an extra step to generate the exact same error.

The idea behind the MRE is that Eye doesn't compute static shapes/broadcastable information in the TensorType outputs it produces in Eye.make_node, so Elemwise.make_node will get inputs that cause it to produce an output with None for its TensorType.shape. Since the Elemwise's inputs are both broadcastable, its output should be as well; however, the logic in Elemwise.infer_shape unreasonably relies on the broadcastable pattern it computed in Elemwise.make_node and produces garbage, instead of simply computing the actual output shape based on the constant inputs.