Fix DenseInfo type (#400)

patrick-kidger · Apr 18, 2024 · eff5f60 · eff5f60
1 parent 38a40e7
commit eff5f60
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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")