pymc-devs / pymc-devs/nutpie

Commas in variable names cause idata writing to fail if combined with transforms

Open
#284 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
207
Forks
28
Avg merge
7d 21h
Merged PRs (30d)
1

Description

Reproducible example

import numpy as np
import pymc as pm

coords = {
    "age_group": [1, 2],
    "education": ["basic", "secondary", "higher"],
    "language_outp": ["et", "other"],
}

with pm.Model(coords=coords):
    x = pm.Normal(
        "alpha_age_group,education_offset_language",
        mu=0.0,
        sigma=1.0,
        dims=("age_group", "education", "language_outp"),
        transform=pm.distributions.transforms.ZeroSumTransform([0, 1]),
    )
    pm.Normal("y", mu=x[0, 0, 0], sigma=1.0, observed=np.array(0.0))

    pm.sample(
        draws=20,
        tune=20,
        chains=1,
        cores=1,
        progressbar=False,
        nuts_sampler="nutpie",
    )

fails with

Traceback (most recent call last):
  File "/home/velochy/salk/salk_internal_package/one_offs/replicate/nutpie_comma_dim_bug_repro.py", line 22, in <module>
    pm.sample(
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/site-packages/pymc/sampling/mcmc.py", line 832, in sample
    return _sample_external_nuts(
           ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/site-packages/pymc/sampling/mcmc.py", line 377, in _sample_external_nuts
    idata = nutpie.sample(
            ^^^^^^^^^^^^^^
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/site-packages/nutpie/sample.py", line 865, in sample
    result = sampler.wait()
             ^^^^^^^^^^^^^^
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/site-packages/nutpie/sample.py", line 528, in wait
    return self._extract(results)
           ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/site-packages/nutpie/sample.py", line 573, in _extract
    return _arrow_to_arviz(
           ^^^^^^^^^^^^^^^^
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/site-packages/nutpie/sample.py", line 99, in _arrow_to_arviz
    return arviz.from_dict(
           ^^^^^^^^^^^^^^^^
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/site-packages/arviz/data/io_dict.py", line 459, in from_dict
    ).to_inference_data()
      ^^^^^^^^^^^^^^^^^^^
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/site-packages/arviz/data/io_dict.py", line 334, in to_inference_data
    "posterior": self.posterior_to_xarray(),
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/site-packages/arviz/data/base.py", line 71, in wrapped
    return func(cls)
           ^^^^^^^^^
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/site-packages/arviz/data/io_dict.py", line 97, in posterior_to_xarray
    dict_to_dataset(
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/site-packages/arviz/data/base.py", line 406, in dict_to_dataset
    key: numpy_to_data_array(
         ^^^^^^^^^^^^^^^^^^^^
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/site-packages/arviz/data/base.py", line 305, in numpy_to_data_array
    coords = {key: xr.IndexVariable((key,), data=np.asarray(coords[key])) for key in dims}
                                                            ~~~~~~^^^^^
KeyError: 'education_offset_language_zerosum___dim_1'

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by running the reproducible Python example, then inspect nutpie/sample.py at _arrow_to_arviz and the ArviZ conversion shown in the traceback. The fix is complete when the example with the comma-containing variable name and ZeroSumTransform writes inference data without the missing-coordinate KeyError.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
data
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.