Workarounds for `arviz` no longer supporting `MultiTrace`?

I recently updated from pymc3 to pymc v5. I’m also using arviz version 0.17.1. I have a model with a large number of parameters and due to memory considerations, I’m trying to stick with MultiTrace instead of InferenceData by using trace = pm.sample(return_inferencedata=False). In previous versions of arviz I could pass a MultiTrace to diagnostic functions like az.summary and az.plot_trace. Now I get

File ~/Documents/Research_Journal/MSP_Populations/pulse_shapes/single_profile_bayes_pymc.py:203, in gaussian_pwrlaw_k_prior_real_imag(profile_dict, tunesamp, target_accept, plot)
    197 map_estimate = pm.find_MAP(model=model)
    198 trace = pm.sample(10000,
    199                   tune=tunesamp,
    200                   target_accept=target_accept,
    201                   init="adapt_diag",
    202                   return_inferencedata=False)
--> 203 summ = az.summary(trace)
    204 print(az.summary(trace, round_to=2))
    205 trace.injected_profile = profile_dict

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/arviz/stats/stats.py:1345, in summary(data, var_names, filter_vars, group, fmt, kind, round_to, circ_var_names, stat_focus, stat_funcs, extend, hdi_prob, skipna, labeller, coords, index_origin, order)
   1343         raise TypeError(f"InferenceData does not contain group: {group}")
   1344 else:
-> 1345     dataset = convert_to_dataset(data, group="posterior")
   1346 var_names = _var_names(var_names, dataset, filter_vars)
   1347 dataset = dataset if var_names is None else dataset[var_names]

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/arviz/data/converters.py:176, in convert_to_dataset(obj, group, coords, dims)
    138 def convert_to_dataset(obj, *, group="posterior", coords=None, dims=None):
    139     """Convert a supported object to an xarray dataset.
    140 
    141     This function is idempotent, in that it will return xarray.Dataset functions
   (...)
    174     xarray.Dataset
    175     """
--> 176     inference_data = convert_to_inference_data(obj, group=group, coords=coords, dims=dims)
    177     dataset = getattr(inference_data, group, None)
    178     if dataset is None:

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/arviz/data/converters.py:130, in convert_to_inference_data(obj, group, coords, dims, **kwargs)
    116 else:
    117     allowable_types = (
    118         "xarray dataarray",
    119         "xarray dataset",
   (...)
    128         "cmdstanpy fit",
    129     )
--> 130     raise ValueError(
    131         f'Can only convert {", ".join(allowable_types)} to InferenceData, '
    132         f"not {obj.__class__.__name__}"
    133     )
    135 return InferenceData(**{group: dataset})

ValueError: Can only convert xarray dataarray, xarray dataset, dict, netcdf filename, numpy array, pystan fit, emcee fit, pyro mcmc fit, numpyro mcmc fit, cmdstan fit csv filename, cmdstanpy fit to InferenceData, not MultiTrace

And it looks like newer versions of arviz no longer contain from_pymc. Is there a version of arviz that works with pymc v5 and still supports MultiTrace? Are there any other workarounds?

There is a to_inference_data function in PyMC. Any reason do you want to stick to MultiTrace data? Using InferenceData has been the recomended way even before pymc 5.x

The main reason is memory limitations, the InferenceData object is ~100 MB, but I suspect there might be a memory leak (or something is being cached?) since the size grows in subsequent runs of the sampler. I frequently run into OSError: [Errno 12] Cannot allocate memory (full traceback below) when using InferenceData but I have no issues when using MultiTrace. A less important reason is needing to refactor my analysis pipelines to use the syntax associated with InferenceData.

