Neural Processes

https://github.com/wesselb/neuralprocesses
Admin
A framework for composing Neural Processes in Python.
Tokens:118,341
Snippets:616
Trust Score:9.7
License:MIT
Update:16 hours ago
Show doc for...
Tokens:
Raw
### Install neuralprocesses with TensorFlow or PyTorch

Source: https://github.com/wesselb/neuralprocesses/blob/main/README.md

Install the neuralprocesses package along with the required backend. Use the first command for TensorFlow and the second for PyTorch.

```bash
pip install neuralprocesses tensorflow tensorflow-probability  # For use with TensorFlow
```

```bash
pip install neuralprocesses torch                              # For use with PyTorch
```

--------------------------------

### SetConv usage examples

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/coders-nn.md

Shows how to instantiate SetConv with default learnable scale and with a fixed scale, and how to apply it using nps.code.

```python
set_conv = nps.SetConv(scale=0.1)          # length scale learnable by default
set_conv = nps.SetConv(scale=1 / 64, learnable=False)
xz, z = nps.code(set_conv, xz, z, x_grid, root=True)
```

--------------------------------

### MLP usage example

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/coders-nn.md

Shows how to instantiate and apply an MLP on (batch, in_dim, n) data. The example uses in_dim=3, layers=(128, 128), out_dim=16, and ReLU nonlinearity.

```python
# MLP operating on (batch, in_dim, n) data:
mlp = nps.MLP(in_dim=3, layers=(128, 128), out_dim=16, nonlinearity="ReLU")
z_new = mlp(z)  # (batch, 16, n)
```

--------------------------------

### Construct a ConvGNP model with two context sets

Source: https://github.com/wesselb/neuralprocesses/blob/main/docs/advanced_usage.md

Sets up a ConvGNP model with two context sets, one gridded and one non-gridded, and creates sample context and target inputs. This is the running example for the masking section.

```python
import lab as B
import torch

import neuralprocesses.torch as nps

cnp = nps.construct_convgnp(
    dim_x=2,
    dim_yc=(1, 1),  # Two context sets, both with one channel
    dim_yt=1, 
)

# Construct two sample context sets with one on a grid.
xc = B.randn(torch.float32, 1, 2, 20)
yc = B.randn(torch.float32, 1, 1, 20)
xc_grid = (B.randn(torch.float32, 1, 1, 10), B.randn(torch.float32, 1, 1, 15))
yc_grid = B.randn(torch.float32, 1, 1, 10, 15)

# Contruct sample target inputs
xt = B.randn(torch.float32, 1, 2, 50)
```

--------------------------------

### Usage Example

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/distributions.md

Example of sampling from a model output and applying a transformation.

```APIDOC
## Usage Example

```python
import neuralprocesses.torch as nps

# Sample from a multi-output normal built from a model output `pred`:
state = B.global_random_state(B.dtype_float(xt))
state, samples = pred.sample(state, num_samples=10)
logpdfs = pred.logpdf(yt)

# Transformed predictions (e.g. positive data):
transformed = nps.Transform.exp()(pred)
mean = transformed.mean  # exp of the underlying mean
```
```

--------------------------------

### Discretisation usage example

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/utilities.md

Creates a Discretisation instance with points_per_unit=32, multiple=1, margin=0.1, and dim=1, then applies it to xt to produce a grid of shape (batch, 1, n_grid).

```python
# Grid for 1D targets:
grid = nps.Discretisation(points_per_unit=32, multiple=1, margin=0.1, dim=1)(xt)
# grid has shape (batch, 1, n_grid)
```

--------------------------------

### predict usage example

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/model.md

Shows how to call predict to obtain mean, variance, noiseless samples, and noisy samples. The example uses num_samples=100.

```python
mean, var, noiseless_samples, noisy_samples = nps.predict(
    model,
    xc, yc, xt,
    num_samples=100,
)

```

--------------------------------

### Construct GNP model and inspect parameters

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/backends.md

Imports the backend module (PyTorch or TensorFlow), constructs a Gaussian neural process with heteroscedastic likelihood, prints the number of parameters, and demonstrates using core symbols like Chain, DeepSet, MLP, and DeterministicLikelihood. The example shows that the model is an instance of the framework-native Module.

```python
import neuralprocesses.torch as nps  # or: import neuralprocesses.tensorflow as nps

model = nps.construct_gnp(dim_x=1, dim_y=1, likelihood="het")
print(nps.num_params(model))         # int: number of parameters/variables

# All core symbols are available on the backend namespace:
encoder = nps.Chain(nps.DeepSet(nps.MLP(in_dim=2, layers=(64,), out_dim=64)), nps.DeterministicLikelihood())
print(isinstance(model, nps.Module)) # True for the framework-native Module
```

--------------------------------

### Basic model call with context set and sampling

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/model.md

Shows how to construct a ConvGNP model and call it with a context set, target inputs, and num_samples. The prediction is a MultiOutputNormal; sample from it to get samples of shape (16, batch, 2, n).

