Function that returns activation function, as well as its derivative

Question

I have the following function that returns one of several functions along with their derivative for training neural networks.
Can this function be improved and made more pythonic in any way?
Are there other things like reducing the redundancy of the list, etc.?
Could a better way be to instead enclose the inner functions inside a class?

Note: I am not looking to optimize the computational aspects of the functions.

import numpy as np


def get_activation_fn_with_deriv(fn_name):
    """ Returns function objects for activation function and its derivative.

    Returns function objects for the specified activation function and its derivative function.

    Args:
        fn_name: Name of the activation function.
    
    Returns:
        (fn, fn_deriv): A tuple of the activation function and its derivative function.
    
    Raises:
        ValueError: If an unavailable activation function is retrieved.
    
    Examples:
        >>> relu, relu_deriv = activation_fn_with_deriv('relu')
        >>> relu(np.array([1, -1, 2.3]))
        array([1. , 0. , 2.3])
    """

    available_fn = ['identity', 'relu', 'sigmoid', 'tanh']

    match fn_name.lower():
        case 'identity':
            def identity(x):
                return x
            
            def identity_deriv(x):
                return np.ones(x.shape)
            
            return identity, identity_deriv
        
        case 'relu':
            def relu(x):
                return np.maximum(x, 0.0)

            def relu_deriv(x):
                return np.where(x >= 0.0, 1.0, 0.0)

            return relu, relu_deriv
        
        case 'sigmoid':
            def sigmoid(x):
                return 1 / (1 + np.exp(-x))

            def sigmoid_deriv(x):
                sigmoid_x = sigmoid(x)
                return sigmoid_x * (1 - sigmoid_x)
            
            return sigmoid, sigmoid_deriv
        
        case 'tanh':
            def tanh(x):
                exp_2x = np.exp(2 * x)
                return (exp_2x - 1) / (exp_2x + 1)

            def tanh_deriv(x):
                return 1 - tanh(x) ** 2
            
            return tanh, tanh_deriv

        case _:
            raise ValueError(f'Specified activation function not available. Choose from - {[name.lower() for name in available_fn]}.')

Answer 1

I like it, this code is beautiful, and it comes with a nice doctest. Ship it as-is. Of course, there's always more to say about any given piece of code.

(Oh, my! $ python -m doctest *.py fails. Having written one, it's important to actually run the doctest.)

namedtuple

Returning a (fn, deriv) tuple is just Fine. But consider naming the [0] and [1] elements.

inner functions

While they have their uses, generally I'm not keen on inner functions, due to coupling of namespace and difficulty of calling them. Here you're actually returning the function, mitigating that second concern.

But suppose a maintenance engineer was writing a test suite that called them. It would still be more convenient to see them up at module level.

Let's assume we do that, creating a new module devoted just to activation functions and their derivatives.

DRY

The parallel structure of foo() and foo_deriv() is perfect; we should keep that bit of repetition. But then we repeat foo in a list, and again in a match clause.

    available_fn = ['identity', 'relu', 'sigmoid', 'tanh']

When raising a ValueError we could build this list by scanning the activation module. (Throw fatal error if we find a foo which lacks foo_deriv.)

too permissive

    match fn_name.lower():

Delete the .lower() call, please. Your heart was in the right place, but I feel this encourages callers to be too creative. Better to offer a narrow Public API which you can choose to widen later if there really is a use case for supporting variant spellings.

With functions accessible at module level, consider requiring callers to pass in the function of interest, and now your responsibility is simply to getattr() the corresponding derivative function from that module.

Or we might choose to populate a dict that maps from fn to deriv.

auto differentiation

Demanding domain knowledge, and an implementation of the derivative, is entirely reasonable and produces fast code.

However, you might consider supporting an activation function foo without an accompanying foo_deriv. Sympy in many cases can find a symbolic derivative, though that involves a different representation of the function. Or one could use numeric differencing methods to turn a vector of values into the corresponding slopes.

Answer 2

Accepting a string fn_name is poorly typed ("stringly typed"). From least to most preferable solutions:

• Type-hint the function parameter as a Literal; or
• Accept an Enum instead of a string; or
• Don't do any of this if you can avoid it. Refer to separate implementations where they're needed.

Could a better way be to instead enclose the inner functions inside a class?

I believe so. Since your functions are all stateless (good), they don't need to be instance methods; you can write them as class methods implementing an abstract parent.

It's important that you hint your signatures. Numpy type-hinting is still in development and so is only half-working, but is better than nothing. The goal is to hint that your functions accept some array type (with the dtype unspecified), and return an array type with the same dtype that they accept. Mypy understands this to a limit.

With this code:

import typing
from abc import ABC, abstractmethod

import numpy as np


ActT = typing.TypeVar('ActT', bound=np.ndarray)


class Activation(ABC):
    @classmethod
    @abstractmethod
    def fn(cls, x: ActT) -> ActT: ...

    @classmethod
    @abstractmethod
    def fn_deriv(cls, x: ActT) -> ActT: ...


class Identity(Activation):
    @classmethod
    def fn(cls, x: ActT) -> ActT:
        return x

    @classmethod
    def fn_deriv(cls, x: ActT) -> ActT:
        return np.ones_like(x)


class Relu(Activation):
    @classmethod
    def fn(cls, x: ActT) -> ActT:
        return x.clip(a_min=0)

    @classmethod
    def fn_deriv(cls, x: ActT) -> ActT:
        return (x >= 0).astype(x.dtype)


class Sigmoid(Activation):
    @classmethod
    def fn(cls, x: ActT) -> ActT:
        return 1/(1 + np.exp(-x))

    @classmethod
    def fn_deriv(cls, x: ActT) -> ActT:
        sigmoid_x = cls.fn(x)
        return sigmoid_x*(1 - sigmoid_x)


class Tanh(Activation):
    @classmethod
    def fn(cls, x: ActT) -> ActT:
        exp_2x = np.exp(2*x)
        return (exp_2x - 1)/(exp_2x + 1)

    @classmethod
    def fn_deriv(cls, x: ActT) -> ActT:
        return 1 - cls.fn(x)**2

invocation looks like

print(
    Tanh.fn_deriv(np.array((1, 2, 3)))
)

Answer 3

You have nice consistency in how you declare your inner functions.

Here are some points to consider.

The function name in the docstring should match the name of the function (get_activation_fn_with_deriv). Change:

    >>> relu, relu_deriv = activation_fn_with_deriv('relu')

to:

    >>> relu, relu_deriv = get_activation_fn_with_deriv('relu')

Since all the function names in the available_fn list are already lower-case, there is no need to loop through them and call lower(). The raise statement can be simplified, also shortening a long line:

raise ValueError(f'Specified activation function not available. Choose from: {available_fn}')