Problem occurred during compilation with the command line below:
/lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/bin/g++ -shared -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 -march=alderlake -mmmx -mpopcnt -msse -msse2 -msse3 -mssse3 -msse4.1 -msse4.2 -mavx -mavx2 -mno-sse4a -mno-fma4 -mno-xop -mfma -mno-avx512f -mbmi -mbmi2 -maes -mpclmul -mno-avx512vl -mno-avx512bw -mno-avx512dq -mno-avx512cd -mno-avx512er -mno-avx512pf -mno-avx512vbmi -mno-avx512ifma -mno-avx5124vnniw -mno-avx5124fmaps -mno-avx512vpopcntdq -mno-avx512vbmi2 -mgfni -mvpclmulqdq -mno-avx512vnni -mno-avx512bitalg -mno-avx512bf16 -mno-avx512vp2intersect -mno-3dnow -madx -mabm -mno-cldemote -mclflushopt -mclwb -mno-clzero -mcx16 -mno-enqcmd -mf16c -mfsgsbase -mfxsr -mno-hle -msahf -mno-lwp -mlzcnt -mmovbe -mmovdir64b -mmovdiri -mno-mwaitx -mpconfig -mpku -mno-prefetchwt1 -mprfchw -mptwrite -mrdpid -mrdrnd -mrdseed -mno-rtm -mserialize -mno-sgx -msha -mshstk -mno-tbm -mno-tsxldtrk -mvaes -mwaitpkg -mno-wbnoinvd -mxsave -mxsavec -mxsaveopt -mxsaves -mno-amx-tile -mno-amx-int8 -mno-amx-bf16 -mno-uintr -mhreset -mno-kl -mno-widekl -mavxvnni -mno-avx512fp16 --param l1-cache-size=32 --param l1-cache-line-size=64 --param l2-cache-size=25600 -mtune=alderlake -DNPY_NO_DEPRECATED_API=NPY_1_7_API_VERSION -m64 -fPIC -I/lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/numpy/core/include -I/lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/include/python3.9 -I/lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/c/c_code -L/lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib -fvisibility=hidden -o /users/tcohen/.pytensor/compiledir_Linux-4.18-el8_10.x86_64-x86_64-with-glibc2.28-x86_64-3.9.21-64/tmpo0be69ms/me00b73b710beaa3c6ed4d22d8941ba754bf43854f31f406a769cc9da89091508.so /users/tcohen/.pytensor/compiledir_Linux-4.18-el8_10.x86_64-x86_64-with-glibc2.28-x86_64-3.9.21-64/tmpo0be69ms/mod.cpp -lpython3.9
ERROR (pytensor.link.c.cmodule): [Errno 12] Cannot allocate memory
---------------------------------------------------------------------------
OSError                                   Traceback (most recent call last)
File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/vm.py:1235, in VMLinker.make_all(self, profiler, input_storage, output_storage, storage_map)
   1231 # no-recycling is done at each VM.__call__ So there is
   1232 # no need to cause duplicate c code by passing
   1233 # no_recycling here.
   1234 thunks.append(
-> 1235     node.op.make_thunk(node, storage_map, compute_map, [], impl=impl)
   1236 )
   1237 linker_make_thunk_time[node] = time.perf_counter() - thunk_start

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/c/op.py:119, in COp.make_thunk(self, node, storage_map, compute_map, no_recycling, impl)
    118 try:
--> 119     return self.make_c_thunk(node, storage_map, compute_map, no_recycling)
    120 except (NotImplementedError, MethodNotDefined):
    121     # We requested the c code, so don't catch the error.

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/c/op.py:84, in COp.make_c_thunk(self, node, storage_map, compute_map, no_recycling)
     83         raise NotImplementedError("float16")
---> 84 outputs = cl.make_thunk(
     85     input_storage=node_input_storage, output_storage=node_output_storage
     86 )
     87 thunk, node_input_filters, node_output_filters = outputs

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/c/basic.py:1189, in CLinker.make_thunk(self, input_storage, output_storage, storage_map, cache, **kwargs)
   1188 init_tasks, tasks = self.get_init_tasks()
-> 1189 cthunk, module, in_storage, out_storage, error_storage = self.__compile__(
   1190     input_storage, output_storage, storage_map, cache
   1191 )
   1193 res = _CThunk(cthunk, init_tasks, tasks, error_storage, module)

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/c/basic.py:1109, in CLinker.__compile__(self, input_storage, output_storage, storage_map, cache)
   1108 output_storage = tuple(output_storage)
