BUG: sampling with Numba backend fails (BUG IN FGRAPH.REPLACE OR A LISTENER)

[fonnesbeck]

Describe the issue:

Models that sample without error using the default backend fail when using Numba via compile_kwargs. Details below.

Reproduceable code example:

Running any of the models in this notebook after modifying pm.sample as follows:

with basic_ordered_logistic_model:
    basic_ordered_logistic_inference = pm.sample(compile_kwargs=dict(mode="NUMBA"))

generates the error below.

Error message:

Initializing NUTS using jitter+adapt_diag...
<<!! BUG IN FGRAPH.REPLACE OR A LISTENER !!>> <class 'TypeError'> Cannot convert Type Tensor3(float64, shape=(?, 1, ?)) (of Variable Reshape{3}.0) into Type Matrix(float64, shape=(?, 1)). You can try to manually convert Reshape{3}.0 into a Matrix(float64, shape=(?, 1)). ravel_multidimensional_int_idx
ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: ravel_multidimensional_int_idx
ERROR (pytensor.graph.rewriting.basic): node: AdvancedSubtensor(Sub.0, ExpandDims{axis=1}.0, [[3] [2] [ ... ] [1] [1]])
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/var/home/fonnesbeck/repos/pymc-examples/.pixi/envs/default/lib/python3.12/site-packages/pytensor/graph/rewriting/basic.py", line 1959, in process_node
    fgraph.replace_all_validate_remove(  # type: ignore
  File "/var/home/fonnesbeck/repos/pymc-examples/.pixi/envs/default/lib/python3.12/site-packages/pytensor/graph/features.py", line 627, in replace_all_validate_remove
    chk = fgraph.replace_all_validate(replacements, reason=reason, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/var/home/fonnesbeck/repos/pymc-examples/.pixi/envs/default/lib/python3.12/site-packages/pytensor/graph/features.py", line 572, in replace_all_validate
    fgraph.replace(r, new_r, reason=reason, verbose=False, **kwargs)
  File "/var/home/fonnesbeck/repos/pymc-examples/.pixi/envs/default/lib/python3.12/site-packages/pytensor/graph/fg.py", line 498, in replace
    new_var = var.type.filter_variable(new_var, allow_convert=True)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/var/home/fonnesbeck/repos/pymc-examples/.pixi/envs/default/lib/python3.12/site-packages/pytensor/tensor/type.py", line 278, in filter_variable
    raise TypeError(
TypeError: Cannot convert Type Tensor3(float64, shape=(?, 1, ?)) (of Variable Reshape{3}.0) into Type Matrix(float64, shape=(?, 1)). You can try to manually convert Reshape{3}.0 into a Matrix(float64, shape=(?, 1)).

<<!! BUG IN FGRAPH.REPLACE OR A LISTENER !!>> <class 'TypeError'> Cannot convert Type Tensor3(float64, shape=(?, 1, ?)) (of Variable Reshape{3}.0) into Type Matrix(float64, shape=(?, 1)). You can try to manually convert Reshape{3}.0 into a Matrix(float64, shape=(?, 1)). ravel_multidimensional_int_idx
ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: ravel_multidimensional_int_idx
ERROR (pytensor.graph.rewriting.basic): node: AdvancedSubtensor(Sub.0, ExpandDims{axis=1}.0, [[3] [2] [ ... ] [1] [1]])
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/var/home/fonnesbeck/repos/pymc-examples/.pixi/envs/default/lib/python3.12/site-packages/pytensor/graph/rewriting/basic.py", line 1959, in process_node
    fgraph.replace_all_validate_remove(  # type: ignore
  File "/var/home/fonnesbeck/repos/pymc-examples/.pixi/envs/default/lib/python3.12/site-packages/pytensor/graph/features.py", line 627, in replace_all_validate_remove
    chk = fgraph.replace_all_validate(replacements, reason=reason, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/var/home/fonnesbeck/repos/pymc-examples/.pixi/envs/default/lib/python3.12/site-packages/pytensor/graph/features.py", line 572, in replace_all_validate
    fgraph.replace(r, new_r, reason=reason, verbose=False, **kwargs)
  File "/var/home/fonnesbeck/repos/pymc-examples/.pixi/envs/default/lib/python3.12/site-packages/pytensor/graph/fg.py", line 498, in replace
    new_var = var.type.filter_variable(new_var, allow_convert=True)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/var/home/fonnesbeck/repos/pymc-examples/.pixi/envs/default/lib/python3.12/site-packages/pytensor/tensor/type.py", line 278, in filter_variable
    raise TypeError(
TypeError: Cannot convert Type Tensor3(float64, shape=(?, 1, ?)) (of Variable Reshape{3}.0) into Type Matrix(float64, shape=(?, 1)). You can try to manually convert Reshape{3}.0 into a Matrix(float64, shape=(?, 1)).

<<!! BUG IN FGRAPH.REPLACE OR A LISTENER !!>> <class 'TypeError'> Cannot convert Type Tensor3(float64, shape=(?, 1, ?)) (of Variable Reshape{3}.0) into Type Matrix(float64, shape=(?, 1)). You can try to manually convert Reshape{3}.0 into a Matrix(float64, shape=(?, 1)). ravel_multidimensional_int_idx
ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: ravel_multidimensional_int_idx
ERROR (pytensor.graph.rewriting.basic): node: AdvancedSubtensor(Sub.0, ExpandDims{axis=1}.0, [[3] [2] [ ... ] [1] [1]])
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/var/home/fonnesbeck/repos/pymc-examples/.pixi/envs/default/lib/python3.12/site-packages/pytensor/graph/rewriting/basic.py", line 1959, in process_node
    fgraph.replace_all_validate_remove(  # type: ignore
  File "/var/home/fonnesbeck/repos/pymc-examples/.pixi/envs/default/lib/python3.12/site-packages/pytensor/graph/features.py", line 627, in replace_all_validate_remove
    chk = fgraph.replace_all_validate(replacements, reason=reason, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/var/home/fonnesbeck/repos/pymc-examples/.pixi/envs/default/lib/python3.12/site-packages/pytensor/graph/features.py", line 572, in replace_all_validate
    fgraph.replace(r, new_r, reason=reason, verbose=False, **kwargs)
  File "/var/home/fonnesbeck/repos/pymc-examples/.pixi/envs/default/lib/python3.12/site-packages/pytensor/graph/fg.py", line 498, in replace
    new_var = var.type.filter_variable(new_var, allow_convert=True)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/var/home/fonnesbeck/repos/pymc-examples/.pixi/envs/default/lib/python3.12/site-packages/pytensor/tensor/type.py", line 278, in filter_variable
    raise TypeError(
TypeError: Cannot convert Type Tensor3(float64, shape=(?, 1, ?)) (of Variable Reshape{3}.0) into Type Matrix(float64, shape=(?, 1)). You can try to manually convert Reshape{3}.0 into a Matrix(float64, shape=(?, 1)).

Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_action, beta_intent, beta_contact, alpha]

### PyMC version information:

PyMC 5.19.1
PyTensor 2.26.12
OS: Fedora
Python: 3.12.5
install: conda

### Context for the issue:

_No response_

[ricardoV94]

That's a pytensor bug, moving the issue to the repo

[danieltomasz]

I have similar issue #1134, update of base conda and clean reinstall helped

[fonnesbeck]

Thanks, @ricardoV94!