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")