class liesel.goose.chain.Chain(*args, **kwargs)[source]#

Bases: Protocol[TPyTree]

A Chain stores multiple chucks of pytrees and concatenates them along a time axis.

The Chain always assume multiple independent chains that are indexed via the first axis. The second dimension represents the time. Consequently, the leaves in the pytree must have a dimension of two (i.e., [chain, time, …]).

A Chain operates on the assumption that all chunks are pytrees with the same structure. However, the time dimension is allowed to vary in size.



Appends a chunk to the chain.


Returns all chunks combined into one pytree.