pymc_marketing.prior.handle_dims does not raise an Error when dims has undesired dimensions

[lucianopaz]

I came across this issue when I made a mistake writing a transformer's dimensions. I had set the parameter dimensions to one sequence, but I had a typo in one of the transformer's dimensions. This raised an obscure error when trying to run inference or sample from the model's prior predictive. Here's the minimal reproducible example that produces this error:

import numpy as np
import pytensor
from pymc_marketing.prior import handle_dims
from pytensor import tensor as pt


shape1 = pytensor.shared(np.array(4, dtype="int"), name="dim1")
shape2 = pytensor.shared(np.array(6, dtype="int"), name="dim2")
x_dynamic = pt.random.normal(size=(shape1, shape2))
x_dynamic.name = "x"

hx = handle_dims(x_dynamic, dims=("a", "b"), desired_dims=("a", "B"))
hx.dprint(print_type=True)

Note that the desired_dims has a dimension called "B" instead of "b". handle_dims completely ignores the fact that a dimension called "b" exists, and tries to soldier through.

The debug print gives this

DimShuffle{order=[0,x]} [id A] <Matrix(float64, shape=(?, 1))>
 └─ normal_rv{"(),()->()"}.1 [id B] <Matrix(float64, shape=(?, ?))> 'x'
    ├─ RNG(<Generator(PCG64) at 0x1657A51C0>) [id C] <RandomGeneratorType>
    ├─ MakeVector{dtype='int64'} [id D] <Vector(int64, shape=(2,))>
    │  ├─ dim1 [id E] <Scalar(int64, shape=())>
    │  └─ dim2 [id F] <Scalar(int64, shape=())>
    ├─ ExpandDims{axes=[0, 1]} [id G] <Matrix(float32, shape=(1, 1))>
    │  └─ 0.0 [id H] <Scalar(float32, shape=())>
    └─ ExpandDims{axes=[0, 1]} [id I] <Matrix(float32, shape=(1, 1))>
       └─ 1.0 [id J] <Scalar(float32, shape=())>

where the DimShuffle Op is trying to add a new dimension of shape 1 at the end of x while dropping the existing dimension. If we actually try to eval the resulting tensor, we get an error:

>>> hx.eval().shape
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
File pytensor/compile/function/types.py:960, in Function.__call__(self, *args, **kwargs)
    958 try:
    959     outputs = (
--> 960         self.vm()
    961         if output_subset is None
    962         else self.vm(output_subset=output_subset)
    963     )
    964 except Exception:

ValueError: cannot reshape array of size 24 into shape (4,1)

During handling of the above exception, another exception occurred:

ValueError                                Traceback (most recent call last)
Cell In[13], line 1
----> 1 hx.eval().shape

File pytensor/graph/basic.py:658, in Variable.eval(self, inputs_to_values, **kwargs)
    652         warnings.warn(
    653             "Keyword arguments could not be used to create a cache key for the underlying variable. "
    654             f"A function will be recompiled on every call with such keyword arguments.\n{exc}"
    655         )
    657 args = [parsed_inputs_to_values[param] for param in inputs]
--> 658 return fn(*args)

File pytensor/compile/function/types.py:973, in Function.__call__(self, *args, **kwargs)
    971     if hasattr(self.vm, "thunks"):
    972         thunk = self.vm.thunks[self.vm.position_of_error]
--> 973     raise_with_op(
    974         self.maker.fgraph,
    975         node=self.vm.nodes[self.vm.position_of_error],
    976         thunk=thunk,
    977         storage_map=getattr(self.vm, "storage_map", None),
    978     )
    979 else:
    980     # old-style linkers raise their own exceptions
    981     raise

File pytensor/link/utils.py:524, in raise_with_op(fgraph, node, thunk, exc_info, storage_map)
    519     warnings.warn(
    520         f"{exc_type} error does not allow us to add an extra error message"
    521     )
    522     # Some exception need extra parameter in inputs. So forget the
    523     # extra long error message in that case.
--> 524 raise exc_value.with_traceback(exc_trace)

File pytensor/compile/function/types.py:960, in Function.__call__(self, *args, **kwargs)
    957     t0_fn = time.perf_counter()
    958 try:
    959     outputs = (
--> 960         self.vm()
    961         if output_subset is None
    962         else self.vm(output_subset=output_subset)
    963     )
    964 except Exception:
    965     self._restore_defaults()

ValueError: cannot reshape array of size 24 into shape (4,1)
Apply node that caused the error: DimShuffle{order=[0,x]}(x)
Toposort index: 2
Inputs types: [TensorType(float64, shape=(None, None))]
Inputs shapes: [(4, 6)]
Inputs strides: [(48, 8)]
Inputs values: ['not shown']
Outputs clients: [[output[0](DimShuffle{order=[0,x]}.0)]]

Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
  File "/var/folders/p0/81r8yvt526ldqjrrf14b4bd00000gn/T/ipykernel_37472/2106782338.py", line 6, in <module>
    hx = handle_dims(x_dynamic, dims=("a", "b"), desired_dims=("a", "B"))
  File "pymc_marketing/prior.py", line 168, in handle_dims
    return x.dimshuffle(*args)

HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.

Selection deleted
x_static = pt.matrix("x", shape=(4, 6))

handle_dims(x_static, dims=("a", "b"), desired_dims=("a", "B")).dprint(print_type=True)

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[3], line 3
      1 x_static = pt.matrix("x", shape=(4, 6))
----> 3 handle_dims(x_static, dims=("a", "b"), desired_dims=("a", "B")).dprint(print_type=True)

File pymc_marketing/prior.py:168, in handle_dims(x, dims, desired_dims)
    163 args = [
    164     "x" if missing else idx
    165     for (idx, missing) in zip(new_idx, missing_dims, strict=False)
    166 ]
    167 args = _remove_leading_xs(args)
--> 168 return x.dimshuffle(*args)

File pytensor/tensor/variable.py:348, in _tensor_py_operators.dimshuffle(self, *pattern)
    346     pattern = pattern[0]
    347 ds_op = pt.elemwise.DimShuffle(input_ndim=self.type.ndim, new_order=pattern)
--> 348 return ds_op(self)

File pytensor/graph/op.py:293, in Op.__call__(self, name, return_list, *inputs, **kwargs)
    249 def __call__(
    250     self, *inputs: Any, name=None, return_list=False, **kwargs
    251 ) -> Variable | list[Variable]:
    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     node = self.make_node(*inputs, **kwargs)
    294     if name is not None:
    295         if len(node.outputs) == 1:

File pytensor/tensor/elemwise.py:202, in DimShuffle.make_node(self, inp)
    200 for d in self.drop:
    201     if input_static_shape[d] not in (1, None):
--> 202         raise TypeError(
    203             f"Input dropped dimension {d} must have length 1 but has {input_static_shape[d]}"
    204         )
    206 out_static_shape = []
    207 for dim_idx in self.new_order:

TypeError: Input dropped dimension 1 must have length 1 but has 6

If the underlying tensor had static shape information, the error would have been picked up immediately:

>>> x_static = pt.matrix("x", shape=(4, 6))
>>> handle_dims(x_static, dims=("a", "b"), desired_dims=("a", "B")).dprint(print_type=True)

TypeError                                 Traceback (most recent call last)
      1 x_static = pt.matrix("x", shape=(4, 6))
----> 2 handle_dims(x_static, dims=("a", "b"), desired_dims=("a", "B")).dprint(print_type=True)

File ~pymc_marketing/prior.py:168, in handle_dims(x, dims, desired_dims)
    163 args = [
    164     "x" if missing else idx
    165     for (idx, missing) in zip(new_idx, missing_dims, strict=False)
    166 ]
    167 args = _remove_leading_xs(args)
--> 168 return x.dimshuffle(*args)

File pytensor/tensor/variable.py:348, in _tensor_py_operators.dimshuffle(self, *pattern)
    346     pattern = pattern[0]
    347 ds_op = pt.elemwise.DimShuffle(input_ndim=self.type.ndim, new_order=pattern)
--> 348 return ds_op(self)

File pytensor/graph/op.py:293, in Op.__call__(self, name, return_list, *inputs, **kwargs)
    249 def __call__(
    250     self, *inputs: Any, name=None, return_list=False, **kwargs
    251 ) -> Variable | list[Variable]:
    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     node = self.make_node(*inputs, **kwargs)
    294     if name is not None:
    295         if len(node.outputs) == 1:

File pytensor/tensor/elemwise.py:202, in DimShuffle.make_node(self, inp)
    200 for d in self.drop:
    201     if input_static_shape[d] not in (1, None):
--> 202         raise TypeError(
    203             f"Input dropped dimension {d} must have length 1 but has {input_static_shape[d]}"
    204         )
    206 out_static_shape = []
    207 for dim_idx in self.new_order:

TypeError: Input dropped dimension 1 must have length 1 but has 6

Summary

handle_dims ignored the existence of dimensions in dims that were not included in desired_dims, leading to hard to identify errors creeping up much later in the workflow. We should add a check to verify that the desired_dims set includes the supplied set of dims.

[wd60622]

Would you recommend raising early to catch the b != B in your two dims example?

[wd60622]

I will make a PR

[lucianopaz]

I recommend adding an assert like this:

assert set(dims) <= set(desired_dims)

If there are dimension names in dims that are not in desired_dims, that assertion should pick it up. The error message could be nicer by listing the dimensions that were in dims but not in desired_dims. Something like set(dims) - set(desired_dims).

[lucianopaz]

So I guess that

if set(dims) <= set(desired_dims):
    raise ValueError(
        "An error occurred while handling the dimensions. "
        "The requested desired dimensions do not include the following "
        "dimensions that are part of the input tensor: "
        f"{set(dims) - set(desired_dims)}".
    )

would be a good enough error message to add into handle_dims.