-> 1109 thunk, module = self.cthunk_factory(
   1110     error_storage,
   1111     input_storage,
   1112     output_storage,
   1113     storage_map,
   1114     cache,
   1115 )
   1116 return (
   1117     thunk,
   1118     module,
   (...)
   1127     error_storage,
   1128 )

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/c/basic.py:1631, in CLinker.cthunk_factory(self, error_storage, in_storage, out_storage, storage_map, cache)
   1630         cache = get_module_cache()
-> 1631     module = cache.module_from_key(key=key, lnk=self)
   1633 vars = self.inputs + self.outputs + self.orphans

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/c/cmodule.py:1231, in ModuleCache.module_from_key(self, key, lnk)
   1230 location = dlimport_workdir(self.dirname)
-> 1231 module = lnk.compile_cmodule(location)
   1232 name = module.__file__

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/c/basic.py:1532, in CLinker.compile_cmodule(self, location)
   1531     _logger.debug(f"LOCATION {location}")
-> 1532     module = c_compiler.compile_str(
   1533         module_name=mod.code_hash,
   1534         src_code=src_code,
   1535         location=location,
   1536         include_dirs=self.header_dirs(),
   1537         lib_dirs=self.lib_dirs(),
   1538         libs=libs,
   1539         preargs=preargs,
   1540     )
   1541 except Exception as e:

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/c/cmodule.py:2589, in GCC_compiler.compile_str(module_name, src_code, location, include_dirs, lib_dirs, libs, preargs, py_module, hide_symbols)
   2588 try:
-> 2589     p_out = output_subprocess_Popen(cmd)
   2590     compile_stderr = p_out[1].decode()

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/utils.py:209, in output_subprocess_Popen(command, **params)
    208 params["stderr"] = subprocess.PIPE
--> 209 p = subprocess_Popen(command, **params)
    210 # we need to use communicate to make sure we don't deadlock around
    211 # the stdout/stderr pipe.

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/utils.py:175, in subprocess_Popen(command, **params)
    174 try:
--> 175     proc = subprocess.Popen(command, startupinfo=startupinfo, **params)
    176 finally:

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/subprocess.py:951, in Popen.__init__(self, args, bufsize, executable, stdin, stdout, stderr, preexec_fn, close_fds, shell, cwd, env, universal_newlines, startupinfo, creationflags, restore_signals, start_new_session, pass_fds, user, group, extra_groups, encoding, errors, text, umask)
    948             self.stderr = io.TextIOWrapper(self.stderr,
    949                     encoding=encoding, errors=errors)
--> 951     self._execute_child(args, executable, preexec_fn, close_fds,
    952                         pass_fds, cwd, env,
    953                         startupinfo, creationflags, shell,
    954                         p2cread, p2cwrite,
    955                         c2pread, c2pwrite,
    956                         errread, errwrite,
    957                         restore_signals,
    958                         gid, gids, uid, umask,
    959                         start_new_session)
    960 except:
    961     # Cleanup if the child failed starting.

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/subprocess.py:1770, in Popen._execute_child(self, args, executable, preexec_fn, close_fds, pass_fds, cwd, env, startupinfo, creationflags, shell, p2cread, p2cwrite, c2pread, c2pwrite, errread, errwrite, restore_signals, gid, gids, uid, umask, start_new_session)
   1769 fds_to_keep.add(errpipe_write)
-> 1770 self.pid = _posixsubprocess.fork_exec(
   1771         args, executable_list,
   1772         close_fds, tuple(sorted(map(int, fds_to_keep))),
   1773         cwd, env_list,
   1774         p2cread, p2cwrite, c2pread, c2pwrite,
   1775         errread, errwrite,
   1776         errpipe_read, errpipe_write,
   1777         restore_signals, start_new_session,
   1778         gid, gids, uid, umask,
   1779         preexec_fn)
   1780 self._child_created = True

OSError: [Errno 12] Cannot allocate memory

During handling of the above exception, another exception occurred:

OSError                                   Traceback (most recent call last)
Cell In[11], line 1
----> 1 single_profile_bayes_pymc.test_widebandportrait_prior_real_imag()

