From 4a94bd44cd96258cd7ac23b69282f4ab7f3c91ad Mon Sep 17 00:00:00 2001 From: Evgenii Zheltonozhskii Date: Tue, 16 Apr 2024 09:45:52 +0300 Subject: [PATCH] Fix DenseInfo type --- diffrax/_custom_types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diffrax/_custom_types.py b/diffrax/_custom_types.py index ffccf64a..1ebbbb01 100644 --- a/diffrax/_custom_types.py +++ b/diffrax/_custom_types.py @@ -49,7 +49,7 @@ Args = PyTree[Any] DenseInfo = dict[str, PyTree[Array]] -DenseInfos = dict[str, PyTree[Shaped[Array, "times ..."]]] +DenseInfos = dict[str, PyTree[Shaped[Array, "times-1 ..."]]] BufferDenseInfos = dict[str, PyTree[eqxi.MaybeBuffer[Shaped[Array, "times ..."]]]] sentinel: Any = eqxi.doc_repr(object(), "sentinel")