View source on GitHub |
Flattens the inputs collapsing all ending dimensions.
Inherits From: Layer
, Module
, Pytree
oryx.experimental.nn.Flatten(
layer_params, name=None
)
Attributes | |
---|---|
info
|
Returns the info for this Layer .
|
params
|
Returns the parameters of this Layer .
|
state
|
Returns the state of this Layer .
|
Methods
call
call(
*args, **kwargs
)
Calls the Layer
's call_and_update
and returns the first result.
call_and_update
call_and_update(
*args, rng=None, **kwargs
)
Uses the layer_cau
primitive to call `self._call_and_update.
flatten
flatten()
Converts the Layer to a tuple suitable for PyTree.
initialize
@classmethod
initialize( rng, in_spec )
Initializes Flatten Layer.
Args | |
---|---|
rng
|
Random key. |
in_spec
|
Input Spec. |
Returns | |
---|---|
Tuple with the output shape and the LayerParams. |
new
@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
replace(
params=None, state=None, info=None
)
Returns a copy of the layer with replaced properties.
spec
@classmethod
spec( in_spec )
unflatten
@classmethod
unflatten( data, xs )
Reconstruct the Layer from a flattened version.
update
update(
*args, **kwargs
)
Calls the Layer
's call_and_update
and returns the second result.
variables
variables()
Returns the variables dictionary for this Layer
.
__call__
__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. |