Layers

Layers#

Implements Bayesian Layers using Jax and Numpyro.

Design:

There are three levels of complexity here: class-level, instance-level, and call-level
The class-level handles things like choosing generic model form and how to multiply coefficents with data. Defined by the class Layer(BLayer) def itself.
The instance-level handles specific distributions that fit into a generic model and the initial parameters for those distributions. Defined by creating an instance of the class: Layer(*args, **kwargs).
The call-level handles seeing a batch of data, sampling from the distributions defined on the class and multiplying coefficients and data to produce an output, works like result = Layer(*args, **kwargs)(data)

Notation:

n: observations in a batch
c: number of categories of things for time, random effects, etc
d: number of coefficients
l: low rank dimension of low rank models
m: embedding dimension
u: units aka output dimension

class blayers.layers.BLayer(*args)[source]#

Bases: ABC

Abstract base class for Bayesian layers. Lays out an interface.

Parameters:: args (Any)

abstractmethod __init__(*args)[source]#

Initialize layer parameters. This is the Bayesian model.

Parameters:: args (Any)
Return type:: None

abstractmethod __call__(*args)[source]#

Run the layer’s forward pass.

Parameters:: *args (Any) – Inputs to the layer.
Returns:: The result of the forward computation.
Return type:: jax.Array

class blayers.layers.AdaptiveLayer(lmbda_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, lmbda_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Bayesian layer with adaptive prior using hierarchical modeling.

Generates coefficients from the hierarchical model

\[\lambda \sim HalfNormal(1.)\]

\[\beta \sim Normal(0., \lambda)\]

Parameters:

lmbda_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
lmbda_kwargs (dict[str, float])

__init__(lmbda_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, lmbda_kwargs={'scale': 1.0})[source]#

Parameters:

lmbda_dist (Distribution) – NumPyro distribution class for the scale (λ) of the prior.
coef_dist (Distribution) – NumPyro distribution class for the coefficient prior.
coef_kwargs (dict[str, float]) – Parameters for the prior distribution.
lmbda_kwargs (dict[str, float]) – Parameters for the scale distribution.

__call__(name, x, units=1, activation=<PjitFunction of <function identity>>)[source]#

Forward pass with adaptive prior on coefficients.

Parameters:

name (str) – Variable name.
x (Array) – Input data array of shape (n, d).
units (int) – Number of outputs.
activation (Callable[[Array], Array]) – Activation function to apply to output.

Returns:

Output array of shape (n, u).

Return type:

jax.Array

class blayers.layers.FixedPriorLayer(coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0, 'scale': 1.0})[source]#

Bases: BLayer

Bayesian layer with a fixed prior distribution over coefficients.

Generates coefficients from the model

\[\beta \sim Normal(0., 1.)\]

Parameters:

coef_dist (Distribution)
coef_kwargs (dict[str, float])

__init__(coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0, 'scale': 1.0})[source]#

Parameters:

coef_dist (Distribution) – NumPyro distribution class for the coefficients.
coef_kwargs (dict[str, float]) – Parameters to initialize the prior distribution.

__call__(name, x, units=1, activation=<PjitFunction of <function identity>>)[source]#

Forward pass with fixed prior.

Parameters:

name (str) – Variable name.
x (Array) – Input data array of shape (n, d).
units (int) – Number of outputs.
activation (Callable[[Array], Array]) – Activation function to apply to output.

Returns:

Output array of shape (n, u).

Return type:

jax.Array

class blayers.layers.InterceptLayer(coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0, 'scale': 1.0})[source]#

Bases: BLayer

Bayesian layer with a fixed prior distribution over coefficients.

Generates coefficients from the model

\[\beta \sim Normal(0., 1.)\]

Parameters:

coef_dist (Distribution)
coef_kwargs (dict[str, float])

__init__(coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0, 'scale': 1.0})[source]#

Parameters:

coef_dist (Distribution) – NumPyro distribution class for the coefficients.
coef_kwargs (dict[str, float]) – Parameters to initialize the prior distribution.

__call__(name, units=1, activation=<PjitFunction of <function identity>>)[source]#

Forward pass with fixed prior.

Parameters:

name (str) – Variable name.
units (int) – Number of outputs.
activation (Callable[[Array], Array]) – Activation function to apply to output.

Returns:

Output array of shape (1, u).

Return type:

jax.Array

