Boltz.Layers and Boltz.Basis API Reference

Layers API

# Boltz.Layers.ConvBatchNormActivation — Function.

ConvBatchNormActivation(kernel_size::Dims, (in_filters, out_filters)::Pair{Int, Int},
    depth::Int, act::F; use_norm::Bool=true, conv_kwargs=(;),
    last_layer_activation::Bool=true, norm_kwargs=(;), flatten_model=false) where {F}

This function is a convenience wrapper around ConvNormActivation that constructs a chain with norm_layer set to Lux.BatchNorm if use_norm is true and nothing otherwise. In most cases, users should use ConvNormActivation directly for a more flexible interface.

source

# Boltz.Layers.ConvNormActivation — Function.

ConvNormActivation(kernel_size::Dims, in_chs::Integer, hidden_chs::Dims{N},
    activation; norm_layer=nothing, conv_kwargs=(;), norm_kwargs=(;),
    last_layer_activation::Bool=false, flatten_model::Bool=false) where {N}

Construct a Chain of convolutional layers with normalization and activation functions.

Arguments

• kernel_size: size of the convolutional kernel

• in_chs: number of input channels

• hidden_chs: dimensions of the hidden layers

• activation: activation function

Keyword Arguments

• norm_layer: Function with signature f(i::Integer, dims::Integer, act::F; kwargs...). i is the location of the layer in the model, dims is the channel dimension of the input, and act is the activation function. kwargs are forwarded from the norm_kwargs input, The function should return a normalization layer. Defaults to nothing, which means no normalization layer is used

• conv_kwargs: keyword arguments for the convolutional layers

• norm_kwargs: keyword arguments for the normalization layers

• last_layer_activation: set to true to apply the activation function to the last layer

Internal Keyword Arguments

Don't rely on these, they are for internal use only.

• flatten_model: set to true construct a flat chain without internal chains (not recommended)

source

# Boltz.Layers.ClassTokens — Type.

ClassTokens(dim; init=zeros32)

Appends class tokens to an input with embedding dimension dim for use in many vision transformer models.

source

# Boltz.Layers.HamiltonianNN — Type.

HamiltonianNN{FST}(model; autodiff=nothing) where {FST}

Constructs a Hamiltonian Neural Network [1]. This neural network is useful for learning symmetries and conservation laws by supervision on the gradients of the trajectories. It takes as input a concatenated vector of length 2n containing the position (of size n) and momentum (of size n) of the particles. It then returns the time derivatives for position and momentum.

Arguments

• FST: If true, then the type of the state returned by the model must be same as the type of the input state. See the documentation on StatefulLuxLayer for more information.

• model: A Lux.AbstractExplicitLayer neural network that returns the Hamiltonian of the system. The model must return a "batched scalar", i.e. all the dimensions of the output except the last one must be equal to 1. The last dimension must be equal to the batchsize of the input.

Keyword Arguments

• autodiff: The autodiff framework to be used for the internal Hamiltonian computation. The default is nothing, which selects the best possible backend available. The available options are AutoForwardDiff and AutoZygote.

Autodiff Backends

autodiff	Package Needed	Notes
AutoZygote	Zygote.jl	Preferred Backend. Chosen if Zygote is loaded and autodiff is nothing.
AutoForwardDiff	ForwardDiff.jl	Chosen if ForwardDiff is loaded, Zygote is not loaded and autodiff is nothing.

Note

This layer uses nested autodiff. Please refer to the manual entry on Nested Autodiff for more information and known limitations.

References

[1] Greydanus, Samuel, Misko Dzamba, and Jason Yosinski. "Hamiltonian Neural Networks." Advances in Neural Information Processing Systems 32 (2019): 15379-15389.

source

# Boltz.Layers.MultiHeadSelfAttention — Function.

MultiHeadSelfAttention(in_planes::Int, number_heads::Int; qkv_bias::Bool=false,
    attention_dropout_rate::T=0.0f0, projection_dropout_rate::T=0.0f0)

Multi-head self-attention layer

Arguments

• planes: number of input channels

• nheads: number of heads

• qkv_bias: whether to use bias in the layer to get the query, key and value

• attn_dropout_prob: dropout probability after the self-attention layer

• proj_dropout_prob: dropout probability after the projection layer

source

# Boltz.Layers.MLP — Function.

MLP(in_dims::Integer, hidden_dims::Dims{N}, activation=NNlib.relu; norm_layer=nothing,
    dropout_rate::Real=0.0f0, dense_kwargs=(;), norm_kwargs=(;),
    last_layer_activation=false) where {N}

Construct a multi-layer perceptron (MLP) with dense layers, optional normalization layers, and dropout.

Arguments

• in_dims: number of input dimensions

• hidden_dims: dimensions of the hidden layers

• activation: activation function (stacked after the normalization layer, if present else after the dense layer)

Keyword Arguments

• norm_layer: Function with signature f(i::Integer, dims::Integer, act::F; kwargs...). i is the location of the layer in the model, dims is the channel dimension of the input, and act is the activation function. kwargs are forwarded from the norm_kwargs input, The function should return a normalization layer. Defaults to nothing, which means no normalization layer is used

• dropout_rate: dropout rate (default: 0.0f0)

• dense_kwargs: keyword arguments for the dense layers

• norm_kwargs: keyword arguments for the normalization layers

• last_layer_activation: set to true to apply the activation function to the last layer

source

# Boltz.Layers.SplineLayer — Type.

