Type's shape not compatible with the data's

[aloctavodia]

This used to work on PyMC-BART 0.2.1 and PyMC 4.1.7

with pm.Model() as model
    w = pmb.BART("w", X, Y, shape=(2, 200))
    y = pm.Normal("y", w[0], pm.math.abs(w[1]), observed=Y)
    idatal = pm.sample()

also this

with pm.Model() as model
    w = pmb.BART("w", X, Y,  size=2)
    y = pm.Normal("y", w[0], pm.math.abs(w[1]), observed=Y)
    idata = pm.sample()

Now it returns the following error

traceback

--------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
File ~/anaconda3/envs/test/lib/python3.10/site-packages/pytensor/link/basic.py:111, in Container.__set__(self, value)
    109 try:
    110     # Use in-place filtering when/if possible
--> 111     self.storage[0] = self.type.filter_inplace(
    112         value, self.storage[0], **kwargs
    113     )
    114 except NotImplementedError:

File ~/anaconda3/envs/test/lib/python3.10/site-packages/pytensor/graph/type.py:130, in Type.filter_inplace(self, value, storage, strict, allow_downcast)
    111 """Return data or an appropriately wrapped/converted data by converting it in-place.
    112 
    113 This method allows one to reuse old allocated memory.  If this method
   (...)
    128 NotImplementedError
    129 """
--> 130 raise NotImplementedError()

NotImplementedError: 

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
Cell In[11], line 4
      2     w = pmb.BART("w", X, Y, m=2, shape=(2, 200))
      3     y = pm.Normal("y", w[0], pm.math.abs(w[1]), observed=Y)
----> 4     idata_marketing_full = pm.sample(compute_convergence_checks=False, return_inferencedata=False)
      5 #w.eval().shape
      6 #pm.model_to_graphviz(model_marketing_full)
      7 model_marketing_full.initial_point()["w"].shape

File ~/proyectos/00_BM/pymc/pymc/sampling/mcmc.py:574, in sample(draws, step, init, n_init, initvals, trace, chains, cores, tune, progressbar, model, random_seed, discard_tuned_samples, compute_convergence_checks, callback, jitter_max_retries, return_inferencedata, keep_warning_stat, idata_kwargs, mp_ctx, **kwargs)
    572 else:
    573     traces, length = _choose_chains(traces, 0)
--> 574 mtrace = MultiTrace(traces)[:length]
    576 # count the number of tune/draw iterations that happened
    577 # ideally via the "tune" statistic, but not all samplers record it!
    578 if "tune" in mtrace.stat_names:

File ~/proyectos/00_BM/pymc/pymc/backends/base.py:310, in MultiTrace.__getitem__(self, idx)
    308 def __getitem__(self, idx):
    309     if isinstance(idx, slice):
--> 310         return self._slice(idx)
    312     try:
    313         return self.point(int(idx))

File ~/proyectos/00_BM/pymc/pymc/backends/base.py:466, in MultiTrace._slice(self, slice)
    464 def _slice(self, slice):
    465     """Return a new MultiTrace object sliced according to `slice`."""
--> 466     new_traces = [trace._slice(slice) for trace in self._straces.values()]
    467     trace = MultiTrace(new_traces)
    468     idxs = slice.indices(len(self))

File ~/proyectos/00_BM/pymc/pymc/backends/base.py:466, in <listcomp>(.0)
    464 def _slice(self, slice):
    465     """Return a new MultiTrace object sliced according to `slice`."""
--> 466     new_traces = [trace._slice(slice) for trace in self._straces.values()]
    467     trace = MultiTrace(new_traces)
    468     idxs = slice.indices(len(self))

File ~/proyectos/00_BM/pymc/pymc/backends/ndarray.py:167, in NDArray._slice(self, idx)
    159 def _slice(self, idx):
    160     # Slicing directly instead of using _slice_as_ndarray to
    161     # support stop value in slice (which is needed by
    162     # iter_sample).
    163 
    164     # Only the first `draw_idx` value are valid because of preallocation
    165     idx = slice(*idx.indices(len(self)))
--> 167     sliced = NDArray(model=self.model, vars=self.vars)
    168     sliced.chain = self.chain
    169     sliced.samples = {varname: values[idx] for varname, values in self.samples.items()}

File ~/proyectos/00_BM/pymc/pymc/backends/ndarray.py:45, in NDArray.__init__(self, name, model, vars, test_point)
     44 def __init__(self, name=None, model=None, vars=None, test_point=None):
---> 45     super().__init__(name, model, vars, test_point)
     46     self.draw_idx = 0
     47     self.draws = None

File ~/proyectos/00_BM/pymc/pymc/backends/base.py:78, in BaseTrace.__init__(self, name, model, vars, test_point)
     76     test_point_.update(test_point)
     77     test_point = test_point_