```python
import neuralprocesses.torch as nps

model = nps.construct_convgnp(dim_x=1, dim_y=2, likelihood="het")
pred = model(
    [(xc, yc)],        # context set
    xt,                # target inputs, shape (batch, 1, n)
    num_samples=16,
)
# pred is a MultiOutputNormal; sample or evaluate logpdf on it.
samples = pred.sample()          # (16, batch, 2, n) without state
```

--------------------------------

### Usage example: sampling and transforming predictions

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/distributions.md

Shows how to sample from a multi-output normal built from a model output `pred` and how to apply an exponential transform for positive data. The example uses neuralprocesses.torch and B.global_random_state.

```python
import neuralprocesses.torch as nps

# Sample from a multi-output normal built from a model output `pred`:
state = B.global_random_state(B.dtype_float(xt))
state, samples = pred.sample(state, num_samples=10)
logpdfs = pred.logpdf(yt)

# Transformed predictions (e.g. positive data):
transformed = nps.Transform.exp()(pred)
mean = transformed.mean  # exp of the underlying mean
```

--------------------------------

### Construct and use a ConvGNP with auxiliary variables

Source: https://github.com/wesselb/neuralprocesses/blob/main/docs/basic_usage.md

Constructs a ConvGNP with auxiliary variables using `nps.construct_convgnp`, specifying `dim_yc` for observed and auxiliary output dimensions, `dim_aux_t` for target-specific auxiliary dimensions, `dim_yt` for prediction dimensions, `num_basis_functions`, and `likelihood="lowrank"`. The example defines three auxiliary variables: one at arbitrary inputs, one on a grid, and one specific to target inputs. The model is called with a list of observed data and auxiliary variables, target inputs, and the target-specific auxiliary variable as a keyword argument. The mean and variance are extracted, and the example prints the log probability of a random target, a sample, the KL divergence to itself, and the entropy.

```python
import lab as B
import tensorflow as tf

import neuralprocesses.tensorflow as nps

cnp = nps.construct_convgnp(
    dim_x=2,
    dim_yc=(
        3,  # Observed data has three dimensions.
        1,  # First auxiliary variable has one dimension.
        2,  # Second auxiliary variable has two dimensions.
    ),
    # Third auxiliary variable has four dimensions and is auxiliary information specific
    # to the target inputs.
    dim_aux_t=4,
    dim_yt=3,  # Predictions have three dimensions.
    num_basis_functions=64, 
    likelihood="lowrank",
)

observed_data = (
    B.randn(tf.float32, 16, 2, 10),
    B.randn(tf.float32, 16, 3, 10),
)

# Define three auxiliary variables. The first one is specified like the observed data
# at arbitrary inputs.
aux_var1 = (
    B.randn(tf.float32, 16, 2, 12),
    B.randn(tf.float32, 16, 1, 12),  # Has one dimension.
)
# The second one is specified on a grid.
aux_var2 = (
    (B.randn(tf.float32, 16, 1, 25), B.randn(tf.float32, 16, 1, 35)),
    B.randn(tf.float32, 16, 2, 25, 35),  # Has two dimensions.
)
# The third one is specific to the target inputs. We could encode it like the first
# auxiliary variable `aux_var1`, but we illustrate how an MLP-style encoding can
# also be used. The number must match the number of target inputs!
aux_var_t = B.randn(tf.float32, 16, 4, 15)  # Has four dimensions.

dist = cnp(
    [observed_data, aux_var1, aux_var2],
    B.randn(tf.float32, 16, 2, 15),
    aux_t=aux_var_t,  # This must be given as a keyword argument.
)
mean, var = dist.mean, dist.var

print(dist.logpdf(B.randn(tf.float32, 16, 3, 15)))
print(dist.sample())
print(dist.kl(dist))
print(dist.entropy())
```

--------------------------------

### Training usage with loglik and elbo

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/model.md

Example from the README showing how to compute the loss for training. The first line uses `loglik` for non-latent models, and the second uses `elbo` for latent-variable models with `num_samples=4`.

```python
loss = -torch.mean(nps.loglik(model, xc, yc, xt, yt, normalise=True))
# or, for a latent-variable model:
loss = -torch.mean(nps.elbo(model, xc, yc, xt, yt, num_samples=4, normalise=True))
```

--------------------------------

### Quick-start: Construct, train, and predict with a ConvCNP in PyTorch

Source: https://github.com/wesselb/neuralprocesses/blob/main/README.md

This example shows how to construct a ConvCNP model using neuralprocesses.torch, train it with an Adam optimizer for 32 batches, and make predictions. The context and target sets are randomly generated; replace them with your own data. The loss is the negative mean log-likelihood, and predictions return mean, variance, and samples.

