oryx.experimental.nn.Layer

Base class for neural network layers.

Inherits From: Module, Pytree

View aliases

Main aliases

oryx.experimental.nn.base.Layer

oryx.experimental.nn.Layer(
    layer_params, name=None
)

A Layer is a subclass of Module with some additional functionality. Like Modules, Layers have a variables() method that returns a dictionary mapping names to state values. It also has a call_and_update function that returns the output of a computation and a new Layer with updated state. Underneath the hood, Layers do a couple extra things beyond Modules.

Attributes
info	Returns the info for this Layer.
params	Returns the parameters of this Layer.
state	Returns the state of this Layer.

Methods

call

View source

call(
    *args, **kwargs
)

Calls the Layer's call_and_update and returns the first result.

call_and_update

View source

call_and_update(
    *args, rng=None, **kwargs
)

Uses the layer_cau primitive to call `self._call_and_update.

flatten

View source

flatten()

Converts the Layer to a tuple suitable for PyTree.

initialize

View source

@classmethod
@abc.abstractmethod
initialize(
    init_key, in_spec
)

Initializes a Layer from an init_key and input specification.

new

View source

@classmethod
new(
    layer_params, name=None
)

Creates Layer given a LayerParams namedtuple.

Args
layer_params	LayerParams namedtuple that defines the Layer.
name	a string name for the Layer.

Returns
A Layer object.

replace

View source

replace(
    params=None, state=None, info=None
)

Returns a copy of the layer with replaced properties.

unflatten

View source

@classmethod
unflatten(
    data, xs
)

Reconstruct the Layer from a flattened version.

update

View source

update(
    *args, **kwargs
)

Calls the Layer's call_and_update and returns the second result.

variables

View source

variables()

Returns the variables dictionary for this Layer.

__call__

View source

__call__(
    *args, **kwargs
) -> Any

Emulates a regular function call.

A Module's dunder call will ensure state is updated after the function call by calling assign on the updated state before returning the output of the function.

Args
*args	The arguments to the module.
**kwargs	The keyword arguments to the module.

Returns
The output of the module.