From eff5f60f6aab3c71a724defe35877a474becb679 Mon Sep 17 00:00:00 2001 From: Evgenii Zheltonozhskii Date: Thu, 18 Apr 2024 22:45:01 +0300 Subject: [PATCH] Fix DenseInfo type (#400) --- 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")