oryx.experimental.nn.Conv

Neural network layer for 2D convolution.

Inherits From: Layer, Module, Pytree

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

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
initialize(
    key,
    in_spec,
    out_chan,
    filter_shape,
    strides=None,
    padding='VALID',
    kernel_init=None,
    bias_init=stax.randn(1e-06),
    use_bias=True
)

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.

spec

View source

@classmethod
spec(
    in_spec,
    out_chan,
    filter_shape,
    strides=None,
    padding='VALID',
    kernel_init=None,
    bias_init=stax.randn(1e-06),
    use_bias=True
)

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.