Bayesian VAR with missing data: PyTensor error when sampling with blackjax or numpyro

Questions v5
,
1

Hi,

I am adapting the official PyMC Bayesian VAR example to work with missing data, see the attached minimal self-contained example at the end of this message. If I use the default nuts_sampler="pymc" when sampling, the code runs smoothly. Instead with nuts_sampler="blackjax or nuts_sampler="numpyro", I get multiple error messages related to PyTensor, like:

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: boolean_indexing_set_or_inc
ERROR (pytensor.graph.rewriting.basic): node: AdvancedIncSubtensor{inplace=False,  set_instead_of_inc=True}(TensorConstant{[[-0.81296..4479119 ]]}, X_unobserved, TensorConstant{[[False Fa..se False]]})
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/ubuntu/mambaforge/envs/pymc5_mamba/lib/python3.11/site-packages/pytensor/tensor/elemwise.py", line 441, in get_output_info
    [
  File "/home/ubuntu/mambaforge/envs/pymc5_mamba/lib/python3.11/site-packages/pytensor/tensor/elemwise.py", line 442, in <listcomp>
    get_most_specialized_shape(shape)
  File "/home/ubuntu/mambaforge/envs/pymc5_mamba/lib/python3.11/site-packages/pytensor/tensor/elemwise.py", line 434, in get_most_specialized_shape
    raise ValueError
ValueError

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/ubuntu/mambaforge/envs/pymc5_mamba/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1918, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/mambaforge/envs/pymc5_mamba/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1078, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/mambaforge/envs/pymc5_mamba/lib/python3.11/site-packages/pytensor/tensor/rewriting/jax.py", line 34, in boolean_indexing_set_or_inc
    out = at.where(cond, y, x)
          ^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/mambaforge/envs/pymc5_mamba/lib/python3.11/site-packages/pytensor/graph/op.py", line 295, in __call__
    node = self.make_node(*inputs, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/mambaforge/envs/pymc5_mamba/lib/python3.11/site-packages/pytensor/tensor/elemwise.py", line 485, in make_node
    out_dtypes, out_shapes, inputs = self.get_output_info(DimShuffle, *inputs)
                                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/mambaforge/envs/pymc5_mamba/lib/python3.11/site-packages/pytensor/tensor/elemwise.py", line 447, in get_output_info
    raise ValueError(
ValueError: Incompatible Elemwise input shapes [(100, 2), (1, 12), (100, 2)]

[...cut...]

…yet the code arrives at the end and the result looks OK.

What is the reason for such problems?

Note: the last line of the error message, “ValueError: Incompatible Elemwise input shapes [(100, 2), (1, 12), (100, 2)]” refers to the timeseries X which is (100,2) - and has 12 NaNs - and to the missing value tensor X_unobserved which is then of 12 elements.

I use a fresh install of PyMC+blackjax+numpyro as per the standard installation instructions, on Ubuntu 20.04.

This question/issue may be of interest to @ricardoV94 .

bug2.py (3.0 KB)

2

In addition to observing the issue on x86_64 Ubuntu 20.04, I can reproduce the exact same issue on M1 MacBook Pro too - via mambaforge for arm64.

3

I am a bit swamped this week, maybe worth opening an issue on PyMC? That way it won’t be accidentally forgotten/missed

1 Like
5

Done: BUG: Bayesian VAR with missing data: PyTensor error when sampling with blackjax or numpyro · Issue #6753 · pymc-devs/pymc · GitHub

2 Likes