From c6e0f97a47692e460ec4503a389fe08631811f0c Mon Sep 17 00:00:00 2001
From: Michael Osthege <m.osthege@fz-juelich.de>
Date: Mon, 19 Dec 2022 17:11:21 +0100
Subject: [PATCH 1/3] Bump for 0.2.6

---
 mcbackend/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mcbackend/__init__.py b/mcbackend/__init__.py
index a044c75..74c0718 100644
--- a/mcbackend/__init__.py
+++ b/mcbackend/__init__.py
@@ -20,4 +20,4 @@
     pass
 
 
-__version__ = "0.2.5"
+__version__ = "0.2.6"

From 170797f1edcde48b13e76f762b7605b4c8c53a98 Mon Sep 17 00:00:00 2001
From: Michael Osthege <m.osthege@fz-juelich.de>
Date: Mon, 19 Dec 2022 18:46:45 +0100
Subject: [PATCH 2/3] Fix ragged-array conversion for compatibility with
 NumPy>=1.24.0

---
 mcbackend/core.py       | 11 ++++++++++-
 mcbackend/test_utils.py | 13 +++++++++++++
 mcbackend/utils.py      | 11 +++++++++++
 3 files changed, 34 insertions(+), 1 deletion(-)
 create mode 100644 mcbackend/utils.py

diff --git a/mcbackend/core.py b/mcbackend/core.py
index b962437..9f3631a 100644
--- a/mcbackend/core.py
+++ b/mcbackend/core.py
@@ -9,6 +9,7 @@
 
 from .meta import ChainMeta, RunMeta, Variable
 from .npproto.utils import ndarray_to_numpy
+from .utils import as_array_from_ragged
 
 InferenceData = TypeVar("InferenceData")
 try:
@@ -252,7 +253,15 @@ def to_inferencedata(self, *, equalize_chain_lengths: bool = True, **kwargs) ->
                 warmup_sample_stats[svar.name].append(stats[tune])
                 sample_stats[svar.name].append(stats[~tune])
 
-        kwargs.setdefault("save_warmup", True)
+        if not equalize_chain_lengths:
+            # Convert ragged arrays to object-dtyped ndarray because NumPy >=1.24.0 no longer does that automatically
+            warmup_posterior = {k: as_array_from_ragged(v) for k, v in warmup_posterior.items()}
+            warmup_sample_stats = {
+                k: as_array_from_ragged(v) for k, v in warmup_sample_stats.items()
+            }
+            posterior = {k: as_array_from_ragged(v) for k, v in posterior.items()}
+            sample_stats = {k: as_array_from_ragged(v) for k, v in sample_stats.items()}
+
         idata = from_dict(
             warmup_posterior=warmup_posterior,
             warmup_sample_stats=warmup_sample_stats,
diff --git a/mcbackend/test_utils.py b/mcbackend/test_utils.py
index 13a90b2..2704785 100644
--- a/mcbackend/test_utils.py
+++ b/mcbackend/test_utils.py
@@ -10,6 +10,7 @@
 import pytest
 
 import mcbackend
+from mcbackend import utils as mutils
 from mcbackend.meta import ChainMeta, DataVariable, RunMeta, Variable
 from mcbackend.npproto import utils
 
@@ -407,3 +408,15 @@ def test__big_variables(self):
         speed = self.measure_big_variables()
         assert speed.draws_per_second > 500 or speed.mib_per_second > 5
         pass
+
+
+def test_as_array_from_ragged():
+    even = mutils.as_array_from_ragged(
+        [
+            numpy.ones(2),
+            numpy.ones(3),
+        ]
+    )
+    assert isinstance(even, numpy.ndarray)
+    assert even.dtype == numpy.dtype(object)
+    pass
diff --git a/mcbackend/utils.py b/mcbackend/utils.py
new file mode 100644
index 0000000..5af8126
--- /dev/null
+++ b/mcbackend/utils.py
@@ -0,0 +1,11 @@
+"""Contains helper functions that are independent of McBackend components."""
+from typing import Sequence
+
+import numpy as np
+
+
+def as_array_from_ragged(arrs: Sequence[np.ndarray]) -> np.ndarray:
+    shapes = {np.shape(arr) for arr in arrs}
+    if len(shapes) > 1:
+        return np.array(arrs, dtype=object)
+    return np.array(arrs)

From a97d84d1c055fddb140af9b279e75085aca080da Mon Sep 17 00:00:00 2001
From: Michael Osthege <m.osthege@fz-juelich.de>
Date: Mon, 19 Dec 2022 18:54:14 +0100
Subject: [PATCH 3/3] adapt to new ArviZ defaults

---
 mcbackend/core.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mcbackend/core.py b/mcbackend/core.py
index 9f3631a..bb1b118 100644
--- a/mcbackend/core.py
+++ b/mcbackend/core.py
@@ -272,6 +272,7 @@ def to_inferencedata(self, *, equalize_chain_lengths: bool = True, **kwargs) ->
             attrs=self.meta.attributes,
             constant_data=self.constant_data,
             observed_data=self.observed_data,
+            save_warmup=True,
             **kwargs,
         )
         return idata