Thanks for this. Its now letting me work with the LKJcorr, however I can’t seem to insert the covariance into MvNormal. From the link you sent, it looks like batched pm.MvNormal is not yet working?
In this post, you mentioned that pm.MvNormal accepts batched params.
From my attempt below, am I doing something wrong? Or is the batched MvNormal not yet working?
from pymc.distributions.transforms import Interval
class MultivariateIntervalTransform(Interval):
name = "interval"
def log_jac_det(self, *args):
return super().log_jac_det(*args).sum(-1)
tr = MultivariateIntervalTransform(-1.0, 1.0)
with pm.Model(coords=coords) as m:
corr_packed = pm.LKJCorr("x", n=n_groups, eta=1, transform=tr)
sigma = pm.HalfCauchy("sigma", 1, dims=("groups", "predictors"))
triu_idx = pt.triu_indices(n_predictors, k=1)
corr_upper = pt.set_subtensor(pt.zeros((n_predictors, n_predictors))[triu_idx], corr_packed)
corr = pm.Deterministic("corr", pt.eye(n_predictors) + corr_upper + corr_upper.T, dims=("predictors", "predictors_I"))
# n_groups x n_predictors x n_predictors
sigma_diag = pt.stack([pt.eye(n_predictors) * sigma[i] for i in range(n_groups)])
cov = pm.Deterministic("cov",
pt.stack([sigma_diag[i] @ corr @ sigma_diag[i] \
for i in range(n_groups)]),
dims=("groups", "predictors", "predictors_I"))
# batched dim on the far left
# cov[group]: N x n_predictors x n_predictors (n_groups disappear)
# obs: N x n_predictors (n_groups disappear)
likelihood = pm.MvNormal("likelihood", mu=0, cov=cov[group], observed=obs, dims=("obs_id", "predictors"))
with m:
idata_prior = pm.sample_prior_predictive()
# idata = pm.sample(tune=1000, draws=1000, nuts_sampler="numpyro")
{
"name": "IndexError",
"message": "boolean index did not match indexed array along dimension 0; dimension is 2 but corresponding boolean dimension is 100
Apply node that caused the error: AdvancedSubtensor(cov, [ True Fal ... lse True])
Toposort index: 22
Inputs types: [TensorType(float64, shape=(2, None, 7)), TensorType(bool, shape=(100,))]
Inputs shapes: [(2, 7, 7), (100,)]
Inputs strides: [(392, 56, 8), (1,)]
Inputs values: ['not shown', 'not shown']
Outputs clients: [[multivariate_normal_rv{1, (1, 2), floatX, True}(RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FEA9C0FDC40>), [53], 11, [[0 0 0 0 ... 0 0 0 0]], AdvancedSubtensor.0)]]
Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
File \"/home/cao/miniconda3/envs/pymc/lib/python3.11/site-packages/ipykernel/zmqshell.py\", line 549, in run_cell
return super().run_cell(*args, **kwargs)
File \"/home/cao/miniconda3/envs/pymc/lib/python3.11/site-packages/IPython/core/interactiveshell.py\", line 3048, in run_cell
result = self._run_cell(
File \"/home/cao/miniconda3/envs/pymc/lib/python3.11/site-packages/IPython/core/interactiveshell.py\", line 3103, in _run_cell
result = runner(coro)
File \"/home/cao/miniconda3/envs/pymc/lib/python3.11/site-packages/IPython/core/async_helpers.py\", line 129, in _pseudo_sync_runner
coro.send(None)
File \"/home/cao/miniconda3/envs/pymc/lib/python3.11/site-packages/IPython/core/interactiveshell.py\", line 3308, in run_cell_async
has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
File \"/home/cao/miniconda3/envs/pymc/lib/python3.11/site-packages/IPython/core/interactiveshell.py\", line 3490, in run_ast_nodes
if await self.run_code(code, result, async_=asy):
File \"/home/cao/miniconda3/envs/pymc/lib/python3.11/site-packages/IPython/core/interactiveshell.py\", line 3550, in run_code
exec(code_obj, self.user_global_ns, self.user_ns)
File \"/tmp/ipykernel_8875/1979937421.py\", line 26, in <module>
likelihood = pm.MvNormal(\"likelihood\", mu=0, cov=cov[group==0], observed=obs[group==0])
HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.",
"stack": "---------------------------------------------------------------------------
IndexError Traceback (most recent call last)
File ~/miniconda3/envs/pymc/lib/python3.11/site-packages/pytensor/compile/function/types.py:970, in Function.__call__(self, *args, **kwargs)
968 try:
969 outputs = (
--> 970 self.vm()
971 if output_subset is None
972 else self.vm(output_subset=output_subset)
973 )
974 except Exception:
File ~/miniconda3/envs/pymc/lib/python3.11/site-packages/pytensor/graph/op.py:518, in Op.make_py_thunk.<locals>.rval(p, i, o, n)
516 @is_thunk_type
517 def rval(p=p, i=node_input_storage, o=node_output_storage, n=node):
--> 518 r = p(n, [x[0] for x in i], o)
519 for o in node.outputs:
File ~/miniconda3/envs/pymc/lib/python3.11/site-packages/pytensor/tensor/subtensor.py:2678, in AdvancedSubtensor.perform(self, node, inputs, out_)
2677 (out,) = out_
-> 2678 check_advanced_indexing_dimensions(inputs[0], inputs[1:])
2679 rval = inputs[0].__getitem__(tuple(inputs[1:]))
File ~/miniconda3/envs/pymc/lib/python3.11/site-packages/pytensor/tensor/subtensor.py:2580, in check_advanced_indexing_dimensions(input, idx_list)
2579 if index.shape[i] != input.shape[dim_seen + i]:
-> 2580 raise IndexError(
2581 \"boolean index did not match indexed array \"
2582 f\"along dimension {int(dim_seen + i)}; dimension is \"
2583 f\"{int(input.shape[dim_seen + i])} but \"
2584 f\"corresponding boolean dimension is {index.shape[i]}\"
2585 )
2586 dim_seen += index.ndim
IndexError: boolean index did not match indexed array along dimension 0; dimension is 2 but corresponding boolean dimension is 100
During handling of the above exception, another exception occurred:
IndexError Traceback (most recent call last)
/home/cao/projects/pymc_experiments/test.ipynb Cell 7 line 2
<a href='vscode-notebook-cell://wsl%2Bubuntu-20.04/home/cao/projects/pymc_experiments/test.ipynb#Y101sdnNjb2RlLXJlbW90ZQ%3D%3D?line=25'>26</a> likelihood = pm.MvNormal(\"likelihood\", mu=0, cov=cov[group==0], observed=obs[group==0])
<a href='vscode-notebook-cell://wsl%2Bubuntu-20.04/home/cao/projects/pymc_experiments/test.ipynb#Y101sdnNjb2RlLXJlbW90ZQ%3D%3D?line=27'>28</a> with m:
---> <a href='vscode-notebook-cell://wsl%2Bubuntu-20.04/home/cao/projects/pymc_experiments/test.ipynb#Y101sdnNjb2RlLXJlbW90ZQ%3D%3D?line=28'>29</a> idata_prior = pm.sample_prior_predictive()
<a href='vscode-notebook-cell://wsl%2Bubuntu-20.04/home/cao/projects/pymc_experiments/test.ipynb#Y101sdnNjb2RlLXJlbW90ZQ%3D%3D?line=29'>30</a> # idata = pm.sample(tune=1000, draws=1000, nuts_sampler=\"numpyro\")
File ~/miniconda3/envs/pymc/lib/python3.11/site-packages/pymc/sampling/forward.py:425, in sample_prior_predictive(samples, model, var_names, random_seed, return_inferencedata, idata_kwargs, compile_kwargs)
423 # All model variables have a name, but mypy does not know this
424 _log.info(f\"Sampling: {list(sorted(volatile_basic_rvs, key=lambda var: var.name))}\") # type: ignore
--> 425 values = zip(*(sampler_fn() for i in range(samples)))
427 data = {k: np.stack(v) for k, v in zip(names, values)}
428 if data is None:
File ~/miniconda3/envs/pymc/lib/python3.11/site-packages/pymc/sampling/forward.py:425, in <genexpr>(.0)
423 # All model variables have a name, but mypy does not know this
424 _log.info(f\"Sampling: {list(sorted(volatile_basic_rvs, key=lambda var: var.name))}\") # type: ignore
--> 425 values = zip(*(sampler_fn() for i in range(samples)))
427 data = {k: np.stack(v) for k, v in zip(names, values)}
428 if data is None:
File ~/miniconda3/envs/pymc/lib/python3.11/site-packages/pytensor/compile/function/types.py:983, in Function.__call__(self, *args, **kwargs)
981 if hasattr(self.vm, \"thunks\"):
982 thunk = self.vm.thunks[self.vm.position_of_error]
--> 983 raise_with_op(
984 self.maker.fgraph,
985 node=self.vm.nodes[self.vm.position_of_error],
986 thunk=thunk,
987 storage_map=getattr(self.vm, \"storage_map\", None),
988 )
989 else:
990 # old-style linkers raise their own exceptions
991 raise
File ~/miniconda3/envs/pymc/lib/python3.11/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 ~/miniconda3/envs/pymc/lib/python3.11/site-packages/pytensor/compile/function/types.py:970, in Function.__call__(self, *args, **kwargs)
967 t0_fn = time.perf_counter()
968 try:
969 outputs = (
--> 970 self.vm()
971 if output_subset is None
972 else self.vm(output_subset=output_subset)
973 )
974 except Exception:
975 restore_defaults()
File ~/miniconda3/envs/pymc/lib/python3.11/site-packages/pytensor/graph/op.py:518, in Op.make_py_thunk.<locals>.rval(p, i, o, n)
516 @is_thunk_type
517 def rval(p=p, i=node_input_storage, o=node_output_storage, n=node):
--> 518 r = p(n, [x[0] for x in i], o)
519 for o in node.outputs:
520 compute_map[o][0] = True
File ~/miniconda3/envs/pymc/lib/python3.11/site-packages/pytensor/tensor/subtensor.py:2678, in AdvancedSubtensor.perform(self, node, inputs, out_)
2676 def perform(self, node, inputs, out_):
2677 (out,) = out_
-> 2678 check_advanced_indexing_dimensions(inputs[0], inputs[1:])
2679 rval = inputs[0].__getitem__(tuple(inputs[1:]))
2680 # When there are no arrays, we are not actually doing advanced
2681 # indexing, so __getitem__ will not return a copy.
2682 # Since no view_map is set, we need to copy the returned value
File ~/miniconda3/envs/pymc/lib/python3.11/site-packages/pytensor/tensor/subtensor.py:2580, in check_advanced_indexing_dimensions(input, idx_list)
2578 for i in range(index.ndim):
2579 if index.shape[i] != input.shape[dim_seen + i]:
-> 2580 raise IndexError(
2581 \"boolean index did not match indexed array \"
2582 f\"along dimension {int(dim_seen + i)}; dimension is \"
2583 f\"{int(input.shape[dim_seen + i])} but \"
2584 f\"corresponding boolean dimension is {index.shape[i]}\"
2585 )
2586 dim_seen += index.ndim
2587 else:
IndexError: boolean index did not match indexed array along dimension 0; dimension is 2 but corresponding boolean dimension is 100
Apply node that caused the error: AdvancedSubtensor(cov, [ True Fal ... lse True])
Toposort index: 22
Inputs types: [TensorType(float64, shape=(2, None, 7)), TensorType(bool, shape=(100,))]
Inputs shapes: [(2, 7, 7), (100,)]
Inputs strides: [(392, 56, 8), (1,)]
Inputs values: ['not shown', 'not shown']
Outputs clients: [[multivariate_normal_rv{1, (1, 2), floatX, True}(RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FEA9C0FDC40>), [53], 11, [[0 0 0 0 ... 0 0 0 0]], AdvancedSubtensor.0)]]
Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
File \"/home/cao/miniconda3/envs/pymc/lib/python3.11/site-packages/ipykernel/zmqshell.py\", line 549, in run_cell
return super().run_cell(*args, **kwargs)
File \"/home/cao/miniconda3/envs/pymc/lib/python3.11/site-packages/IPython/core/interactiveshell.py\", line 3048, in run_cell
result = self._run_cell(
File \"/home/cao/miniconda3/envs/pymc/lib/python3.11/site-packages/IPython/core/interactiveshell.py\", line 3103, in _run_cell
result = runner(coro)
File \"/home/cao/miniconda3/envs/pymc/lib/python3.11/site-packages/IPython/core/async_helpers.py\", line 129, in _pseudo_sync_runner
coro.send(None)
File \"/home/cao/miniconda3/envs/pymc/lib/python3.11/site-packages/IPython/core/interactiveshell.py\", line 3308, in run_cell_async
has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
File \"/home/cao/miniconda3/envs/pymc/lib/python3.11/site-packages/IPython/core/interactiveshell.py\", line 3490, in run_ast_nodes
if await self.run_code(code, result, async_=asy):
File \"/home/cao/miniconda3/envs/pymc/lib/python3.11/site-packages/IPython/core/interactiveshell.py\", line 3550, in run_code
exec(code_obj, self.user_global_ns, self.user_ns)
File \"/tmp/ipykernel_8875/1979937421.py\", line 26, in <module>
likelihood = pm.MvNormal(\"likelihood\", mu=0, cov=cov[group==0], observed=obs[group==0])
HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node."
}