Lux.jl

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.AbstractExplicitLayerType. ```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.AbstractExplicitContainerLayerType. ```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.applyFunction. ```julia apply(model, x, ps, st) ``` Simply calls `model(x, ps, st)`


# LuxCore.stateless_applyFunction. ```julia stateless_apply(model, x, ps) ``` Calls `apply` and only returns the first argument.


# LuxCore.check_fmap_conditionFunction. ```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_layerFunction. ```julia contains_lux_layer(l) -> Bool ``` Check if the structure `l` is a Lux AbstractExplicitLayer or a container of such a layer.


# LuxCore.display_nameFunction. ```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.replicateFunction. ```julia replicate(rng::AbstractRNG) ``` Creates a copy of the `rng` state depending on its type.


# LuxCore.setupFunction. ```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.initialparametersFunction. ```julia initialparameters(rng::AbstractRNG, layer) ``` Generate the initial parameters of the layer `l`.


# LuxCore.parameterlengthFunction. ```julia parameterlength(layer) ``` Return the total number of parameters of the layer `l`.


States

# LuxCore.initialstatesFunction. ```julia initialstates(rng::AbstractRNG, layer) ``` Generate the initial states of the layer `l`.


# LuxCore.statelengthFunction. ```julia statelength(layer) ``` Return the total number of states of the layer `l`.


# LuxCore.testmodeFunction. ```julia testmode(st::NamedTuple) ``` Make all occurances of `training` in state `st` – `Val(false)`.


# LuxCore.trainmodeFunction. ```julia trainmode(st::NamedTuple) ``` Make all occurances of `training` in state `st` – `Val(true)`.


# LuxCore.update_stateFunction. ```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.inputsizeFunction. ```julia inputsize(layer) ``` Return the input size of the layer.


# LuxCore.outputsizeFunction. ```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).