```python
import torch

import neuralprocesses.torch as nps

# Construct a ConvCNP.
convcnp = nps.construct_convgnp(dim_x=1, dim_y=2, likelihood="het")

# Construct optimiser.
opt = torch.optim.Adam(convcnp.parameters(), 1e-3)

# Training: optimise the model for 32 batches.
for _ in range(32):
    # Sample a batch of new context and target sets. Replace this with your data. The
    # shapes are `(batch_size, dimensionality, num_data)`.
    xc = torch.randn(16, 1, 10)  # Context inputs
    yc = torch.randn(16, 2, 10)  # Context outputs
    xt = torch.randn(16, 1, 15)  # Target inputs
    yt = torch.randn(16, 2, 15)  # Target output

    # Compute the loss and update the model parameters.
    loss = -torch.mean(nps.loglik(convcnp, xc, yc, xt, yt, normalise=True))
    opt.zero_grad(set_to_none=True)
    loss.backward()
    opt.step()

# Testing: make some predictions.
mean, var, noiseless_samples, noisy_samples = nps.predict(
    convcnp,
    torch.randn(16, 1, 10),  # Context inputs
    torch.randn(16, 2, 10),  # Context outputs
    torch.randn(16, 1, 15),  # Target inputs
)
```

--------------------------------

### Construct and use a ConvGNP with TensorFlow

Source: https://github.com/wesselb/neuralprocesses/blob/main/docs/basic_usage.md

Constructs a Convolutional Gaussian Neural Process (ConvGNP) with `nps.construct_convgnp`, using the same parameters as the GNP example. The model is called with context and target inputs to produce a distribution, and the mean and variance are extracted. The example prints the log probability of a random target, a sample, the KL divergence to itself, and the entropy.

```python
import lab as B
import tensorflow as tf

import neuralprocesses.tensorflow as nps

cnp = nps.construct_convgnp(dim_x=2, dim_y=3, likelihood="lowrank")
dist = cnp(
    B.randn(tf.float32, 16, 2, 10),  # Context inputs
    B.randn(tf.float32, 16, 3, 10),  # Context outputs
    B.randn(tf.float32, 16, 2, 15),  # Target inputs
)
mean, var = dist.mean, dist.var  # Prediction for target outputs

print(dist.logpdf(B.randn(tf.float32, 16, 3, 15)))
print(dist.sample())
print(dist.kl(dist))
print(dist.entropy())
```

--------------------------------

### Construct and evaluate a ConvGNP with two context sets

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/architectures.md

Builds a ConvGNP model with two context sets (one gridded) and a heteroscedastic likelihood, then computes predictions and the negative log-likelihood loss. Requires PyTorch and neuralprocesses.torch. The example uses random tensors for context and target sets.

```python
import torch
import neuralprocesses.torch as nps

# ConvGNP with two context sets (one gridded), heteroscedastic likelihood.
model = nps.construct_convgnp(
    dim_x=2,
    dim_yc=(1, 1),
    dim_yt=1,
    likelihood="het",
    points_per_unit=32,
)

xc = torch.randn(16, 2, 20)
yc = torch.randn(16, 1, 20)
xc_grid = (torch.randn(16, 1, 10), torch.randn(16, 1, 15))
yc_grid = torch.randn(16, 1, 10, 15)
xt = torch.randn(16, 2, 50)

pred = model([(xc, yc), (xc_grid, yc_grid)], xt)
loss = -torch.mean(nps.loglik(model, [(xc, yc), (xc_grid, yc_grid)], xt, yt, normalise=True))
```

--------------------------------

### Construct and use a GNP with TensorFlow

Source: https://github.com/wesselb/neuralprocesses/blob/main/docs/basic_usage.md

Constructs a Gaussian Neural Process (GNP) with `nps.construct_gnp`, specifying `dim_x=2`, `dim_y=3`, and `likelihood="lowrank"`. The model is called with context and target inputs to produce a distribution, from which the mean and variance are extracted. The example then prints the log probability of a random target, a sample, the KL divergence to itself, and the entropy.

```python
import lab as B
import tensorflow as tf

import neuralprocesses.tensorflow as nps

cnp = nps.construct_gnp(dim_x=2, dim_y=3, likelihood="lowrank")
dist = cnp(
    B.randn(tf.float32, 16, 2, 10),  # Context inputs
    B.randn(tf.float32, 16, 3, 10),  # Context outputs
    B.randn(tf.float32, 16, 2, 15),  # Target inputs
)
mean, var = dist.mean, dist.var  # Prediction for target outputs

print(dist.logpdf(B.randn(tf.float32, 16, 3, 15)))
print(dist.sample())
print(dist.kl(dist))
print(dist.entropy())
```

--------------------------------

### Import NeuralProcesses backend (PyTorch or TensorFlow)

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/README.md

Shows how to select the backend by importing the corresponding submodule. Use `neuralprocesses.torch` for PyTorch or `neuralprocesses.tensorflow` for TensorFlow.

```python
import neuralprocesses.torch as nps          # PyTorch
import neuralprocesses.tensorflow as nps    # TensorFlow
```

--------------------------------

### merge_contexts usage example

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/utilities.md

Example of merging two context sets of different sizes into one masked context. The resulting yc is a Masked object, which can be passed to a model along with the context inputs.

