codecov[bot]: # [Codecov](https://codecov.io/gh/eigenfoo/littlemcmc/pull/63?src=pr&el=h1) Report
> Merging [#63](https://codecov.io/gh/eigenfoo/littlemcmc/pull/63?src=pr&el=desc) into [master](https://codecov.io/gh/eigenfoo/littlemcmc/commit/6ab62cc644efde9aec9258c7a5477c3ecdc16284&el=desc) will **decrease** coverage by `0.02%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/eigenfoo/littlemcmc/pull/63/graphs/tree.svg?width=650&height=150&src=pr&token=7TboOy71MZ)](https://codecov.io/gh/eigenfoo/littlemcmc/pull/63?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master      #63      +/-   ##
==========================================
- Coverage   78.68%   78.66%   -0.03%     
==========================================
  Files          10       10              
  Lines         807      806       -1     
==========================================
- Hits          635      634       -1     
  Misses        172      172              
```


| [Impacted Files](https://codecov.io/gh/eigenfoo/littlemcmc/pull/63?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [littlemcmc/sampling.py](https://codecov.io/gh/eigenfoo/littlemcmc/pull/63/diff?src=pr&el=tree#diff-bGl0dGxlbWNtYy9zYW1wbGluZy5weQ==) | `87.00% <100.00%> (-0.13%)` | :arrow_down: |

diff --git a/littlemcmc/sampling.py b/littlemcmc/sampling.py
index 50a3c96..74e22fd 100644
--- a/littlemcmc/sampling.py
+++ b/littlemcmc/sampling.py
@@ -210,13 +210,18 @@ def sample(
         for i, seed in enumerate(random_seed)
     )
 
-    # Flatten `trace` to have shape [num_variables, num_chains * num_samples]
-    trace = np.hstack([np.atleast_2d(chain_trace) for (chain_trace, _) in results])
+    # Reshape `trace` to have shape [num_chains, num_samples, num_variables]
+    trace = np.array([np.atleast_2d(chain_trace).T for (chain_trace, _) in results])
 
-    # Reshape `stats` to a dictionary
-    stats_ = [iter_stats for (_, chain_stats) in results for iter_stats in chain_stats]
+    # Reshape `stats` to a dictionary with keys = string of sampling stat name,
+    # values = np.array with shape [num_chains, num_samples, num_variables]
     stats = {
-        name: np.squeeze(np.array([iter_stats[name] for iter_stats in stats_])).astype(dtype)
+        name: np.array(
+            [
+                [np.atleast_1d(iter_stats[name]) for iter_stats in chain_stats]
+                for (_, chain_stats) in results
+            ]
+        ).astype(dtype)
         for (name, dtype) in step.stats_dtypes[0].items()
     }
 

Use a better shape for trace object
Currently we flatten the `trace` object to have shape `[num_variables, num_chains * num_samples]`. Really, it ought to be `[num_chains, num_samples, num_variables]` (or whatever reasonable choice ArviZ makes: we should target compatibility with them.

https://github.com/eigenfoo/littlemcmc/blob/master/littlemcmc/sampling.py#L213-L214
diff --git a/tests/test_sampling.py b/tests/test_sampling.py
index 19fca98..7b1d71e 100644
--- a/tests/test_sampling.py
+++ b/tests/test_sampling.py
@@ -31,27 +31,47 @@ def test_init_nuts(method):
 def test_hmc_sampling_runs():
     model_ndim = 1
     step = lmc.HamiltonianMC(logp_dlogp_func=logp_dlogp_func, model_ndim=model_ndim)
-    draws = 2
+    draws = 3
     tune = 1
-    chains = 1
+    chains = 2
     cores = 1
+
+    expected_shape = (2, 3, 1)
+
     trace, stats = lmc.sample(
         logp_dlogp_func, model_ndim, draws, tune, step=step, chains=chains, cores=cores
     )
-    assert trace.shape == (1, 2)
+    assert trace.shape == expected_shape
+    assert all([stats[name].shape == expected_shape for (name, _) in step.stats_dtypes[0].items()])
+    assert all(
+        [
+            stats[name].dtype == expected_dtype
+            for (name, expected_dtype) in step.stats_dtypes[0].items()
+        ]
+    )
 
 
 def test_nuts_sampling_runs():
     model_ndim = 1
     step = lmc.NUTS(logp_dlogp_func=logp_dlogp_func, model_ndim=model_ndim)
-    draws = 2
+    draws = 3
     tune = 1
-    chains = 1
+    chains = 2
     cores = 1
+
+    expected_shape = (2, 3, 1)
+
     trace, stats = lmc.sample(
         logp_dlogp_func, model_ndim, draws, tune, step=step, chains=chains, cores=cores
     )
-    assert trace.shape == (1, 2)
+    assert trace.shape == (2, 3, 1)
+    assert all([stats[name].shape == expected_shape for (name, _) in step.stats_dtypes[0].items()])
+    assert all(
+        [
+            stats[name].dtype == expected_dtype
+            for (name, expected_dtype) in step.stats_dtypes[0].items()
+        ]
+    )
 
 
 def test_multiprocess_sampling_runs():

pip install -e .[dev]
parse_log_pytest
requirements.txt
pytest pytest-cov numpy>=1.20.0
apt-get update apt-get install -y gcc
3.9
requirements/base.txt
pytest --no-header -rA --tb=line --color=no -p no:cacheprovider -W ignore::DeprecationWarning
coverage==7.8.0
exceptiongroup==1.2.2
iniconfig==2.1.0
joblib==1.4.2
-e git+https://github.com/eigenfoo/littlemcmc.git@6ab62cc644efde9aec9258c7a5477c3ecdc16284#egg=littlemcmc
numpy==2.0.2
packaging==24.2
pluggy==1.5.0
pytest==8.3.5
pytest-cov==6.0.0
scipy==1.13.1
tomli==2.2.1
tqdm==4.67.1

name: littlemcmc
channels:
  - defaults
  - https://repo.anaconda.com/pkgs/main
  - https://repo.anaconda.com/pkgs/r
  - conda-forge
dependencies:
  - _libgcc_mutex=0.1=main
  - _openmp_mutex=5.1=1_gnu
  - ca-certificates=2025.2.25=h06a4308_0
  - ld_impl_linux-64=2.40=h12ee557_0
  - libffi=3.4.4=h6a678d5_1
  - libgcc-ng=11.2.0=h1234567_1
  - libgomp=11.2.0=h1234567_1
  - libstdcxx-ng=11.2.0=h1234567_1
  - ncurses=6.4=h6a678d5_0
  - openssl=3.0.16=h5eee18b_0
  - pip=25.0=py39h06a4308_0
  - python=3.9.21=he870216_1
  - readline=8.2=h5eee18b_0
  - setuptools=75.8.0=py39h06a4308_0
  - sqlite=3.45.3=h5eee18b_0
  - tk=8.6.14=h39e8969_0
  - tzdata=2025a=h04d1e81_0
  - wheel=0.45.1=py39h06a4308_0
  - xz=5.6.4=h5eee18b_1
  - zlib=1.2.13=h5eee18b_1
  - pip:
      - coverage==7.8.0
      - exceptiongroup==1.2.2
      - iniconfig==2.1.0
      - joblib==1.4.2
      - littlemcmc==0.1.0
      - numpy==2.0.2
      - packaging==24.2
      - pluggy==1.5.0
      - pytest==8.3.5
      - pytest-cov==6.0.0
      - scipy==1.13.1
      - tomli==2.2.1
      - tqdm==4.67.1
prefix: /opt/conda/envs/littlemcmc

tests/test_sampling.py::test_hmc_sampling_runs tests/test_sampling.py::test_nuts_sampling_runs tests/test_sampling.py::test_hmc_recovers_1d_normal tests/test_sampling.py::test_nuts_recovers_1d_normal
tests/test_sampling.py::test_init_nuts[adapt_diag] tests/test_sampling.py::test_init_nuts[jitter+adapt_diag] tests/test_sampling.py::test_init_nuts[adapt_full] tests/test_sampling.py::test_init_nuts[jitter+adapt_full] tests/test_sampling.py::test_multiprocess_sampling_runs
littlemcmc/sampling.py
littlemcmc/sampling.py