PyMC Compilation Nesting with pytensor.scan

@ricardoV94,
That recommended article was useful. I wrapped my " Scan inside a CustomDist" as you recommended. Here’s the updated version of my code:

from pymc.pytensorf import collect_default_updates


def step(coef_mode, task_array, acc):
    # Create LogNormal prior for coefficient
    lognormal_std = coef_mode / 3  # Set the standard deviation of prior in lognormal space to 1/3 of the mode
    coef_sigma = np.sqrt(np.log(1 + (lognormal_std / coef_mode) ** 2))  # Convert lognormal std to normal std
    coef_mu = np.log(coef_mode) + coef_sigma ** 2  # for lognormal \mu = \ln(\text{Mode}) + \sigma^2
    coef_prior = pm.Lognormal.dist(mu=coef_mu, sigma=coef_sigma)

    new_acc = acc + coef_prior * task_array
    updates = collect_default_updates(inputs=[acc], outputs=[coef_prior, new_acc])

    return new_acc, updates


coef_modes_tensor = pt.as_tensor_variable(list(task_length_guesses.values()))
operation_arrays_tensor = pt.stack([pt.as_tensor_variable(df[task_col].values) for task_col in task_columns])
initial_acc = pt.zeros_like(operation_arrays_tensor[0], dtype='float64')


def custom_dist_func(coef_modes, task_arrays):
    expected_work_hours, _ = pytensor.scan(
        fn=step,
        sequences=[coef_modes, task_arrays],
        outputs_info=initial_acc
    )
    return expected_work_hours[-1]  # Return the final value


# Define a custom distribution using pm.CustomDist
with pm.Model() as model:
    
    custom_dist = pm.CustomDist(
        'custom_workhours',
        coef_modes_tensor,
        operation_arrays_tensor,
        dist=lambda coef_modes, task_arrays, size=None: custom_dist_func(coef_modes, task_arrays)
    )
    coef_sigma = pm.Exponential(name='modelsigma', lam=sigma_lam, transform=None)
    observed_workhours = pm.Normal(
        name='observed_workhours',
        mu=custom_dist,
        sigma=coef_sigma,
        observed=df['time_taken'].values,
    )

    trace = pm.sample(draws=2100, tune=1100, return_inferencedata=True, progressbar=True, cores=4,
                      init="adapt_diag", initvals=start)

Even after this change I’m getting the following error:

ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/opt/miniconda3/envs/myenv/lib/python3.9/site-packages/pytensor/graph/rewriting/basic.py", line 1922, in process_node
    replacements = node_rewriter.transform(fgraph, node)
  File "/opt/miniconda3/envs/myenv/lib/python3.9/site-packages/pytensor/graph/rewriting/basic.py", line 1081, in transform
    return self.fn(fgraph, node)
  File "/opt/miniconda3/envs/myenv/lib/python3.9/site-packages/pytensor/tensor/rewriting/basic.py", line 1110, in constant_folding
    thunk = node.op.make_thunk(node, storage_map, compute_map, no_recycling=[])
  File "/opt/miniconda3/envs/myenv/lib/python3.9/site-packages/pytensor/link/c/op.py", line 119, in make_thunk
    return self.make_c_thunk(node, storage_map, compute_map, no_recycling)
  File "/opt/miniconda3/envs/myenv/lib/python3.9/site-packages/pytensor/link/c/op.py", line 84, in make_c_thunk
    outputs = cl.make_thunk(
  File "/opt/miniconda3/envs/myenv/lib/python3.9/site-packages/pytensor/link/c/basic.py", line 1189, in make_thunk
    cthunk, module, in_storage, out_storage, error_storage = self.__compile__(
  File "/opt/miniconda3/envs/myenv/lib/python3.9/site-packages/pytensor/link/c/basic.py", line 1109, in __compile__
    thunk, module = self.cthunk_factory(
  File "/opt/miniconda3/envs/myenv/lib/python3.9/site-packages/pytensor/link/c/basic.py", line 1631, in cthunk_factory
    module = cache.module_from_key(key=key, lnk=self)
  File "/opt/miniconda3/envs/myenv/lib/python3.9/site-packages/pytensor/link/c/cmodule.py", line 1231, in module_from_key
    module = lnk.compile_cmodule(location)
  File "/opt/miniconda3/envs/myenv/lib/python3.9/site-packages/pytensor/link/c/basic.py", line 1532, in compile_cmodule
    module = c_compiler.compile_str(
  File "/opt/miniconda3/envs/myenv/lib/python3.9/site-packages/pytensor/link/c/cmodule.py", line 2639, in compile_str
    raise CompileError(
pytensor.link.c.exceptions.CompileError: Compilation failed (return status=1):
/usr/bin/clang++ -dynamiclib -g -O3 -fno-math-errno -Wno-unused-label -Wno-unused-variable -Wno-write-strings -Wno-c++11-narrowing -fno-exceptions -fno-unwind-tables -fno-asynchronous-unwind-tables -DNPY_NO_DEPRECATED_API=NPY_1_7_API_VERSION -m64 -fPIC -undefined dynamic_lookup -I/opt/miniconda3/envs/myenv/lib/python3.9/site-packages/numpy/core/include -I/opt/miniconda3/envs/myenv/include/python3.9 -I/opt/miniconda3/envs/myenv/lib/python3.9/site-packages/pytensor/link/c/c_code -L/opt/miniconda3/envs/myenv/lib -fvisibility=hidden -o /Users/myuser/.pytensor/compiledir_macOS-14.6.1-x86_64-i386-64bit-i386-3.9.19-64/tmpmfcfbpqd/m5827e5447ad45b3501581ba1dab81f701bd42767840fb02f1414980322660bb2.so /Users/myuser/.pytensor/compiledir_macOS-14.6.1-x86_64-i386-64bit-i386-3.9.19-64/tmpmfcfbpqd/mod.cpp
/Users/myuser/.pytensor/compiledir_macOS-14.6.1-x86_64-i386-64bit-i386-3.9.19-64/tmpmfcfbpqd/mod.cpp:5712:5: fatal error: bracket nesting level exceeded maximum of 256
    {Py_XINCREF(py_V507);}
    ^
/Users/myuser/.pytensor/compiledir_macOS-14.6.1-x86_64-i386-64bit-i386-3.9.19-64/tmpmfcfbpqd/mod.cpp:5712:5: note: use -fbracket-depth=N to increase maximum nesting level
1 error generated.