```python
# Batch two contexts of different sizes into one masked context:
xc1, yc1 = torch.randn(16, 1, 10), torch.randn(16, 1, 10)
xc2, yc2 = torch.randn(16, 1, 30), torch.randn(16, 1, 30)
xc, yc = nps.merge_contexts((xc1, yc1), (xc2, yc2))   # yc is Masked
pred = model([(xc, yc)], xt)
```

--------------------------------

### ELBO usage example combining prior and posterior

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/coding.md

Usage example from neuralprocesses/model/elbo.py showing how elbo uses recode_stochastic to construct the approximate posterior from the prior encoding. It first codes the prior with history, then recodes the stochastic parts for the merged context set, and finally samples from the posterior.

```python
# Prior (with history):
xz, pz, h = nps.code_track(model.encoder, *nps.compress_contexts(contexts), xt, root=True, dtype_lik=dtype_lik)
# Posterior: recode the stochastic parts for the merged context set:
qz = nps.recode_stochastic(model.encoder, pz, *nps.compress_contexts(contexts_q), h, root=True, dtype_lik=dtype_lik)
# Sample from the posterior:
state, z = nps.sample(state, qz, num_samples)
```

--------------------------------

### PredPreyGenerator.__init__

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/data.md

Initializes a generator for simulated hare-lynx predator-prey data (Lotka-Volterra with noise). Samples are drawn from big simulated batches cached across calls.

```APIDOC
## PredPreyGenerator.__init__

### Description
Initializes a generator for simulated hare-lynx predator-prey data (Lotka-Volterra with noise). Produces two outputs (hare, lynx). Samples are drawn from big simulated batches (`big_batch_size`) cached across calls.

### Method
Constructor

### Signature
```python
class PredPreyGenerator(DataGenerator):
    def __init__(
        self,
        dtype,
        seed=0,
        num_tasks=2**14,
        batch_size=16,
        big_batch_size=2048,
        dist_x=UniformContinuous(0, 100),
        num_data=UniformDiscrete(150, 250),
        num_target=UniformDiscrete(100, 100),
        forecast_start=UniformContinuous(25, 75),
        mode="random",
        device="cpu",
    ):
```

### Parameters
- **dtype** - Required - Data type.
- **seed** (int) - Optional, default `0` - Seed.
- **num_tasks** (int) - Optional, default `2**14` - Number of tasks.
- **batch_size** (int) - Optional, default `16` - Batch size.
- **big_batch_size** (int) - Optional, default `2048` - Size of big simulated batches.
- **dist_x** (distribution) - Optional, default `UniformContinuous(0, 100)` - Distribution for x values.
- **num_data** (distribution) - Optional, default `UniformDiscrete(150, 250)` - Distribution for number of data points.
- **num_target** (distribution) - Optional, default `UniformDiscrete(100, 100)` - Distribution for number of target points.
- **forecast_start** (distribution) - Optional, default `UniformContinuous(25, 75)` - Distribution for forecast start.
- **mode** (str) - Optional, default `"random"` - Passed to `apply_task`.
- **device** (str) - Optional, default `"cpu"` - Device to use.

### Notes
- Simulated hare–lynx predator–prey data (Lotka–Volterra with noise) — two outputs (hare, lynx).
```

--------------------------------

### construct_fullconvgnp — FullConvGNP options

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/configuration.md

Constructor options for the FullConvGNP model, including additional options beyond the shared set conv/discretisation options.

```APIDOC
## construct_fullconvgnp — FullConvGNP options

### Description
Constructor options for the FullConvGNP model. It shares the set conv/discretisation options with ConvGNP (minus `dim_aux_t`, `aux_t_mlp_layers`, and likelihood family — the output is a dense Gaussian), plus the following additional options.

### Parameters
- **kernel_factor** (int) - Optional - Channel/PPU reduction factor of the kernel CNN. Default: `2`.
- **dim_x** (int) - Optional - Only `1` supported (`NotImplementedError` otherwise). Default: `1`.
- **dim_lv** (int) - Optional - Only `0` supported. Default: `0`.
```

--------------------------------

### Evaluation with fixed noise

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/model.md

Example call to nps.loglik with fix_noise=1e-4 and normalise=True, returning state and log-likelihood. Used for evaluation with fixed noise.

```python
state, ll = nps.loglik(state, model, contexts, xt, yt, fix_noise=1e-4, normalise=True)
```

--------------------------------

### ConvNet.__init__

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/coders-nn.md

Initializes a ConvNet instance with the given configuration. The kernel size is automatically computed from points_per_unit and receptive_field if not provided.

