lakonlab/models/diffusions/piflow_policies/base.py

from abc import ABCMeta, abstractmethod


class BasePolicy(metaclass=ABCMeta):

    @abstractmethod
    def pi(self, x_t, sigma_t):
        """Compute the flow velocity at (x_t, t).

        Args:
            x_t (torch.Tensor): Noisy input at time t.
            sigma_t (torch.Tensor): Noise level at time t.

        Returns:
            torch.Tensor: The computed flow velocity u_t.
        """
        pass

    @abstractmethod
    def detach(self):
        pass