LuxCore
LuxCore.jl
defines the abstract layers for Lux. Allows users to be compatible with the entirely of Lux.jl
without having such a heavy dependency. If you are depending on Lux.jl
directly, you do not need to depend on LuxCore.jl
(all the functionality is exported via Lux.jl
).
Index
Abstract Types
# LuxCore.AbstractExplicitLayer —
Type.
```julia
abstract type AbstractExplicitLayer
```
Abstract Type for all Lux Layers
Users implementing their custom layer, **must** implement
* `initialparameters(rng::AbstractRNG, layer::CustomAbstractExplicitLayer)` – This returns a `NamedTuple` containing the trainable parameters for the layer.
* `initialstates(rng::AbstractRNG, layer::CustomAbstractExplicitLayer)` – This returns a NamedTuple containing the current state for the layer. For most layers this is typically empty. Layers that would potentially contain this include `BatchNorm`, `LSTM`, `GRU` etc.
Optionally:
* `parameterlength(layer::CustomAbstractExplicitLayer)` – These can be automatically calculated, but it is recommended that the user defines these.
* `statelength(layer::CustomAbstractExplicitLayer)` – These can be automatically calculated, but it is recommended that the user defines these.
See also [`AbstractExplicitContainerLayer`](LuxCore#LuxCore.AbstractExplicitContainerLayer)
# LuxCore.AbstractExplicitContainerLayer —
Type.
```julia
abstract type AbstractExplicitContainerLayer{layers} <: LuxCore.AbstractExplicitLayer
```
Abstract Container Type for certain Lux Layers. `layers` is a tuple containing fieldnames for the layer, and constructs the parameters and states using those.
Users implementing their custom layer can extend the same functions as in [`AbstractExplicitLayer`](LuxCore#LuxCore.AbstractExplicitLayer).
::: tip
Advanced structure manipulation of these layers post construction is possible via `Functors.fmap`. For a more flexible interface, we recommend using the experimental feature [`Lux.Experimental.@layer_map`](@ref).
:::
General
# LuxCore.apply —
Function.
```julia
apply(
model::LuxCore.AbstractExplicitLayer,
x,
ps,
st::NamedTuple
) -> Any
```
Simply calls `model(x, ps, st)`
# LuxCore.display_name —
Function.
```julia
display_name(layer::AbstractExplicitLayer)
```
Printed Name of the `layer`. If the `layer` has a field `name` that is used, else the type name is used.
# LuxCore.setup —
Function.
```julia
setup(
rng::Random.AbstractRNG,
l::LuxCore.AbstractExplicitLayer
) -> Tuple{Any, Any}
```
Shorthand for getting the parameters and states of the layer `l`. Is equivalent to `(initialparameters(rng, l), initialstates(rng, l))`.
::: warning
This function is not pure, it mutates `rng`.
:::
Parameters
# LuxCore.initialparameters —
Function.
```julia
initialparameters(
_::Random.AbstractRNG,
_::LuxCore.AbstractExplicitLayer
) -> Union{NamedTuple{(), Tuple{}}, NamedTuple{(:bias, :scale), _A} where _A<:Tuple{Any, Any}}
```
Generate the initial parameters of the layer `l`.
# LuxCore.parameterlength —
Function.
```julia
parameterlength(l::LuxCore.AbstractExplicitLayer) -> Any
```
Return the total number of parameters of the layer `l`.
States
# LuxCore.initialstates —
Function.
```julia
initialstates(
_::Random.AbstractRNG,
_::LuxCore.AbstractExplicitLayer
) -> NamedTuple{(:cell, :carry), _A} where _A<:Tuple{Any, Nothing}
```
Generate the initial states of the layer `l`.
# LuxCore.statelength —
Function.
```julia
statelength(l::LuxCore.AbstractExplicitLayer) -> Any
```
Return the total number of states of the layer `l`.
# LuxCore.testmode —
Function.
```julia
testmode(st::NamedTuple) -> Any
```
Make all occurances of `training` in state `st` – `Val(false)`.
# LuxCore.trainmode —
Function.
```julia
trainmode(st::NamedTuple) -> Any
```
Make all occurances of `training` in state `st` – `Val(true)`.
# LuxCore.update_state —
Function.
```julia
update_state(
st::NamedTuple,
key::Symbol,
value;
layer_check
) -> Any
```
Recursively update all occurances of the `key` in the state `st` with the `value`.