Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix bugs in local_useless_slice and local_reduce_join rewrites #1057

Merged
merged 2 commits into from
Nov 1, 2024

Conversation

ricardoV94
Copy link
Member

@ricardoV94 ricardoV94 commented Oct 31, 2024

@ricardoV94 ricardoV94 added bug Something isn't working graph rewriting labels Oct 31, 2024
@ricardoV94 ricardoV94 changed the title Fix bug in local_useless_slice Fix bug in local_useless_slice rewrite for negative steps Oct 31, 2024
@ricardoV94 ricardoV94 force-pushed the useless_slice_bug branch 2 times, most recently from b1778f3 to 711f94e Compare October 31, 2024 09:07
Copy link

codecov bot commented Oct 31, 2024

Codecov Report

All modified and coverable lines are covered by tests ✅

Please upload report for BASE (main@4f7d709). Learn more about missing BASE report.
Report is 4 commits behind head on main.

Additional details and impacted files

Impacted file tree graph

@@           Coverage Diff           @@
##             main    #1057   +/-   ##
=======================================
  Coverage        ?   81.95%           
=======================================
  Files           ?      182           
  Lines           ?    47864           
  Branches        ?     8621           
=======================================
  Hits            ?    39226           
  Misses          ?     6472           
  Partials        ?     2166           
Files with missing lines Coverage Δ
pytensor/tensor/rewriting/math.py 89.82% <100.00%> (ø)
pytensor/tensor/rewriting/subtensor.py 90.39% <100.00%> (ø)

@tomicapretto
Copy link

Not sure if this is related to PyMC or not but I can confirm the following

Example 1: works without issues

import pymc as pm

coords = {"threshold_dim": [0, 1]}

with pm.Model(coords=coords) as model:
    threshold = pm.Normal(
        "threshold",
        mu=[-2, 2],
        sigma=1,
        transform=pm.distributions.transforms.ordered,
        dims="threshold_dim",
    )

    idata = pm.sample()

Example 2: works without issues

import numpy as np
import pymc as pm
import pytensor.tensor as pt

coords = {
    "threshold_dim": [0, 1],
    "to_predict_dim": [0, 1, 2],
    "__obs__": [0, 1, 2],
}

predictor = np.array([1, 0, 1])
observed = np.array([0, 1, 2])

with pm.Model(coords=coords) as model:
    b_predictor = pm.Normal("b_predictor")
    threshold = pm.Normal(
        "threshold",
        mu=[-2, 2],
        sigma=1,
        transform=pm.distributions.transforms.ordered,
        dims="threshold_dim"
    )

    eta = b_predictor * np.array([1, 0, 1])
    eta_shifted = threshold - pt.shape_padright(eta)
    p = pm.math.sigmoid(eta_shifted)
    p = pt.concatenate(
        [
            pt.shape_padright(p[..., 0]),
            p[..., 1:] - p[..., :-1],
            pt.shape_padright(1 - p[..., -1]),
        ],
        axis=-1,
    )

    p = pm.Deterministic("p", p, dims=("__obs__", "to_predict_dim"))

    pm.Categorical("to_predict", p=p, observed=observed, dims="__obs__")
    idata = pm.sample()

Example 3: Seems to work, but there are some error messages. The difference is that we're not using dims.

import numpy as np
import pymc as pm
import pytensor.tensor as pt

coords = {
    "threshold_dim": [0, 1],
    "to_predict_dim": [0, 1, 2],
    "__obs__": [0, 1, 2],
}

predictor = np.array([1, 0, 1])
observed = np.array([0, 1, 2])

with pm.Model(coords=coords) as model:
    b_predictor = pm.Normal("b_predictor")
    threshold = pm.Normal(
        "threshold",
        mu=[-2, 2],
        sigma=1,
        transform=pm.distributions.transforms.ordered,
    )

    eta = b_predictor * np.array([1, 0, 1])
    eta_shifted = threshold - pt.shape_padright(eta)
    p = pm.math.sigmoid(eta_shifted)
    p = pt.concatenate(
        [
            pt.shape_padright(p[..., 0]),
            p[..., 1:] - p[..., :-1],
            pt.shape_padright(1 - p[..., -1]),
        ],
        axis=-1,
    )

    p = pm.Deterministic("p", p, dims=("__obs__", "to_predict_dim"))

    pm.Categorical("to_predict", p=p, observed=observed, dims="__obs__")

    idata = pm.sample()
threshold K (2,)
threshold_dim (2)
__obs__ (3)
threshold
~
Normal
to_predict
~
Categorical
predictor
~
Normal
/home/tomas/oss/bambinos/bambi/bambi/formula.py:111: UserWarning: The intercept is omitted in ordinal families
  warnings.warn("The intercept is omitted in ordinal families")
threshold K (2,)
'0+untagged.30830.g711f94e'
Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [threshold]
Sampling 4 chains, 0 divergences ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:00 / 0:00:01
Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 1 seconds.
ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_reduce_join
ERROR (pytensor.graph.rewriting.basic): node: Sum{axis=1}(Join.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1909, in process_node
    replacements = node_rewriter.transform(fgraph, node)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/math.py", line 1668, in local_reduce_join
    new_input = apply_local_dimshuffle_lift(None, inp.squeeze(join_axis))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 374, in apply_local_dimshuffle_lift
    new = local_dimshuffle_lift.transform(fgraph, var.owner)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tomas/miniconda3/envs/bambi-dev/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py", line 432, in local_dimshuffle_lift
    and (len(fgraph.clients[inp]) == 1)
             ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'clients'

Multiprocess sampling (4 chains in 4 jobs)
NUTS: [b_predictor, threshold]

@ricardoV94
Copy link
Member Author

Yeah I forgot that other bug you had found. I'll fix it here as well. It's something simple

Canonical slice start and stop values depend on the sign of the step. The rewrite wrongly assumed they were always 0:len(dim)
The helper `apply_local_dimshuffle_lift` requires a FunctionGraph when elemwise inputs are involved.
@ricardoV94
Copy link
Member Author

@tomicapretto pushed a commit that fixes that other bug

@ricardoV94 ricardoV94 changed the title Fix bug in local_useless_slice rewrite for negative steps Fix bugs in local_useless_slice and reduce_join rewrites Oct 31, 2024
@ricardoV94 ricardoV94 changed the title Fix bugs in local_useless_slice and reduce_join rewrites Fix bugs in local_useless_slice and local_reduce_join rewrites Oct 31, 2024
Copy link

@tomicapretto tomicapretto left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The changes fix the bug I reported.

@ricardoV94 ricardoV94 merged commit e73258b into pymc-devs:main Nov 1, 2024
61 checks passed
@ricardoV94 ricardoV94 deleted the useless_slice_bug branch November 1, 2024 17:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug Something isn't working graph rewriting
Projects
None yet
2 participants