File ~/Documents/Research_Journal/MSP_Populations/pulse_shapes/single_profile_bayes_pymc.py:145, in test_widebandportrait_prior_real_imag(td_noise_sigma, p0, nbins, plot)
    138 uscale_smooth = SinglePulse(profile, period=p0).getUscale()
    139 injected_profile = {"weff_smooth" : weff_smooth,
    140                     "uscale_smooth" : uscale_smooth,
    141                     "spin_period" : nbins,
    142                     "profile" : profile,
    143                     "noise_sigma" : noise_sigma,
    144                     "noiseless_profile" : noiseless_prof}
--> 145 trace = gaussian_pwrlaw_k_prior_real_imag(injected_profile, tunesamp=4000,
    146                                          plot=plot)
    147 return trace

File ~/Documents/Research_Journal/MSP_Populations/pulse_shapes/single_profile_bayes_pymc.py:201, in gaussian_pwrlaw_k_prior_real_imag(profile_dict, tunesamp, target_accept, plot)
    196 Y_obs_imag = pm.Normal('Y_obs_imag',
    197                        mu=b_k,
    198                        sigma=noise_sigma,
    199                        observed=Uobs_imag)
    200 # Compute the Maximum A Posteriori estimate
--> 201 map_estimate = pm.find_MAP(model=model)
    202 trace = pm.sample(10000,
    203                   tune=tunesamp,
    204                   target_accept=target_accept,
    205                   init="adapt_diag")
    206 with open("single_profile_bayes_pymc.map", "wb") as f:

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pymc/tuning/starting.py:127, in find_MAP(start, vars, method, return_raw, include_transformed, progressbar, maxeval, model, seed, *args, **kwargs)
    120 ipfn = make_initial_point_fn(
    121     model=model,
    122     jitter_rvs=set(),
    123     return_transformed=True,
    124     overrides=start,
    125 )
    126 start = ipfn(seed)
--> 127 model.check_start_vals(start)
    129 vars_dict = {var.name: var for var in vars}
    130 x0 = DictToArrayBijection.map(
    131     {var_name: value for var_name, value in start.items() if var_name in vars_dict}
    132 )

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pymc/model/core.py:1678, in Model.check_start_vals(self, start)
   1672     valid_keys = ", ".join(value_names_set)
   1673     raise KeyError(
   1674         "Some start parameters do not appear in the model!\n"
   1675         f"Valid keys are: {valid_keys}, but {extra_keys} was supplied"
   1676     )
-> 1678 initial_eval = self.point_logps(point=elem)
   1680 if not all(np.isfinite(v) for v in initial_eval.values()):
   1681     raise SamplingError(
   1682         "Initial evaluation of model at starting point failed!\n"
   1683         f"Starting values:\n{elem}\n\n"
   1684         f"Logp initial evaluation results:\n{initial_eval}\n"
   1685         "You can call `model.debug()` for more details."
   1686     )

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pymc/model/core.py:1713, in Model.point_logps(self, point, round_vals)
   1707 factors = self.basic_RVs + self.potentials
   1708 factor_logps_fn = [pt.sum(factor) for factor in self.logp(factors, sum=False)]
   1709 return {
   1710     factor.name: np.round(np.asarray(factor_logp), round_vals)
   1711     for factor, factor_logp in zip(
   1712         factors,
-> 1713         self.compile_fn(factor_logps_fn)(point),
   1714     )
   1715 }

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pymc/model/core.py:1552, in Model.compile_fn(self, outs, inputs, mode, point_fn, **kwargs)
   1549     inputs = inputvars(outs)
   1551 with self:
-> 1552     fn = compile_pymc(
   1553         inputs,
   1554         outs,
   1555         allow_input_downcast=True,
   1556         accept_inplace=True,
   1557         mode=mode,
   1558         **kwargs,
   1559     )
   1561 if point_fn:
   1562     return PointFunc(fn)

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pymc/pytensorf.py:985, in compile_pymc(inputs, outputs, random_seed, mode, **kwargs)
    983 opt_qry = mode.provided_optimizer.including("random_make_inplace", check_parameter_opt)
    984 mode = Mode(linker=mode.linker, optimizer=opt_qry)