---> 78 var_values = list(zip(self.varnames, self.fn(test_point)))
     79 self.var_shapes = {var: value.shape for var, value in var_values}
     80 self.var_dtypes = {var: value.dtype for var, value in var_values}

File ~/proyectos/00_BM/pymc/pymc/pytensorf.py:764, in PointFunc.__call__(self, state)
    763 def __call__(self, state):
--> 764     return self.f(**state)

File ~/anaconda3/envs/test/lib/python3.10/site-packages/pytensor/compile/function/types.py:897, in Function.__call__(self, *args, **kwargs)
    895 if kwargs:  # for speed, skip the items for empty kwargs
    896     for k, arg in kwargs.items():
--> 897         self[k] = arg
    899 if (
    900     not self.trust_input
    901     and
   (...)
    904 ):
    905     # Collect aliased inputs among the storage space
    906     args_share_memory = []

File ~/anaconda3/envs/test/lib/python3.10/site-packages/pytensor/compile/function/types.py:549, in Function.__setitem__(self, item, value)
    548 def __setitem__(self, item, value):
--> 549     self.value[item] = value

File ~/anaconda3/envs/test/lib/python3.10/site-packages/pytensor/compile/function/types.py:505, in Function.__init__.<locals>.ValueAttribute.__setitem__(self, item, value)
    499     raise TypeError(
    500         f"Ambiguous name: {item} - please check the "
    501         "names of the inputs of your function "
    502         "for duplicates."
    503     )
    504 if isinstance(s, Container):
--> 505     s.value = value
    506     s.provided += 1
    507 else:

File ~/anaconda3/envs/test/lib/python3.10/site-packages/pytensor/link/basic.py:115, in Container.__set__(self, value)
    111         self.storage[0] = self.type.filter_inplace(
    112             value, self.storage[0], **kwargs
    113         )
    114     except NotImplementedError:
--> 115         self.storage[0] = self.type.filter(value, **kwargs)
    117 except Exception as e:
    118     e.args = e.args + (f'Container name "{self.name}"',)

File ~/anaconda3/envs/test/lib/python3.10/site-packages/pytensor/tensor/type.py:251, in TensorType.filter(self, data, strict, allow_downcast)
    242     raise TypeError(
    243         "The numpy.ndarray object is not aligned."
    244         " PyTensor C code does not support that.",
    245     )
    247 if not all(
    248     ds == ts if ts is not None else True
    249     for ds, ts in zip(data.shape, self.shape)
    250 ):
--> 251     raise TypeError(
    252         f"The type's shape ({self.shape}) is not compatible with the data's ({data.shape})"
    253     )
    255 if self.filter_checks_isfinite and not np.all(np.isfinite(data)):
    256     raise ValueError("Non-finite elements not allowed")

TypeError: ("The type's shape ((2, 200)) is not compatible with the data's ((200, 2))", 'Container name "w"')

Do you have a guess on how to fix this @michaelosthege, @ricardoV94?

[michaelosthege]

What are the shapes and dtypes of X and Y?

[aloctavodia]

X.shape, (200, 1)
Y.shape, (200,)

both are float64

[michaelosthege]

The shape passed to pmb.BART doesn't match the X. If this worked before, this was probably just luck and funny broadcasting

...broadcasting (200, 1) to (2, 200) I mean

[aloctavodia]

sorry for the silly question but when/why do you need to broadcast (200, 1) to (2, 200)?

[michaelosthege]

I don't know how exactly pm.BART works, but you're passing an X=(200, 1), Y=(200,) while telling it to come out as shape=(2, 200).
And the error is also about the w variable and exactly these two shapes being incompatible.

To me this sounds like you were relying on some internal broadcasting mechanisms figuring out how to reshape.

If this was any of the standard distributions I'd try to slice/reshape X myself, or pass a shape=(200, 1) to avoid any behind-the-scenes reshaping/broadcasting.

[aloctavodia]

If the shape is not defined by the user, the shape of a BART variable like w is X.shape[0]. And it used to be possible to define shapes like (X.shape[0], n), or or maybe it was (n, X.shape[0]) (I should check). Other shapes were not allowed.

Thanks for your input.

[aloctavodia]

Ohh more more note. Just to clarify, this error is raised after sampling has finished.

[michaelosthege]

Ohh more more note. Just to clarify, this error is raised after sampling has finished.

That's an important detail. The error originates when during NDArray initialization it recompiles and evaluates a test point to get shape information about the variables.

Side note: We can avoid this unnecessary recompilation by switching to McBackend traces! Please review pymc-devs/pymc#6475 which is a preparation for that. After that it takes just one commit to get optional McBackend support!

[aloctavodia]

I will try to review your PR tomorrow. Thanks again for the feedback.