subset_draws(x, chain = j) has incorrect behaviour following split_chains for type draws_rvars

[jatotterdell]

Describe the bug
subset_draws detects incorrect number of chains/iterations for draws_rvars following a call of split_chains.

To Reproduce

library(cmdstanr)
library(posterior)

file <- file.path(cmdstan_path(), "examples/bernoulli/bernoulli.stan")
mod <- cmdstan_model(file)
stan_data <- list(N = 10, y = c(0, 1, 0, 0, 0, 0, 0, 0, 0, 1))
fit <- mod$sample(data = stan_data)
rvs <- as_draws_rvars(fit$draws())

c("Chains" = nchains(rvs), "Iter" = niterations(rvs))
sapply(
  seq_len(nchains(rvs)),
  \(x) {
    tmp <- subset_draws(rvs, chain = x)
    c("Chains" = nchains(tmp), "Iter" = niterations(tmp))
  }
)

split_rvs <- split_chains(rvs)
c("Chains" = nchains(split_rvs), "Iter" = niterations(split_rvs))
split_rvs_1 <- subset_draws(split_rvs, chain = 1)
c("Chains" = nchains(split_rvs_1), "Iter" = niterations(split_rvs_1))
subset_draws(split_rvs, chain = 2)

Results in

Chains   Iter 
     4   1000 
       [,1] [,2] [,3] [,4]
Chains    1    1    1    1
Iter   1000 1000 1000 1000
Chains   Iter 
     8    500 
Chains   Iter 
     1   4000 
Error: Tried to subset chains up to '2' but the object only has '1' chains.

Expected behavior
Expect subset_draws(split_chains(x), chain = 1) to return draws for selected chain rather than all draws, and for subset_draws(split_chains(x), chain = 2) to return draws for chain 2 rather than fail.

Operating system

r$> sessionInfo()
R version 4.3.1 (2023-06-16)
Platform: x86_64-pc-linux-gnu (64-bit)
Running under: Ubuntu 22.04.3 LTS

posterior version number
1.4.1

[paul-buerkner]

Thanks. I will take a look. Currently occupied with deadlines but will get back to it soon.

[mjskay]

Hmm I think I wrote this before draws_of got a with_chains option, could probably be simplified to share some similarities (and maybe code) with the draws_array implementation by using with_chains = TRUE.

[paul-buerkner]

@mjskay I would like to do a new posterior release soon (end of this week perhaps). @mjskay how complicated would it be for you to fix this? I am not sure I am confident enough to touch the rvar interface unless I have to :-D

[mjskay]

I can try to take a look this weekend --- currently at a conference but there's a long plane ride back 🤷‍♂️

[paul-buerkner]

Thank you!