Skip to content

treex.Lambda

A Module that applies a pure function to its input.

Source code in treex/nn/sequential.py
class Lambda(Module):
    """
    A Module that applies a pure function to its input.
    """

    f: tp.Callable[[jnp.ndarray], jnp.ndarray]

    def __init__(self, f: tp.Callable[[jnp.ndarray], jnp.ndarray]):
        """
        Arguments:
            f: A function to apply to the input.
        """

        self.f = f

    def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
        """
        Arguments:
            x: The input to the function.
        Returns:
            The output of the function.
        """
        return self.f(x)

__call__(self, x) special

Parameters:

Name Type Description Default
x ndarray

The input to the function.

required

Returns:

Type Description
ndarray

The output of the function.

Source code in treex/nn/sequential.py
def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
    """
    Arguments:
        x: The input to the function.
    Returns:
        The output of the function.
    """
    return self.f(x)

__init__(self, f) special

Parameters:

Name Type Description Default
f Callable[[jax._src.numpy.lax_numpy.ndarray], jax._src.numpy.lax_numpy.ndarray]

A function to apply to the input.

required
Source code in treex/nn/sequential.py
def __init__(self, f: tp.Callable[[jnp.ndarray], jnp.ndarray]):
    """
    Arguments:
        f: A function to apply to the input.
    """

    self.f = f