```APIDOC
## ConvNet.__init__

### Description
Initializes a ConvNet instance with the given configuration. The kernel size is automatically computed from points_per_unit and receptive_field if not provided.

### Method
Constructor

### Signature
```python
ConvNet(
    dim: int,
    in_channels: int,
    out_channels: int,
    channels: int,
    num_layers: int,
    kernel: Optional[int] = None,
    points_per_unit: Optional[float] = 1,
    receptive_field: Optional[float] = None,
    separable: bool = True,
    residual: bool = False,
    dtype=None,
)
```

### Parameters
- **dim** (int) - Required - Dimensionality.
- **in_channels** (int) - Required - Input channels.
- **out_channels** (int) - Required - Output channels.
- **channels** (int) - Required - Channels at every intermediate layer.
- **num_layers** (int) - Required - Number of layers.
- **kernel** (int) - Optional - Kernel size; overrides the points_per_unit/receptive_field computation.
- **points_per_unit** (float) - Optional - Density of the discretisation of the inputs. Default: 1.
- **receptive_field** (float) - Optional - Desired receptive field.
- **separable** (bool) - Optional - Depthwise separable convolutions. Default: True.
- **residual** (bool) - Optional - Residual blocks. Default: False.
- **dtype** (dtype) - Optional - Data type.

### Notes
If kernel is None, the kernel size is computed as the odd integer ceil(1 + (receptive_points - 1) / num_layers) with receptive_points = receptive_field * points_per_unit; otherwise receptive_field is derived from kernel. Attributes: dim, kernel, num_halving_layers (0, for Discretisation compatibility), receptive_field, conv_net.
```

--------------------------------

### PredPreyRealGenerator.__init__

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/data.md

Initializes a generator that uses the real hare-lynx dataset loaded via `load()` in `data/predprey.py`.

```APIDOC
## PredPreyRealGenerator.__init__

### Description
Initializes a generator that uses the real hare-lynx dataset loaded via `load()` in `data/predprey.py`.

### Method
Constructor

### Signature
```python
class PredPreyRealGenerator(DataGenerator):
    def __init__(
        self,
        dtype,
        seed=0,
        num_tasks=2**10,
        num_target=UniformDiscrete(1, 20),
        forecast_start=UniformContinuous(25, 75),
        mode="interpolation",
        device="cpu",
    ):
```

### Parameters
- **dtype** - Required - Data type.
- **seed** (int) - Optional, default `0` - Seed.
- **num_tasks** (int) - Optional, default `2**10` - Number of tasks.
- **num_target** (distribution) - Optional, default `UniformDiscrete(1, 20)` - Distribution for number of target points.
- **forecast_start** (distribution) - Optional, default `UniformContinuous(25, 75)` - Distribution for forecast start.
- **mode** (str) - Optional, default `"interpolation"` - Task mode.
- **device** (str) - Optional, default `"cpu"` - Device to use.

### Notes
- Uses the real hare–lynx dataset (`load()` in `data/predprey.py`).
```

--------------------------------

### Applying a chain of coders manually

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/coding.md

Example of manually applying a chain of coders using `nps.code`. The `root=True` keyword is mandatory on the outermost call; the wrapper raises `RuntimeError` if it is missing for `code`, `code_track`, and `recode`.

```python
import neuralprocesses.torch as nps

coder = nps.Chain(
    nps.PrependDensityChannel(),
    nps.SetConv(scale=0.1),
    nps.DivideByFirstChannel(),
    nps.Concatenate(),
    nps.DeterministicLikelihood(),
)
# xz: context inputs (grid: tuple of tensors), z: context outputs
xz_new, z_new = nps.code(coder, xz, z, disc_targets, root=True)
```

--------------------------------

### TemperatureGenerator.__init__

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/data.md

Initializes a German climate (temperature/elevation) data generator. Exposes raw loaded data as `self.data` and supports optional context count in `generate_batch`.

```APIDOC
## TemperatureGenerator.__init__

### Description
Initializes a German climate (temperature/elevation) data generator. Exposes the raw loaded data as `self.data` (a `_TemperatureData`). Supports `generate_batch(nc=None)` (optional context count) and a custom `shuffle()`/`epoch()`.

### Method
Constructor

### Signature
```python
class TemperatureGenerator(DataGenerator):
    def __init__(
        self,
        dtype,
        seed=0,
        batch_size=16,
        cv_stations=True,
        context_sample=False,
        context_sample_factor=10,
        context_sample_factor_at=None,
        context_elev_hr=True,
        target_min=5,
        target_max=None,
        target_dropnan=False,
        target_square=0.0,
        target_elev=False,
        target_elev_interpolate=False,
        subset="train",
        passes=1,
        device="cpu",
        data_task="germany",
        data_fold=5,
        data_path="climate_data",
    ):
```

### Parameters
- **dtype** - Required - Data type.
- **seed** (int) - Optional, default `0` - Seed.
- **batch_size** (int) - Optional, default `16` - Batch size.
- **cv_stations** (bool) - Optional, default `True` - Whether to use cross-validation stations.
- **context_sample** (bool) - Optional, default `False` - Whether to sample context.
- **context_sample_factor** (int) - Optional, default `10` - Context sample factor.
- **context_sample_factor_at** - Optional, default `None` - Context sample factor at.
- **context_elev_hr** (bool) - Optional, default `True` - Whether to use elevation in context.
- **target_min** (int) - Optional, default `5` - Minimum target value.
- **target_max** - Optional, default `None` - Maximum target value.
- **target_dropnan** (bool) - Optional, default `False` - Whether to drop NaN targets.
- **target_square** (float) - Optional, default `0.0` - Target square.
- **target_elev** (bool) - Optional, default `False` - Whether to use elevation as target.
- **target_elev_interpolate** (bool) - Optional, default `False` - Whether to interpolate elevation target.
- **subset** (str) - Optional, default `"train"` - Subset.
- **passes** (int) - Optional, default `1` - Number of passes.
- **device** (str) - Optional, default `"cpu"` - Device to use.
- **data_task** (str) - Optional, default `"germany"` - Data task.
- **data_fold** (int) - Optional, default `5` - Data fold.
- **data_path** (str) - Optional, default `"climate_data"` - Path to climate data.

### Notes
- German climate (temperature/elevation) data generator.
- Exposes the raw loaded data as `self.data` (a `_TemperatureData`).
- Supports `generate_batch(nc=None)` (optional context count) and a custom `shuffle()`/`epoch()`.
```

