function.py

class Function:
    """
    Abstract model of a differentiable function.
    """
    def forward(self, *args, **kwargs):
        """
        Forward pass of the function. Calculates the output value and the
        gradient at the input as well.
        """
        pass

    def local_grad(self, *args, **kwargs):
        """
        Calculates the local gradients of the function at the given input.
        Returns:
            grad: dictionary of local gradients.
        """
        pass