SplineLayer(in_dims, grid_min, grid_max, grid_step, basis::Type{Basis};
    train_grid::Union{Val, Bool}=Val(false), init_saved_points=nothing)

Constructs a spline layer with the given basis function.

Arguments

• in_dims: input dimensions of the layer. This must be a tuple of integers, to construct a flat vector of saved_points pass in ().

• grid_min: minimum value of the grid.

• grid_max: maximum value of the grid.

• grid_step: step size of the grid.

• basis: basis function to use for the interpolation. Currently only the basis functions from DataInterpolations.jl are supported:

  1. ConstantInterpolation

  2. LinearInterpolation

  3. QuadraticInterpolation

  4. QuadraticSpline

  5. CubicSpline

Keyword Arguments

• train_grid: whether to train the grid or not.

• init_saved_points: values of the function at multiples of the time step. Initialized by default to a random vector sampled from the unit normal. Alternatively, can take a function with the signature init_saved_points(rng, in_dims, grid_min, grid_max, grid_step).

Warning

Currently this layer is limited since it relies on DataInterpolations.jl which doesn't work with GPU arrays. This will be fixed in the future by extending support to different basis functions

source

# Boltz.Layers.TensorProductLayer — Function.

TensorProductLayer(model, out_dim::Int; init_weight = randn32)

Constructs the Tensor Product Layer, which takes as input an array of n tensor product basis, $[B_1,B_2,\dots ,B_n]$ a data point x, computes

$$z_i=W_{i,:}\odot [B_1(x_1)\otimes B_2(x_2)\otimes \cdots \otimes B_n(x_n)]$$

where $W$ is the layer's weight, and returns $[z_1,\dots ,z_{out}]$.

Arguments

• basis_fns: Array of TensorProductBasis $[B_1(n_1),\dots ,B_k(n_k)]$, where $k$ corresponds to the dimension of the input.

• out_dim: Dimension of the output.

• init_weight: Initializer for the weight matrix. Defaults to randn32.

Warning

This layer currently only works on CPU and CUDA devices.

source

# Boltz.Layers.ViPosEmbedding — Type.

ViPosEmbedding(embedding_size, number_patches; init = randn32)

Positional embedding layer used by many vision transformer-like models.

source

# Boltz.Layers.VisionTransformerEncoder — Function.

VisionTransformerEncoder(in_planes, depth, number_heads; mlp_ratio = 4.0f0,
    dropout = 0.0f0)

Transformer as used in the base ViT architecture.

Arguments

• in_planes: number of input channels

• depth: number of attention blocks

• number_heads: number of attention heads

Keyword Arguments

• mlp_ratio: ratio of MLP layers to the number of input channels

• dropout_rate: dropout rate

References

[1] Dosovitskiy, Alexey, et al. "An image is worth 16x16 words: Transformers for image recognition at scale." arXiv preprint arXiv:2010.11929 (2020).

source

Basis Functions

Warning

The function calls for these basis functions should be considered experimental and are subject to change without deprecation. However, the functions themselves are stable and can be freely used in combination with the other Layers and Models.

# Boltz.Basis.Cos — Function.

Cos(n; dim::Int=1)

Constructs a cosine basis of the form $[\cos (x),\cos (2x),\dots ,\cos (nx)]$.

Arguments

• n: number of terms in the cosine expansion.

Keyword Arguments

• dim::Int=1: The dimension along which the basis functions are applied.

source

# Boltz.Basis.Chebyshev — Function.

Chebyshev(n; dim::Int=1)

Constructs a Chebyshev basis of the form $[T_0(x),T_1(x),\dots ,T_{n-1}(x)]$ where $T_j(.)$ is the $j^{th}$ Chebyshev polynomial of the first kind.

Arguments

• n: number of terms in the polynomial expansion.

Keyword Arguments

• dim::Int=1: The dimension along which the basis functions are applied.

source

# Boltz.Basis.Fourier — Function.

Fourier(n; dim=1)

Constructs a Fourier basis of the form

$$F_j(x)=\{\begin{array}{ll}cos\left(\frac{j}{2}x\right)& \text{if }j\text{ is even}\\ sin\left(\frac{j}{2}x\right)& \text{if }j\text{ is odd}\end{array}$$

Arguments

• n: number of terms in the Fourier expansion.

Keyword Arguments

• dim::Int=1: The dimension along which the basis functions are applied.

source

# Boltz.Basis.Legendre — Function.

Legendre(n; dim::Int=1)

Constructs a Legendre basis of the form $[P_0(x),P_1(x),\dots ,P_{n-1}(x)]$ where $P_j(.)$ is the $j^{th}$ Legendre polynomial.

Arguments

• n: number of terms in the polynomial expansion.

Keyword Arguments

• dim::Int=1: The dimension along which the basis functions are applied.

source

# Boltz.Basis.Polynomial — Function.

Polynomial(n; dim::Int=1)

Constructs a Polynomial basis of the form $[1,x,\dots ,x^{(n-1)}]$.

Arguments

• n: number of terms in the polynomial expansion.

Keyword Arguments

• dim::Int=1: The dimension along which the basis functions are applied.

source

# Boltz.Basis.Sin — Function.

Sin(n; dim::Int=1)

Constructs a sine basis of the form $[\sin (x),\sin (2x),\dots ,\sin (nx)]$.

Arguments

• n: number of terms in the sine expansion.

Keyword Arguments

• dim::Int=1: The dimension along which the basis functions are applied.

source