--------------------------------

### Catch built-in errors in consumer code

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/errors.md

Shows how to catch built-in exceptions raised by the neuralprocesses library. Includes examples for root-coding guard, likelihood spec validation, ordering validation in autoregressive prediction, and masked/aggregate misuse.

```python
import neuralprocesses.torch as nps
import torch

model = nps.construct_convgnp(dim_x=1, dim_y=1, likelihood="het")

# Root-coding guard:
try:
    pred = nps.code(model.encoder, xc, yc, xt)   # forgot root=True
except RuntimeError as e:
    assert "root = True" in str(e)

# Likelihood spec validation:
try:
    nps.construct_gnp(likelihood="not-a-likelihood")
except ValueError as e:
    assert "Incorrect likelihood specification" in str(e)

# Ordering validation in autoregressive prediction:
try:
    nps.ar_predict(model, contexts, xt, order="diagonal")
except RuntimeError as e:
    assert "Invalid ordering" in str(e)

# Masked/aggregate misuse:
try:
    nps.merge_contexts((xc1, yc1), (xc2, yc2))   # different n per axis
except (ValueError, AssertionError) as e:
    pass
```

--------------------------------

### DataGenerator.__init__

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/data.md

Initializes a DataGenerator with the given parameters.

```APIDOC
## DataGenerator.__init__

### Description
Initializes a DataGenerator with the given parameters.

### Method
__init__

### Signature
```python
def __init__(self, dtype, seed, num_tasks, batch_size, device):
```

### Parameters
- **dtype** (dtype) - Required - Data type of the generated data.
- **seed** (int) - Required - Random seed (creates `self.state` on `device`).
- **num_tasks** (int) - Required - Number of tasks per epoch.
- **batch_size** (int) - Required - Number of tasks per batch.
- **device** (str) - Required - Device (e.g. `"cpu"`).

### Attributes
- `dtype` - Data type.
- `float64` - 64-bit float variant.
- `int64` - 64-bit integer variant.
- `device` - Device.
- `state` - Random state.
- `batch_size` - Batch size.
- `num_batches` - `ceil(num_tasks / batch_size)`.

```

--------------------------------

### construct_convgnp — ConvGNP options

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/configuration.md

Constructor options for the ConvGNP model.

```APIDOC
## construct_convgnp — ConvGNP options

### Description
Constructor options for the ConvGNP model.

### Parameters
- **dim_aux_t** (int) - Optional - Dimensionality of target-specific auxiliary variables. Default: `None`.
- **points_per_unit** (float) - Optional - Density of the internal discretisation. Default: `64`.
- **margin** (float) - Optional - Margin of the internal discretisation. Default: `0.1`.
- **conv_arch** (str) - Optional - `"unet[-res][-sep]"` or `"conv[-res][-sep]"`. Default: `"unet"`.
- **unet_channels** (tuple[int]) - Optional - Channels per UNet layer. Default: `(64,) * 6`.
- **unet_kernels** (int or tuple[int]) - Optional - UNet kernel sizes. Default: `5`.
- **unet_strides** (int or tuple[int]) - Optional - UNet strides. Default: `2`.
- **unet_activations** (object or tuple) - Optional - UNet activations. Default: `None` (ReLU).
- **unet_resize_convs** (bool) - Optional - Resize convolutions instead of transposed. Default: `False`.
- **unet_resize_conv_interp_method** (str) - Optional - Interpolation method (`"bilinear"`). Default: `"nearest"`.
- **conv_receptive_field** (float) - Optional - Receptive field for `"conv"` architectures (required then). Default: `None`.
- **conv_layers** (int) - Optional - Layers of the standard architecture. Default: `6`.
- **conv_channels** (int) - Optional - Channels of the standard architecture. Default: `64`.
- **num_basis_functions** (int) - Optional - Low-rank likelihood rank. Default: `64`.
- **encoder_scales** (float or tuple[float]) - Optional - Initial encoder SetConv length scales. Default: `None` (→ `1/points_per_unit`).
- **encoder_scales_learnable** (bool) - Optional - Encoder length scales learnable. Default: `True`.
- **decoder_scale** (float) - Optional - Initial decoder SetConv length scale. Default: `None` (→ `1/points_per_unit`).
- **decoder_scale_learnable** (bool) - Optional - Decoder length scale learnable. Default: `True`.
- **aux_t_mlp_layers** (tuple[int]) - Optional - Widths of the auxiliary MLP. Default: `(128, 128, 128)`.
- **divide_by_density** (bool) - Optional - Divide by the density channel. Default: `True`.
- **epsilon** (float) - Optional - Addend before density division. Default: `1e-4`.
```

