SPATZ/spatz/transforms/transform.py

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