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} <: 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 `Lux.Experimental.@layer_map`.
General
# LuxCore.apply —
Function.
```julia
apply(model, x, ps, st)
```
Simply calls `model(x, ps, st)`
# LuxCore.stateless_apply —
Function.
```julia
stateless_apply(model, x, ps)
```
Calls `apply` and only returns the first argument.
# LuxCore.check_fmap_condition —
Function.
```julia
check_fmap_condition(cond, tmatch, x) -> Bool
```
`fmap`s into the structure `x` and see if `cond` is statisfied for any of the leaf elements.
**Arguments**
* `cond` - A function that takes a single argument and returns a `Bool`.
* `tmatch` - A shortcut to check if `x` is of type `tmatch`. Can be disabled by passing `nothing`.
* `x` - The structure to check.
**Returns**
A Boolean Value
# LuxCore.contains_lux_layer —
Function.
```julia
contains_lux_layer(l) -> Bool
```
Check if the structure `l` is a Lux AbstractExplicitLayer or a container of such a layer.
# 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.replicate —
Function.
```julia
replicate(rng::AbstractRNG)
```
Creates a copy of the `rng` state depending on its type.
# LuxCore.setup —
Function.
```julia
setup(rng::AbstractRNG, layer)
```
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(rng::AbstractRNG, layer)
```
Generate the initial parameters of the layer `l`.
# LuxCore.parameterlength —
Function.
```julia
parameterlength(layer)
```
Return the total number of parameters of the layer `l`.
States
# LuxCore.initialstates —
Function.
```julia
initialstates(rng::AbstractRNG, layer)
```
Generate the initial states of the layer `l`.
# LuxCore.statelength —
Function.
```julia
statelength(layer)
```
Return the total number of states of the layer `l`.
# LuxCore.testmode —
Function.
```julia
testmode(st::NamedTuple)
```
Make all occurances of `training` in state `st` – `Val(false)`.
# LuxCore.trainmode —
Function.
```julia
trainmode(st::NamedTuple)
```
Make all occurances of `training` in state `st` – `Val(true)`.
# LuxCore.update_state —
Function.
```julia
update_state(st::NamedTuple, key::Symbol, value;
layer_check=_default_layer_check(key))
```
Recursively update all occurances of the `key` in the state `st` with the `value`.
Layer size
::: warning
These specifications have been added very recently and most layers currently do not implement them.
:::
# LuxCore.inputsize —
Function.
```julia
inputsize(layer)
```
Return the input size of the layer.
# LuxCore.outputsize —
Function.
```julia
outputsize(layer, x, rng)
```
Return the output size of the layer. If `outputsize(layer)` is defined, that method takes precedence, else we compute the layer output to determine the final size.
The fallback implementation of this function assumes the inputs were batched, i.e., if any of the outputs are Arrays, with `ndims(A) > 1`, it will return `size(A)[1:(end - 1)]`. If this behavior is undesirable, provide a custom `outputsize(layer, x, rng)` implementation).