--> 985 pytensor_function = pytensor.function(
    986     inputs,
    987     outputs,
    988     updates={**rng_updates, **kwargs.pop("updates", {})},
    989     mode=mode,
    990     **kwargs,
    991 )
    992 return pytensor_function

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/compile/function/__init__.py:315, in function(inputs, outputs, mode, updates, givens, no_default_updates, accept_inplace, name, rebuild_strict, allow_input_downcast, profile, on_unused_input)
    309     fn = orig_function(
    310         inputs, outputs, mode=mode, accept_inplace=accept_inplace, name=name
    311     )
    312 else:
    313     # note: pfunc will also call orig_function -- orig_function is
    314     #      a choke point that all compilation must pass through
--> 315     fn = pfunc(
    316         params=inputs,
    317         outputs=outputs,
    318         mode=mode,
    319         updates=updates,
    320         givens=givens,
    321         no_default_updates=no_default_updates,
    322         accept_inplace=accept_inplace,
    323         name=name,
    324         rebuild_strict=rebuild_strict,
    325         allow_input_downcast=allow_input_downcast,
    326         on_unused_input=on_unused_input,
    327         profile=profile,
    328         output_keys=output_keys,
    329     )
    330 return fn

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/compile/function/pfunc.py:469, in pfunc(params, outputs, mode, updates, givens, no_default_updates, accept_inplace, name, rebuild_strict, allow_input_downcast, profile, on_unused_input, output_keys, fgraph)
    455     profile = ProfileStats(message=profile)
    457 inputs, cloned_outputs = construct_pfunc_ins_and_outs(
    458     params,
    459     outputs,
   (...)
    466     fgraph=fgraph,
    467 )
--> 469 return orig_function(
    470     inputs,
    471     cloned_outputs,
    472     mode,
    473     accept_inplace=accept_inplace,
    474     name=name,
    475     profile=profile,
    476     on_unused_input=on_unused_input,
    477     output_keys=output_keys,
    478     fgraph=fgraph,
    479 )

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/compile/function/types.py:1762, in orig_function(inputs, outputs, mode, accept_inplace, name, profile, on_unused_input, output_keys, fgraph)
   1750     m = Maker(
   1751         inputs,
   1752         outputs,
   (...)
   1759         fgraph=fgraph,
   1760     )
   1761     with config.change_flags(compute_test_value="off"):
-> 1762         fn = m.create(defaults)
   1763 finally:
   1764     if profile and fn:

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/compile/function/types.py:1654, in FunctionMaker.create(self, input_storage, storage_map)
   1651 start_import_time = pytensor.link.c.cmodule.import_time
   1653 with config.change_flags(traceback__limit=config.traceback__compile_limit):
-> 1654     _fn, _i, _o = self.linker.make_thunk(
   1655         input_storage=input_storage_lists, storage_map=storage_map
   1656     )
   1658 end_linker = time.perf_counter()
   1660 linker_time = end_linker - start_linker

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/basic.py:245, in LocalLinker.make_thunk(self, input_storage, output_storage, storage_map, **kwargs)
    238 def make_thunk(
    239     self,
    240     input_storage: Optional["InputStorageType"] = None,
   (...)
    243     **kwargs,
    244 ) -> tuple["BasicThunkType", "InputStorageType", "OutputStorageType"]:
--> 245     return self.make_all(
    246         input_storage=input_storage,
    247         output_storage=output_storage,
    248         storage_map=storage_map,
    249     )[:3]

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/vm.py:1244, in VMLinker.make_all(self, profiler, input_storage, output_storage, storage_map)
   1242             thunks[-1].lazy = False
   1243     except Exception:
-> 1244         raise_with_op(fgraph, node)
   1246 t1 = time.perf_counter()
   1248 if self.profile:

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/utils.py:531, in raise_with_op(fgraph, node, thunk, exc_info, storage_map)
    526     warnings.warn(
    527         f"{exc_type} error does not allow us to add an extra error message"
    528     )
    529     # Some exception need extra parameter in inputs. So forget the
    530     # extra long error message in that case.