class blayers.layers.FMLayer(lmbda_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, lmbda_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Bayesian factorization machine layer with adaptive priors.

Generates coefficients from the hierarchical model

\[\lambda \sim HalfNormal(1.)\]

\[\beta \sim Normal(0., \lambda)\]

The shape of beta is (j, l), where j is the number if input covariates and l is the low rank dim.

Then performs matrix multiplication using the formula in Rendle (2010).

Parameters:

lmbda_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
lmbda_kwargs (dict[str, float])

Parameters:

lmbda_dist (Distribution) – Distribution for scaling factor λ.
coef_dist (Distribution) – Prior for beta parameters.
coef_kwargs (dict[str, float]) – Arguments for prior distribution.
lmbda_kwargs (dict[str, float]) – Arguments for λ distribution.

__call__(name, x, low_rank_dim, units=1, activation=<PjitFunction of <function identity>>)[source]#

Forward pass through the factorization machine layer.

Parameters:

name (str) – Variable name scope.
x (Array) – Input matrix of shape (n, d).
low_rank_dim (int) – Dimensionality of low-rank approximation.
units (int) – Number of outputs.
activation (Callable[[Array], Array]) – Activation function to apply to output.

Returns:

Output array of shape (n, u).

Return type:

jax.Array

class blayers.layers.FM3Layer(lmbda_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, lmbda_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Order 3 FM. See Blondel et al 2016.

Parameters:

lmbda_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
lmbda_kwargs (dict[str, float])

Parameters:

lmbda_dist (Distribution) – Distribution for scaling factor λ.
coef_dist (Distribution) – Prior for beta parameters.
coef_kwargs (dict[str, float]) – Arguments for prior distribution.
lmbda_kwargs (dict[str, float]) – Arguments for λ distribution.

__call__(name, x, low_rank_dim, units=1, activation=<PjitFunction of <function identity>>)[source]#

Forward pass through the factorization machine layer.

Parameters:

name (str) – Variable name scope.
x (Array) – Input matrix of shape (n, d).
low_rank_dim (int) – Dimensionality of low-rank approximation.
units (int) – Number of outputs.
activation (Callable[[Array], Array]) – Activation function to apply to output.

Returns:

Output array of shape (n,).

Return type:

jax.Array

class blayers.layers.LowRankInteractionLayer(lmbda_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, lmbda_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Takes two sets of features and learns a low-rank interaction matrix.

Parameters:

lmbda_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
lmbda_kwargs (dict[str, float])

Initialize layer parameters. This is the Bayesian model.

Parameters:

lmbda_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
lmbda_kwargs (dict[str, float])

__call__(name, x, z, low_rank_dim, units=1, activation=<PjitFunction of <function identity>>)[source]#

Interaction between feature matrices X and Z in a low rank way. UV decomp.

Parameters:

name (str) – Variable name scope.
x (Array) – Input matrix of shape (n, d1).
z (Array) – Input matrix of shape (n, d2).
low_rank_dim (int) – Dimensionality of low-rank approximation.
units (int) – Number of outputs.
activation (Callable[[Array], Array]) – Activation function to apply to output.

Returns:

Output array of shape (n, u).

Return type:

jax.Array

class blayers.layers.InteractionLayer(lmbda_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, lmbda_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Computes every interaction coefficient between two sets of inputs.

Parameters:

lmbda_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
lmbda_kwargs (dict[str, float])

Initialize layer parameters. This is the Bayesian model.

Parameters:

lmbda_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
lmbda_kwargs (dict[str, float])

__call__(name, x, z, units=1, activation=<PjitFunction of <function identity>>)[source]#

Interaction between feature matrices X and Z in a low rank way. UV decomp.

Parameters:

name (str) – Variable name scope.
x (Array) – Input matrix of shape (n, d1).
z (Array) – Input matrix of shape (n, d2).
units (int) – Number of outputs.
activation (Callable[[Array], Array]) – Activation function to apply to output.

Returns:

Output array of shape (n, u).

Return type:

jax.Array

class blayers.layers.BilinearLayer(lmbda_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, lmbda_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Bayesian bilinear interaction layer: computes x^T W z.

Parameters:

lmbda_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
lmbda_kwargs (dict[str, float])

Parameters:

lmbda_dist (Distribution) – prior on scale of coefficients
coef_dist (Distribution) – distribution for coefficients
coef_kwargs (dict[str, float]) – kwargs for coef distribution
lmbda_kwargs (dict[str, float]) – kwargs for scale prior

__call__(name, x, z, units=1, activation=<PjitFunction of <function identity>>)[source]#

Interaction between feature matrices X and Z in a low rank way. UV decomp.

Parameters:

name (str) – Variable name scope.
x (Array) – Input matrix of shape (n, d1).
z (Array) – Input matrix of shape (n, d2).
units (int) – Number of outputs.
activation (Callable[[Array], Array]) – Activation function to apply to output.

Returns:

Output array of shape (n, u).

Return type:

jax.Array

class blayers.layers.LowRankBilinearLayer(lmbda_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, lmbda_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Bayesian bilinear interaction layer: computes x^T W z. W low rank.

Parameters:

lmbda_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
lmbda_kwargs (dict[str, float])

Parameters:

lmbda_dist (Distribution) – prior on scale of coefficients
coef_dist (Distribution) – distribution for coefficients
coef_kwargs (dict[str, float]) – kwargs for coef distribution
lmbda_kwargs (dict[str, float]) – kwargs for scale prior

__call__(name, x, z, low_rank_dim, units=1, activation=<PjitFunction of <function identity>>)[source]#

Interaction between feature matrices X and Z in a low rank way. UV decomp.

Parameters:

name (str) – Variable name scope.
x (Array) – Input matrix of shape (n, d1).
z (Array) – Input matrix of shape (n, d2).
low_rank_dim (int) – Dimensionality of low-rank approximation.
units (int) – Number of outputs.
activation (Callable[[Array], Array]) – Activation function to apply to output.

Returns:

Output array of shape (n, u).

Return type:

jax.Array

class blayers.layers.EmbeddingLayer(lmbda_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, lmbda_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Bayesian embedding layer for sparse categorical features.

Parameters:

lmbda_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
lmbda_kwargs (dict[str, float])

Parameters:

lmbda_dist (Distribution) – NumPyro distribution class for the scale (λ) of the prior.
coef_dist (Distribution) – NumPyro distribution class for the coefficient prior.
coef_kwargs (dict[str, float]) – Parameters for the prior distribution.
lmbda_kwargs (dict[str, float]) – Parameters for the scale distribution.

__call__(name, x, num_categories, embedding_dim)[source]#

Forward pass through embedding lookup.

Parameters:

name (str) – Variable name scope.
x (Array) – Integer indices indicating embeddings to use.
num_categories (int) – The number of distinct things getting an embedding
embedding_dim (int) – The size of each embedding, e.g. 2, 4, 8, etc.

Returns:

Embedding vectors of shape (n, m).

Return type:

jax.Array

class blayers.layers.RandomEffectsLayer(lmbda_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, lmbda_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Exactly like the EmbeddingLayer but with embedding_dim=1.

Parameters:

lmbda_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
lmbda_kwargs (dict[str, float])

Parameters:

num_embeddings – Total number of discrete embedding entries.
embedding_dim – Dimensionality of each embedding vector.
coef_dist (Distribution) – Prior distribution for embedding weights.
coef_kwargs (dict[str, float]) – Parameters for the prior distribution.
lmbda_dist (Distribution)
lmbda_kwargs (dict[str, float])

__call__(name, x, num_categories)[source]#

Forward pass through embedding lookup.

Parameters:

name (str) – Variable name scope.
x (Array) – Integer indicating embeddings to use.
num_categories (int) – The number of distinct things getting an embedding

Returns:

Embedding vectors of shape (n, embedding_dim).

Return type:

jax.Array

class blayers.layers.RandomWalkLayer(lmbda_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, lmbda_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Random walk of embedding dim m, defaults to Gaussian walk.

Parameters:

lmbda_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
lmbda_kwargs (dict[str, float])

Initialize layer parameters. This is the Bayesian model.

Parameters:

lmbda_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
lmbda_kwargs (dict[str, float])

__call__(name, x, num_categories, embedding_dim)[source]#

Forward pass through embedding lookup.

Parameters:

name (str) – Variable name scope.
x (Array) – Integer indices indicating embeddings to use.
num_categories (int) – The number of distinct things getting an embedding
embedding_dim (int) – The size of each embedding, e.g. 2, 4, 8, etc.

Returns:

Embedding vectors of shape (n, m).

Return type:

jax.Array

Layers

Contents

Layers#