pymc icon indicating copy to clipboard operation
pymc copied to clipboard

Consider having coords per variable

Open ricardoV94 opened this issue 5 months ago • 1 comments

Description

This is currently not possible to do correctly:

import pymc as pm

with pm.Model(coords={"time": range(4)}) as m:
  x = pm.Normal("x", dims=("time",))
  pm.Deterministic("x[1:]", x[1:], dims=("time",))
  idata = pm.sample()  # fails

Users have two options: Use a different dim with custom cords for x[1:] or don't specify dims for it, in which case it get's a different dim and default coords after sampling. Either way no option to retain the dim.

I suggest allowing coords to pm.Deterministic and keep track of coords per variable at the model level.

with pm.Model(coords={"time": range(4)}) as m:
  x = pm.Normal("x", dims=("time",))
  pm.Deterministic("x[1:]", x[1:], dims=("time",), coords={"time": range(1, 4)})
  idata = pm.sample()  # correct

With the awareness that coords don't do anything at the model level, only after exporting them

ricardoV94 avatar Jul 11 '25 13:07 ricardoV94