--------------------------------

### MultiOutputNormal.__init__

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/distributions.md

Direct constructor for MultiOutputNormal, taking three AbstractMatrix arguments for mean, variance, and noise, along with the data shape.

```APIDOC
## MultiOutputNormal.__init__

### Description
Direct constructor for MultiOutputNormal, taking three AbstractMatrix arguments for mean, variance, and noise, along with the data shape.

### Method
constructor

### Signature
```python
def __init__(self, mean: AbstractMatrix, var: AbstractMatrix, noise: AbstractMatrix, shape): ...
```

### Parameters
- **mean** (AbstractMatrix) - Required - Mean matrix.
- **var** (AbstractMatrix) - Required - Variance matrix.
- **noise** (AbstractMatrix) - Required - Noise matrix.
- **shape** - Required - Shape of the data before vectorising.
```

--------------------------------

### Define model and dataset configurations

Source: https://github.com/wesselb/neuralprocesses/blob/main/tables.ipynb

Defines a list of model configurations with names, display names, and loss functions, and a list of datasets to benchmark.

```python
models = [
    ("cnp", "CNP", "loglik"),
    ("acnp", "ACNP", "loglik"),
    ("convcnp", "ConvCNP", "loglik"),
    ("gnp", "GNP", "loglik"),
    ("agnp", "AGNP", "loglik"),
    ("convgnp", "ConvGNP", "loglik"),
    ("fullconvgnp", "FullConvGNP", "loglik"),
    ("np", "NP", "loglik"),
    ("anp", "ANP", "loglik"),
    ("convnp", "ConvNP", "loglik"),
    ("np", "NP", "elbo"),
    ("anp", "ANP", "elbo"),
    ("convnp", "ConvNP", "elbo"),
]

datasets = [
    "sawtooth",
    # "mixture",
]
```

--------------------------------

### construct_gnp — GNP / AGNP options

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/configuration.md

Constructor options for the GNP and AGNP models.

```APIDOC
## construct_gnp — GNP / AGNP options

### Description
Constructor options for the GNP and AGNP models.

### Parameters
- **dim_embedding** (int) - Optional - Deterministic embedding dimensionality. Default: `256`.
- **attention** (bool) - Optional - Attention encoder (AGNP: always `True`). Default: `False`.
- **attention_num_heads** (int) - Optional - Number of heads (`num_heads` on `construct_agnp`). Default: `8`.
- **num_enc_layers** (int) - Optional - Encoder layers. Default: `3`.
- **enc_same** (bool) - Optional - Share the encoder across context sets (`ValueError` if `dim_yc` is heterogeneous). Default: `False`.
- **num_dec_layers** (int) - Optional - Decoder layers. Default: `6`.
- **width** (int) - Optional - Width of intermediate MLPs. Default: `512`.
- **nonlinearity** (Callable or str) - Optional - `"ReLU"` or `"LeakyReLU"` (case-insensitive). Default: `"ReLU"`.
- **num_basis_functions** (int) - Optional - Rank of the low-rank likelihood. Default: `512`.
```

--------------------------------

### SpikesSlab.__init__

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/distributions.md

Constructs a spike-and-slab distribution with spike locations, a slab distribution, log-probabilities, and data dimensionality. The log-probabilities are normalized internally via logsumexp.

```APIDOC
## SpikesSlab.__init__

### Description
Constructs a spike-and-slab distribution with spike locations, a slab distribution, log-probabilities, and data dimensionality. The log-probabilities are normalized internally via logsumexp.

### Method
Constructor

### Signature
```python
@parametric
class SpikesSlab(AbstractDistribution):
    @_dispatch
    def __init__(
        self,
        spikes,
        slab: AbstractDistribution,
        logprobs,
        d,
        *,
        epsilon: float = 1e-6,
    ):
```

### Parameters
- **spikes** (vector) - Required - Spike locations (values).
- **slab** (AbstractDistribution) - Required - Slab distribution (e.g. `Beta`, `Gamma`).
- **logprobs** (tensor) - Required - Log-probabilities for the spikes and slab (slab last); normalised in `__init__` via `logsumexp`.
- **d** (int) - Required - Data dimensionality.
- **epsilon** (float) - Optional, default `1e-6` - Tolerance for equality checking (matching spike values).
```

--------------------------------

### MixtureGenerator.__init__

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/data.md

Initializes a mixture generator that combines multiple data generators. Each call to generate_batch picks a uniform random component.

```APIDOC
## MixtureGenerator.__init__

### Description
Initializes a mixture generator that combines multiple data generators. Each call to generate_batch picks a uniform random component.

### Method
Constructor

### Signature
```python
class MixtureGenerator(AbstractGenerator):
    @_dispatch
    def __init__(self, *gens: DataGenerator, seed=0):