--> 531 raise exc_value.with_traceback(exc_trace)

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/vm.py:1235, in VMLinker.make_all(self, profiler, input_storage, output_storage, storage_map)
   1230 thunk_start = time.perf_counter()
   1231 # no-recycling is done at each VM.__call__ So there is
   1232 # no need to cause duplicate c code by passing
   1233 # no_recycling here.
   1234 thunks.append(
-> 1235     node.op.make_thunk(node, storage_map, compute_map, [], impl=impl)
   1236 )
   1237 linker_make_thunk_time[node] = time.perf_counter() - thunk_start
   1238 if not hasattr(thunks[-1], "lazy"):
   1239     # We don't want all ops maker to think about lazy Ops.
   1240     # So if they didn't specify that its lazy or not, it isn't.
   1241     # If this member isn't present, it will crash later.

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/c/op.py:119, in COp.make_thunk(self, node, storage_map, compute_map, no_recycling, impl)
    115 self.prepare_node(
    116     node, storage_map=storage_map, compute_map=compute_map, impl="c"
    117 )
    118 try:
--> 119     return self.make_c_thunk(node, storage_map, compute_map, no_recycling)
    120 except (NotImplementedError, MethodNotDefined):
    121     # We requested the c code, so don't catch the error.
    122     if impl == "c":

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/c/op.py:84, in COp.make_c_thunk(self, node, storage_map, compute_map, no_recycling)
     82         print(f"Disabling C code for {self} due to unsupported float16")
     83         raise NotImplementedError("float16")
---> 84 outputs = cl.make_thunk(
     85     input_storage=node_input_storage, output_storage=node_output_storage
     86 )
     87 thunk, node_input_filters, node_output_filters = outputs
     89 @is_cthunk_wrapper_type
     90 def rval():

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/c/basic.py:1189, in CLinker.make_thunk(self, input_storage, output_storage, storage_map, cache, **kwargs)
   1154 """Compile this linker's `self.fgraph` and return a function that performs the computations.
   1155 
   1156 The return values can be used as follows:
   (...)
   1186 
   1187 """
   1188 init_tasks, tasks = self.get_init_tasks()
-> 1189 cthunk, module, in_storage, out_storage, error_storage = self.__compile__(
   1190     input_storage, output_storage, storage_map, cache
   1191 )
   1193 res = _CThunk(cthunk, init_tasks, tasks, error_storage, module)
   1194 res.nodes = self.node_order

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/c/basic.py:1109, in CLinker.__compile__(self, input_storage, output_storage, storage_map, cache)
   1107 input_storage = tuple(input_storage)
   1108 output_storage = tuple(output_storage)
-> 1109 thunk, module = self.cthunk_factory(
   1110     error_storage,
   1111     input_storage,
   1112     output_storage,
   1113     storage_map,
   1114     cache,
   1115 )
   1116 return (
   1117     thunk,
   1118     module,
   (...)
   1127     error_storage,
   1128 )

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/c/basic.py:1631, in CLinker.cthunk_factory(self, error_storage, in_storage, out_storage, storage_map, cache)
   1629     if cache is None:
   1630         cache = get_module_cache()
-> 1631     module = cache.module_from_key(key=key, lnk=self)
   1633 vars = self.inputs + self.outputs + self.orphans
   1634 # List of indices that should be ignored when passing the arguments
   1635 # (basically, everything that the previous call to uniq eliminated)

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/c/cmodule.py:1231, in ModuleCache.module_from_key(self, key, lnk)
   1229 try:
   1230     location = dlimport_workdir(self.dirname)
-> 1231     module = lnk.compile_cmodule(location)
   1232     name = module.__file__
   1233     assert name.startswith(location)

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/c/basic.py:1532, in CLinker.compile_cmodule(self, location)
   1530 try:
   1531     _logger.debug(f"LOCATION {location}")
-> 1532     module = c_compiler.compile_str(
   1533         module_name=mod.code_hash,
   1534         src_code=src_code,
   1535         location=location,
   1536         include_dirs=self.header_dirs(),
   1537         lib_dirs=self.lib_dirs(),
   1538         libs=libs,
   1539         preargs=preargs,
   1540     )
   1541 except Exception as e:
   1542     e.args += (str(self.fgraph),)

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/link/c/cmodule.py:2589, in GCC_compiler.compile_str(module_name, src_code, location, include_dirs, lib_dirs, libs, preargs, py_module, hide_symbols)
   2586     print(" ".join(cmd), file=sys.stderr)
   2588 try:
-> 2589     p_out = output_subprocess_Popen(cmd)
   2590     compile_stderr = p_out[1].decode()
   2591 except Exception:
   2592     # An exception can occur e.g. if `g++` is not found.

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/utils.py:209, in output_subprocess_Popen(command, **params)
    207 params["stdout"] = subprocess.PIPE
    208 params["stderr"] = subprocess.PIPE
--> 209 p = subprocess_Popen(command, **params)
    210 # we need to use communicate to make sure we don't deadlock around
    211 # the stdout/stderr pipe.
    212 out = p.communicate()

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/site-packages/pytensor/utils.py:175, in subprocess_Popen(command, **params)
    172     params["stdin"] = stdin.fileno()
    174 try:
--> 175     proc = subprocess.Popen(command, startupinfo=startupinfo, **params)
    176 finally:
    177     if stdin is not None:

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/subprocess.py:951, in Popen.__init__(self, args, bufsize, executable, stdin, stdout, stderr, preexec_fn, close_fds, shell, cwd, env, universal_newlines, startupinfo, creationflags, restore_signals, start_new_session, pass_fds, user, group, extra_groups, encoding, errors, text, umask)
    947         if self.text_mode:
    948             self.stderr = io.TextIOWrapper(self.stderr,
    949                     encoding=encoding, errors=errors)
--> 951     self._execute_child(args, executable, preexec_fn, close_fds,
    952                         pass_fds, cwd, env,
    953                         startupinfo, creationflags, shell,
    954                         p2cread, p2cwrite,
    955                         c2pread, c2pwrite,
    956                         errread, errwrite,
    957                         restore_signals,
    958                         gid, gids, uid, umask,
    959                         start_new_session)
    960 except:
    961     # Cleanup if the child failed starting.
    962     for f in filter(None, (self.stdin, self.stdout, self.stderr)):

File /lustre/aoc/students/tcohen/miniconda3/envs/pulse_shapes_pm5/lib/python3.9/subprocess.py:1770, in Popen._execute_child(self, args, executable, preexec_fn, close_fds, pass_fds, cwd, env, startupinfo, creationflags, shell, p2cread, p2cwrite, c2pread, c2pwrite, errread, errwrite, restore_signals, gid, gids, uid, umask, start_new_session)
   1768     fds_to_keep = set(pass_fds)
   1769     fds_to_keep.add(errpipe_write)
-> 1770     self.pid = _posixsubprocess.fork_exec(
   1771             args, executable_list,
   1772             close_fds, tuple(sorted(map(int, fds_to_keep))),
   1773             cwd, env_list,
   1774             p2cread, p2cwrite, c2pread, c2pwrite,
   1775             errread, errwrite,
   1776             errpipe_read, errpipe_write,
   1777             restore_signals, start_new_session,
   1778             gid, gids, uid, umask,
   1779             preexec_fn)
   1780     self._child_created = True
   1781 finally:
   1782     # be sure the FD is closed no matter what

OSError: [Errno 12] Cannot allocate memory
Apply node that caused the error: Composite{((-0.5 * sqr((0.8838834764831842 * (i0 - i1)))) - 1.0423685721704357)}(Y_obs_real{[ 2.904211 ... 99220e-01]}, a_k)
Toposort index: 4
Inputs types: [TensorType(float64, shape=(129,)), TensorType(float64, shape=(129,))]

HINT: Use a linker other than the C linker to print the inputs' shapes and strides.
HINT: Re-running with most PyTensor optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the PyTensor flag 'optimizer=fast_compile'. If that does not work, PyTensor optimizations can be disabled with 'optimizer=None'.
HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.

Are you on the latest version of pymc? Can you share a reproducible snippet that shows the memory leakage?

It would be surprising that switching the backend would lead to leakage

Do we compute deterministics when return_inferencedata=False? That’s usually the culprit for out-of-memory errors

We do. Not sure about the log_likelihood, maybe that’s different?