neural_dynamics.core

Core modules and typing helpers for neural_dynamics.

class TrajectoryDataset(typing.Protocol):

Protocol describing a time-indexed batch of trajectories.

Implementations must expose a monotonic time_grid tensor and a trajectories tensor where the first axis indexes the batch of trajectories and the second axis aligns with time_grid.

TrajectoryDataset(*args, **kwargs)
time_grid: torch.Tensor
trajectories: torch.Tensor