```

### Parameters
- **gens** (DataGenerator) - Required - Components of the mixture; all must have equal `num_batches` (else `ValueError`).
- **seed** (int) - Optional, default `0` - Seed.

### Notes
- `generate_batch` picks a uniform random component per call.
```

--------------------------------

### UNet.__init__

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/coders-nn.md

Initializes a UNet instance with specified dimensions, channels, kernels, strides, activations, and other options. Raises ValueError if kernel, stride, or activation lengths do not match the number of channels.

```APIDOC
## UNet.__init__

### Description
Initializes a UNet instance with specified dimensions, channels, kernels, strides, activations, and other options. Raises ValueError if kernel, stride, or activation lengths do not match the number of channels.

### Method
Constructor

### Signature
```python
UNet(
    dim: int,
    in_channels: int,
    out_channels: int,
    channels: Tuple[int, ...] = (8, 16, 16, 32, 32, 64),
    kernels: Union[int, Tuple[Union[int, Tuple[int, ...]], ...]] = 5,
    strides: Union[int, Tuple[int, ...]] = 2,
    activations: Union[None, object, Tuple[object, ...]] = None,
    separable: bool = False,
    residual: bool = False,
    resize_convs: bool = False,
    resize_conv_interp_method: str = "nearest",
    dtype=None,
)
```

### Parameters
- **dim** (int) - Required - Dimensionality.
- **in_channels** (int) - Required - Input channels.
- **out_channels** (int) - Required - Output channels.
- **channels** (tuple[int]) - Optional - Channels of every layer. Default: `(8, 16, 16, 32, 32, 64)`.
- **kernels** (int or tuple) - Optional - Kernel sizes; int repeats per layer. Default: `5`.
- **strides** (int or tuple) - Optional - Strides; int repeats per layer. Default: `2`.
- **activations** (object or tuple) - Optional - Activations; default ReLU. Default: `None`.
- **separable** (bool) - Optional - Depthwise separable convolutions. Default: `False`.
- **residual** (bool) - Optional - Residual convolutional blocks. Default: `False`.
- **resize_convs** (bool) - Optional - Use resize convolutions instead of transposed ones. Default: `False`.
- **resize_conv_interp_method** (str) - Optional - Interpolation method for resize convolutions. Default: `"nearest"`.
- **dtype** (dtype) - Optional - Data type. Default: `None`.

### Errors
Raises `ValueError` if `kernels`, `strides`, or `activations` have length ≠ `len(channels)`.

### Attributes
- `dim` - Dimensionality.
- `kernels` - Kernel sizes.
- `strides` - Strides.
- `activations` - Activations.
- `num_halving_layers` - Equal to `len(channels)`.
- `receptive_fields` - List of receptive fields.
- `receptive_field` - Final receptive field.
- `before_turn_layers` - Layers before the turn.
- `after_turn_layers` - Layers after the turn.
- `final_linear` - 1×1 conv to `out_channels`.

### Notes
Downsampling layers use strided convolution or average pooling depending on parity of the receptive field; upsampling uses transposed convolutions or `UpSampling + conv`.
```

--------------------------------

### BernoulliGammaLikelihood.__init__

Source: https://github.com/wesselb/neuralprocesses/blob/main/_autodocs/api-reference/likelihoods.md

Constructor for the BernoulliGammaLikelihood class. Initializes the likelihood with a tolerance parameter for equality checking. The likelihood models a mixture of a Gamma slab and a Dirac spike at 0.

```APIDOC
## BernoulliGammaLikelihood.__init__

### Description
Constructor for the BernoulliGammaLikelihood class. Initializes the likelihood with a tolerance parameter for equality checking. The likelihood models a mixture of a Gamma slab and a Dirac spike at 0.

### Method
Constructor

### Signature
```python
@register_module
class BernoulliGammaLikelihood(AbstractLikelihood):
    @_dispatch
    def __init__(self, epsilon: float = 1e-6):
```

### Parameters
- **epsilon** (float) - Optional, default `1e-6` - Tolerance for equality checking.

### Channel Layout
`4 * dim_y` channels = `(k, scale, logp0, logps)` blocks. Parameters: `k = 1e-3 + softplus(z_k)`, `scale = 1e-3 + softplus(z_scale)`. Returns `SpikesSlab(spikes=(0,), Gamma(k, scale, d), stack(logp0, logps), d, epsilon=...)` — mixture of a Gamma slab and a Dirac spike at 0. Also uses `d + 1` data dims.

### Usage Example
```python
import neuralprocesses.torch as nps

# Low-rank Gaussian likelihood with rank-64 factor, inside a decoder chain:
num_channels, selector, lik = nps.construct_likelihood(
    nps, spec="lowrank", dim_y=2, num_basis_functions=64, dtype=None
)

decoder = nps.Chain(
    nps.MLP(in_dim=3, layers=(128,), out_dim=num_channels),
    selector,
    lik,
)
```
```
Tokens:
Raw