import numpy as np from numpy.typing import ArrayLike from typing import Any, Tuple class Transform: def apply(self, t: float, x: np.array): y = self(t, x) assert x.shape == y.shape, "Transform has to maintain the array's shape." return y