From bb3e3e36bf204837c95a50a8fee77627d0aaf8f8 Mon Sep 17 00:00:00 2001
From: Jan <jan.boelts@mailbox.org>
Date: Mon, 26 Oct 2026 22:30:44 +0100
Subject: [PATCH 01/12] docs: fix CNNEmbedding argument names in docstring and
 errors

The docstring and two assertion messages named num_cnn_layers and pool_size, which are not arguments of CNNEmbedding; the real ones are num_conv_layers and pool_kernel_size. The docstring also said the first two layers are convolutional and that pooling uses MaxPool1d, which is wrong for num_conv_layers != 2 and for 2D inputs.

The messages now name the real arguments, state the default out_channels_per_layer=[6, 12] that only fits two layers, and no longer carry the line breaks and indentation of triple-quoted strings. A test checks that each message names the argument that fixes the error.
---
 sbi/neural_nets/embedding_nets/cnn.py | 54 +++++++++++++--------------
 tests/embedding_net_test.py           | 13 +++++++
 2 files changed, 40 insertions(+), 27 deletions(-)

diff --git a/sbi/neural_nets/embedding_nets/cnn.py b/sbi/neural_nets/embedding_nets/cnn.py
index 7ce5e9eb6..246f2e950 100644
--- a/sbi/neural_nets/embedding_nets/cnn.py
+++ b/sbi/neural_nets/embedding_nets/cnn.py
@@ -105,36 +105,35 @@ def __init__(
     ):
         """Convolutional embedding network.
 
-        First two layers are convolutional, followed by fully connected layers.
-
-        Automatically infers whether to apply 1D or 2D convolution depending on
-        input_shape.
-        Allows usage of multiple (color) channels by passing in_channels > 1.
+        Applies `num_conv_layers` blocks of convolution, ReLU and max pooling,
+        followed by fully connected layers. A 1D `input_shape` uses `Conv1d` and
+        `MaxPool1d`, a 2D `input_shape` uses `Conv2d` and `MaxPool2d`. Each
+        convolution (stride 1, padding 1) turns a spatial size `n` into
+        `n + 3 - kernel_size`, and each pooling step then divides it by
+        `pool_kernel_size`, rounding down.
 
         Args:
             input_shape: Spatial input shape without batch or channel dimensions, e.g.,
                 (28,) for 1D or (28, 28) for 2D. Inputs may be flat or use
                 channel-first layout.
-            in_channels: Number of image channels, default 1.
-            out_channels_per_layer: Number of out convolutional out_channels for each
-                layer. Must match the number of layers passed below.
-            num_cnn_layers: Number of convolutional layers.
-            num_linear_layers: Number fully connected layer.
-            num_linear_units: Number of hidden units in fully-connected layers.
+            in_channels: Number of input channels, e.g., 3 for RGB images.
+            out_channels_per_layer: Number of output channels of each convolutional
+                layer, one entry per layer. Defaults to `[6, 12]`, which only fits
+                `num_conv_layers=2`.
+            num_conv_layers: Number of convolutional layers.
+            num_linear_layers: Number of fully connected layers.
+            num_linear_units: Number of hidden units in the fully connected layers.
             output_dim: Number of output units of the final layer.
-            kernel_size: Kernel size for both convolutional layers.
-            pool_size: pool size for MaxPool1d operation after the convolutional
-                layers.
+            kernel_size: Kernel size of every convolutional layer.
+            pool_kernel_size: Kernel size of the max pooling after every
+                convolutional layer.
         """
         super(CNNEmbedding, self).__init__()
 
-        assert isinstance(input_shape, Tuple), (
-            "input_shape must be a Tuple of size 1 or 2, e.g., (width, [height])."
+        assert isinstance(input_shape, Tuple) and 0 < len(input_shape) < 3, (
+            "input_shape must be a tuple of length 1 or 2, e.g., (28,) or (28, 28). "
+            "Pass the number of channels as in_channels."
         )
-        assert (
-            0 < len(input_shape) < 3
-        ), """input_shape must be a Tuple of size 1 or 2, e.g.,
-            (width, [height]). Number of input channels are passed separately"""
 
         use_2d_cnn = len(input_shape) == 2
         conv_module = nn.Conv2d if use_2d_cnn else nn.Conv1d
@@ -143,7 +142,9 @@ def __init__(
         if out_channels_per_layer is None:
             out_channels_per_layer = [6, 12]
         assert len(out_channels_per_layer) == num_conv_layers, (
-            "out_channels needs as many entries as num_cnn_layers."
+            f"out_channels_per_layer has {len(out_channels_per_layer)} entries, but "
+            f"num_conv_layers is {num_conv_layers}. Pass one entry per convolutional "
+            "layer; the default [6, 12] only fits num_conv_layers=2."
         )
 
         # define input shape with channel
@@ -156,7 +157,6 @@ def __init__(
         padding = 1
 
         for ii in range(num_conv_layers):
-            # Defining another 2D convolution layer
             conv_layer = conv_module(
                 in_channels=in_channels if ii == 0 else out_channels_per_layer[ii - 1],
                 out_channels=out_channels_per_layer[ii],
@@ -169,11 +169,11 @@ def __init__(
             # Calculate change of output size of each CNN layer
             cnn_output_size = get_new_cnn_output_size(cnn_output_size, conv_layer, pool)
 
-            assert all(
-                cnn_output_size
-            ), f"""CNN output size is zero at layer {ii + 1}. Either reduce
-                 num_cnn_layers to {ii} or adjust the kernel_size
-                 and pool_kernel_size accordingly."""
+            assert all(cnn_output_size), (
+                f"CNN output size is zero at layer {ii + 1}. Use fewer "
+                "convolutional layers (num_conv_layers) or a smaller kernel_size "
+                "or pool_kernel_size."
+            )
 
         self.cnn_subnet = nn.Sequential(*cnn_layers)
 
diff --git a/tests/embedding_net_test.py b/tests/embedding_net_test.py
index 7aaaea1b5..7dd0bcf9c 100644
--- a/tests/embedding_net_test.py
+++ b/tests/embedding_net_test.py
@@ -200,6 +200,19 @@ def test_cnn_embedding_input_layout(embedding_cls, input_shape, kwargs):
         embedding_net(channel_last_input)
 
 
+@pytest.mark.parametrize(
+    ("input_shape", "kwargs", "match"),
+    [
+        ((32, 32, 3), {}, "in_channels"),
+        ((32,), {"num_conv_layers": 3}, "num_conv_layers is 3"),
+    ],
+)
+def test_cnn_embedding_errors_name_arguments(input_shape, kwargs, match):
+    """CNNEmbedding errors name the constructor arguments that fix them."""
+    with pytest.raises(AssertionError, match=match):
+        CNNEmbedding(input_shape, **kwargs)
+
+
 @pytest.mark.parametrize("input_shape", [(2,), (128,)])
 @pytest.mark.parametrize("num_layers", [2, 4])
 @pytest.mark.parametrize("num_hiddens", [32, 64])

From 050a5be9169721fbdcdcb82f961a5fb8b4451f14 Mon Sep 17 00:00:00 2001
From: Jan <jan.boelts@mailbox.org>
Date: Mon, 26 Oct 2026 22:32:38 +0100
Subject: [PATCH 02/12] fix: CNNEmbedding rejects layers that shrink the input
 below size one

The size guard in CNNEmbedding only caught an output size of exactly zero. A kernel much larger than the padded input can give a negative size, which passed the guard. For 2D inputs two negative sizes multiply to a positive linear input size, so the network was built without error and failed only in the first forward pass, with a torch error instead of the guard's message.

The guard now requires every size to be positive and reports the computed size.
---
 sbi/neural_nets/embedding_nets/cnn.py | 8 ++++----
 tests/embedding_net_test.py           | 3 +++
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/sbi/neural_nets/embedding_nets/cnn.py b/sbi/neural_nets/embedding_nets/cnn.py
index 246f2e950..a88ecc81f 100644
--- a/sbi/neural_nets/embedding_nets/cnn.py
+++ b/sbi/neural_nets/embedding_nets/cnn.py
@@ -169,10 +169,10 @@ def __init__(
             # Calculate change of output size of each CNN layer
             cnn_output_size = get_new_cnn_output_size(cnn_output_size, conv_layer, pool)
 
-            assert all(cnn_output_size), (
-                f"CNN output size is zero at layer {ii + 1}. Use fewer "
-                "convolutional layers (num_conv_layers) or a smaller kernel_size "
-                "or pool_kernel_size."
+            assert all(size > 0 for size in cnn_output_size), (
+                f"CNN output size {cnn_output_size} is not positive at layer "
+                f"{ii + 1}. Use fewer convolutional layers (num_conv_layers) or a "
+                "smaller kernel_size or pool_kernel_size."
             )
 
         self.cnn_subnet = nn.Sequential(*cnn_layers)
diff --git a/tests/embedding_net_test.py b/tests/embedding_net_test.py
index 7dd0bcf9c..8db668542 100644
--- a/tests/embedding_net_test.py
+++ b/tests/embedding_net_test.py
@@ -205,6 +205,9 @@ def test_cnn_embedding_input_layout(embedding_cls, input_shape, kwargs):
     [
         ((32, 32, 3), {}, "in_channels"),
         ((32,), {"num_conv_layers": 3}, "num_conv_layers is 3"),
+        # Kernel larger than the padded input: both sizes are negative, and
+        # their product is positive.
+        ((4, 4), {"kernel_size": 9}, "num_conv_layers"),
     ],
 )
 def test_cnn_embedding_errors_name_arguments(input_shape, kwargs, match):

From 1d837d663256833ae7f200efab99dc6e056d7b3b Mon Sep 17 00:00:00 2001
From: Jan <jan.boelts@mailbox.org>
Date: Mon, 26 Oct 2026 22:32:52 +0100
Subject: [PATCH 03/12] docs: fix typos in CausalCNNEmbedding docstring and
 errors

---
 sbi/neural_nets/embedding_nets/causal_cnn.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/sbi/neural_nets/embedding_nets/causal_cnn.py b/sbi/neural_nets/embedding_nets/causal_cnn.py
index e9827a0e1..045a630bf 100644
--- a/sbi/neural_nets/embedding_nets/causal_cnn.py
+++ b/sbi/neural_nets/embedding_nets/causal_cnn.py
@@ -152,10 +152,10 @@ def __init__(
         aggregator: Optional[nn.Module] = None,
         output_dim: int = 20,
     ):
-        """Intitialize embedding network that uses 1D causal convolutions.
+        """Initialize embedding network that uses 1D causal convolutions.
 
         This is a simplified version of the architecture introduced for
-        the speech recognition task in the WaveNet paper (van den Oord, et al. (2016))
+        the speech recognition task in the WaveNet paper (van den Oord, et al. (2016)).
 
         After several dilated causal convolutions (that maintain the dimensionality
         of the input), an aggregator network is used to bring down the dimensionality.
@@ -163,7 +163,7 @@ def __init__(
         If you do not provide an aggregator network yourself, a default aggregator
         is used. This default aggregator is based on the WaveNet paper's description
         of their Speech Recognition Task, and uses non-causal convolutions and pooling
-        layers, and global average poolingg to obtain a final low dimensional embedding.
+        layers, and global average pooling to obtain a final low dimensional embedding.
 
         Args:
             input_shape: Spatial input shape without batch or channel dimensions,
@@ -216,8 +216,8 @@ def __init__(
                     dilation_per_layer = [1] * num_conv_layers
                 case _:
                     raise ValueError(
-                        f"{dilation} is not a valid option, please use \"none\","
-                        "\"exponential\",or \"exponential_cyclic\", or pass a list "
+                        f"{dilation} is not a valid option, please use \"none\", "
+                        "\"exponential\", or \"exponential_cyclic\", or pass a list "
                         "of dilation sizes."
                     )
         else:
@@ -227,8 +227,8 @@ def __init__(
             dilation_per_layer = dilation
 
         assert max(dilation_per_layer) < total_timepoints, (
-            "Your maximal dilations size used is larger than the number of "
-            "timepoints in your input, please provide a list with smaller dilations."
+            "The largest dilation is not smaller than the number of timepoints "
+            "in your input. Please provide a list with smaller dilations."
         )
         if out_channels_per_layer is None:
             out_channels_per_layer = [16] * num_conv_layers
@@ -256,7 +256,7 @@ def __init__(
                 int(total_timepoints / pool_kernel_size),
             )
             assert aggregator_out_shape[-1] > 1, (
-                "Your dimensionality is already small,"
+                "Your dimensionality is already small. "
                 "Please ensure a larger input size or use a custom aggregator."
             )
             aggregator = WaveNetSRLikeAggregator(

From bdafc1f0703697aa1de437a2f4f712f4fa04e271 Mon Sep 17 00:00:00 2001
From: Jan <jan.boelts@mailbox.org>
Date: Thu, 29 Oct 2026 06:55:11 +0100
Subject: [PATCH 04/12] fix: CNNEmbedding default channels follow
 num_conv_layers

The default out_channels_per_layer was [6, 12], so setting num_conv_layers to anything but 2 without also passing out_channels_per_layer always failed. The default now starts at 6 channels and doubles per layer. For the default two layers it is still [6, 12], so existing networks do not change.
---
 sbi/neural_nets/embedding_nets/cnn.py |  9 +++++----
 tests/embedding_net_test.py           | 15 ++++++++++++++-
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/sbi/neural_nets/embedding_nets/cnn.py b/sbi/neural_nets/embedding_nets/cnn.py
index a88ecc81f..0c726e7b0 100644
--- a/sbi/neural_nets/embedding_nets/cnn.py
+++ b/sbi/neural_nets/embedding_nets/cnn.py
@@ -118,8 +118,9 @@ def __init__(
                 channel-first layout.
             in_channels: Number of input channels, e.g., 3 for RGB images.
             out_channels_per_layer: Number of output channels of each convolutional
-                layer, one entry per layer. Defaults to `[6, 12]`, which only fits
-                `num_conv_layers=2`.
+                layer, one entry per layer. Defaults to 6 channels in the first
+                layer, doubling in each further layer, e.g., `[6, 12]` for two
+                layers.
             num_conv_layers: Number of convolutional layers.
             num_linear_layers: Number of fully connected layers.
             num_linear_units: Number of hidden units in the fully connected layers.
@@ -140,11 +141,11 @@ def __init__(
         pool_module = nn.MaxPool2d if use_2d_cnn else nn.MaxPool1d
 
         if out_channels_per_layer is None:
-            out_channels_per_layer = [6, 12]
+            out_channels_per_layer = [6 * 2**ii for ii in range(num_conv_layers)]
         assert len(out_channels_per_layer) == num_conv_layers, (
             f"out_channels_per_layer has {len(out_channels_per_layer)} entries, but "
             f"num_conv_layers is {num_conv_layers}. Pass one entry per convolutional "
-            "layer; the default [6, 12] only fits num_conv_layers=2."
+            "layer."
         )
 
         # define input shape with channel
diff --git a/tests/embedding_net_test.py b/tests/embedding_net_test.py
index 8db668542..7082819b6 100644
--- a/tests/embedding_net_test.py
+++ b/tests/embedding_net_test.py
@@ -204,7 +204,11 @@ def test_cnn_embedding_input_layout(embedding_cls, input_shape, kwargs):
     ("input_shape", "kwargs", "match"),
     [
         ((32, 32, 3), {}, "in_channels"),
-        ((32,), {"num_conv_layers": 3}, "num_conv_layers is 3"),
+        (
+            (32,),
+            {"num_conv_layers": 3, "out_channels_per_layer": [6, 12]},
+            "num_conv_layers is 3",
+        ),
         # Kernel larger than the padded input: both sizes are negative, and
         # their product is positive.
         ((4, 4), {"kernel_size": 9}, "num_conv_layers"),
@@ -216,6 +220,15 @@ def test_cnn_embedding_errors_name_arguments(input_shape, kwargs, match):
         CNNEmbedding(input_shape, **kwargs)
 
 
+def test_cnn_embedding_default_channels_follow_num_conv_layers():
+    """The default channels double per layer, starting at 6."""
+    embedding_net = CNNEmbedding((32,), num_conv_layers=3)
+    convs = [m for m in embedding_net.cnn_subnet if isinstance(m, nn.Conv1d)]
+
+    assert [conv.out_channels for conv in convs] == [6, 12, 24]
+    assert embedding_net(torch.randn(4, 32)).shape == (4, 20)
+
+
 @pytest.mark.parametrize("input_shape", [(2,), (128,)])
 @pytest.mark.parametrize("num_layers", [2, 4])
 @pytest.mark.parametrize("num_hiddens", [32, 64])

From dedd4737117976324728c1aee045f4286fcf03aa Mon Sep 17 00:00:00 2001
From: Jan <jan.boelts@mailbox.org>
Date: Thu, 29 Oct 2026 06:56:55 +0100
Subject: [PATCH 05/12] fix: embedding nets raise ValueError or TypeError on
 invalid input

Embedding nets checked user arguments with assert, and Python drops asserts under -O. The checks on constructor arguments of CNNEmbedding, CausalCNNEmbedding and PermutationInvariantEmbedding, and on user input in the forward passes of SpectralConvEmbedding and LRU, now raise TypeError for a wrong type and ValueError for a wrong value, as the other embedding nets already do.

CNNEmbedding also rejects num_conv_layers < 1. With the default channels that follow num_conv_layers, zero layers gave an empty channel list and an IndexError. Its size error suggests fewer layers only from layer 2 on, since at layer 1 only a smaller kernel_size or pool_kernel_size helps.

Asserts on values that the classes compute themselves stay asserts.
---
 .../embedding_nets/SC_embedding.py            |  5 +--
 sbi/neural_nets/embedding_nets/causal_cnn.py  | 43 +++++++++++--------
 sbi/neural_nets/embedding_nets/cnn.py         | 43 +++++++++++++------
 sbi/neural_nets/embedding_nets/lru.py         |  9 ++--
 .../embedding_nets/permutation_invariant.py   |  6 +--
 tests/embedding_net_test.py                   |  8 ++--
 6 files changed, 67 insertions(+), 47 deletions(-)

diff --git a/sbi/neural_nets/embedding_nets/SC_embedding.py b/sbi/neural_nets/embedding_nets/SC_embedding.py
index 8829ec7c3..837f26ae5 100644
--- a/sbi/neural_nets/embedding_nets/SC_embedding.py
+++ b/sbi/neural_nets/embedding_nets/SC_embedding.py
@@ -325,9 +325,8 @@ def forward(self, x: Tensor) -> Tensor:
 
         n_points = x.shape[1]
 
-        assert self.modes <= n_points // 2 + 1, (
-            "Modes should be at most floor(n_points/2) + 1"
-        )
+        if self.modes > n_points // 2 + 1:
+            raise ValueError("Modes should be at most floor(n_points/2) + 1")
 
         x = self.fc0(x)  # (batch_size, n_points, in_channels)
 
diff --git a/sbi/neural_nets/embedding_nets/causal_cnn.py b/sbi/neural_nets/embedding_nets/causal_cnn.py
index 045a630bf..468fc48c5 100644
--- a/sbi/neural_nets/embedding_nets/causal_cnn.py
+++ b/sbi/neural_nets/embedding_nets/causal_cnn.py
@@ -191,17 +191,21 @@ def __init__(
         """
 
         super(CausalCNNEmbedding, self).__init__()
-        assert isinstance(input_shape, Tuple), (
-            "input_shape must be a Tuple of size 1, e.g. (timepoints,)."
-        )
-        assert len(input_shape) == 1, "Currently only 1D causal CNNs are supported."
+        if not isinstance(input_shape, tuple):
+            raise TypeError(
+                "input_shape must be a tuple of length 1, e.g. (timepoints,), but "
+                f"got {type(input_shape).__name__}."
+            )
+        if len(input_shape) != 1:
+            raise ValueError("Currently only 1D causal CNNs are supported.")
         self.input_shape = (in_channels, *input_shape)
 
         total_timepoints = input_shape[0]
-        assert total_timepoints >= pool_kernel_size, (
-            "Please ensure that the pool kernel size is not "
-            "larger than the number of observed timepoints."
-        )
+        if total_timepoints < pool_kernel_size:
+            raise ValueError(
+                "Please ensure that the pool kernel size is not "
+                "larger than the number of observed timepoints."
+            )
         if isinstance(dilation, str):
             match dilation.lower():
                 case "exponential_cyclic":
@@ -221,15 +225,15 @@ def __init__(
                         "of dilation sizes."
                     )
         else:
-            assert isinstance(dilation, List), (
-                "Please pass dilation size as list or a string option."
-            )
+            if not isinstance(dilation, list):
+                raise TypeError("Please pass dilation size as list or a string option.")
             dilation_per_layer = dilation
 
-        assert max(dilation_per_layer) < total_timepoints, (
-            "The largest dilation is not smaller than the number of timepoints "
-            "in your input. Please provide a list with smaller dilations."
-        )
+        if max(dilation_per_layer) >= total_timepoints:
+            raise ValueError(
+                "The largest dilation is not smaller than the number of timepoints "
+                "in your input. Please provide a list with smaller dilations."
+            )
         if out_channels_per_layer is None:
             out_channels_per_layer = [16] * num_conv_layers
 
@@ -255,10 +259,11 @@ def __init__(
                 out_channels_per_layer[-1],
                 int(total_timepoints / pool_kernel_size),
             )
-            assert aggregator_out_shape[-1] > 1, (
-                "Your dimensionality is already small. "
-                "Please ensure a larger input size or use a custom aggregator."
-            )
+            if aggregator_out_shape[-1] <= 1:
+                raise ValueError(
+                    "Your dimensionality is already small. "
+                    "Please ensure a larger input size or use a custom aggregator."
+                )
             aggregator = WaveNetSRLikeAggregator(
                 aggregator_out_shape[0],
                 aggregator_out_shape[-1],
diff --git a/sbi/neural_nets/embedding_nets/cnn.py b/sbi/neural_nets/embedding_nets/cnn.py
index 0c726e7b0..08bff4a2c 100644
--- a/sbi/neural_nets/embedding_nets/cnn.py
+++ b/sbi/neural_nets/embedding_nets/cnn.py
@@ -131,22 +131,33 @@ def __init__(
         """
         super(CNNEmbedding, self).__init__()
 
-        assert isinstance(input_shape, Tuple) and 0 < len(input_shape) < 3, (
-            "input_shape must be a tuple of length 1 or 2, e.g., (28,) or (28, 28). "
-            "Pass the number of channels as in_channels."
-        )
+        if not isinstance(input_shape, tuple):
+            raise TypeError(
+                "input_shape must be a tuple, e.g., (28,) or (28, 28), but got "
+                f"{type(input_shape).__name__}."
+            )
+        if not 0 < len(input_shape) < 3:
+            raise ValueError(
+                "input_shape must have length 1 or 2, e.g., (28,) or (28, 28). "
+                "Pass the number of channels as in_channels."
+            )
 
         use_2d_cnn = len(input_shape) == 2
         conv_module = nn.Conv2d if use_2d_cnn else nn.Conv1d
         pool_module = nn.MaxPool2d if use_2d_cnn else nn.MaxPool1d
 
+        if num_conv_layers < 1:
+            raise ValueError(
+                f"num_conv_layers must be at least 1, but got {num_conv_layers}."
+            )
         if out_channels_per_layer is None:
             out_channels_per_layer = [6 * 2**ii for ii in range(num_conv_layers)]
-        assert len(out_channels_per_layer) == num_conv_layers, (
-            f"out_channels_per_layer has {len(out_channels_per_layer)} entries, but "
-            f"num_conv_layers is {num_conv_layers}. Pass one entry per convolutional "
-            "layer."
-        )
+        if len(out_channels_per_layer) != num_conv_layers:
+            raise ValueError(
+                f"out_channels_per_layer has {len(out_channels_per_layer)} entries, "
+                f"but num_conv_layers is {num_conv_layers}. Pass one entry per "
+                "convolutional layer."
+            )
 
         # define input shape with channel
         self.input_shape = (in_channels, *input_shape)
@@ -170,11 +181,15 @@ def __init__(
             # Calculate change of output size of each CNN layer
             cnn_output_size = get_new_cnn_output_size(cnn_output_size, conv_layer, pool)
 
-            assert all(size > 0 for size in cnn_output_size), (
-                f"CNN output size {cnn_output_size} is not positive at layer "
-                f"{ii + 1}. Use fewer convolutional layers (num_conv_layers) or a "
-                "smaller kernel_size or pool_kernel_size."
-            )
+            if not all(size > 0 for size in cnn_output_size):
+                fewer_layers = (
+                    f", or set num_conv_layers to at most {ii}" if ii > 0 else ""
+                )
+                raise ValueError(
+                    f"CNN output size {cnn_output_size} is not positive at layer "
+                    f"{ii + 1}. Use a smaller kernel_size or pool_kernel_size"
+                    f"{fewer_layers}."
+                )
 
         self.cnn_subnet = nn.Sequential(*cnn_layers)
 
diff --git a/sbi/neural_nets/embedding_nets/lru.py b/sbi/neural_nets/embedding_nets/lru.py
index 848a3acbb..714997709 100644
--- a/sbi/neural_nets/embedding_nets/lru.py
+++ b/sbi/neural_nets/embedding_nets/lru.py
@@ -337,10 +337,11 @@ def forward(
             )
         else:
             state = state.to(device=input.device)
-            assert state.shape == expected_state_shape, (
-                f"Invalid state shape {state.shape}, "  # fmt: skip
-                f"expected {expected_state_shape}"  # fmt: skip
-            )
+            if state.shape != expected_state_shape:
+                raise ValueError(
+                    f"Invalid state shape {state.shape}, "
+                    f"expected {expected_state_shape}"
+                )
 
         # Detemine which mode to run the forward method with.
         mode = mode or self.mode
diff --git a/sbi/neural_nets/embedding_nets/permutation_invariant.py b/sbi/neural_nets/embedding_nets/permutation_invariant.py
index e55a1e422..02b45cf91 100644
--- a/sbi/neural_nets/embedding_nets/permutation_invariant.py
+++ b/sbi/neural_nets/embedding_nets/permutation_invariant.py
@@ -70,10 +70,8 @@ def __init__(
         super().__init__()
         self.trial_net = trial_net
         self.aggregation_dim = aggregation_dim
-        assert aggregation_fn in [
-            "mean",
-            "sum",
-        ], "aggregation_fn must be 'mean' or 'sum'."
+        if aggregation_fn not in ("mean", "sum"):
+            raise ValueError("aggregation_fn must be 'mean' or 'sum'.")
         self.aggregation_fn = aggregation_fn
 
         # construct fully connected layers
diff --git a/tests/embedding_net_test.py b/tests/embedding_net_test.py
index 7082819b6..51acad96e 100644
--- a/tests/embedding_net_test.py
+++ b/tests/embedding_net_test.py
@@ -204,19 +204,21 @@ def test_cnn_embedding_input_layout(embedding_cls, input_shape, kwargs):
     ("input_shape", "kwargs", "match"),
     [
         ((32, 32, 3), {}, "in_channels"),
+        ((32,), {"num_conv_layers": 0}, "num_conv_layers must be at least 1"),
         (
             (32,),
             {"num_conv_layers": 3, "out_channels_per_layer": [6, 12]},
             "num_conv_layers is 3",
         ),
-        # Kernel larger than the padded input: both sizes are negative, and
+        # Kernel much larger than the padded input: both sizes are negative, and
         # their product is positive.
-        ((4, 4), {"kernel_size": 9}, "num_conv_layers"),
+        ((4, 4), {"kernel_size": 9}, "smaller kernel_size"),
+        ((8,), {"pool_kernel_size": 4}, "num_conv_layers to at most 1"),
     ],
 )
 def test_cnn_embedding_errors_name_arguments(input_shape, kwargs, match):
     """CNNEmbedding errors name the constructor arguments that fix them."""
-    with pytest.raises(AssertionError, match=match):
+    with pytest.raises(ValueError, match=match):
         CNNEmbedding(input_shape, **kwargs)
 
 

From 8f196a60323e2af57f2df508960385bbffbd9bf5 Mon Sep 17 00:00:00 2001
From: Jan <jan.boelts@mailbox.org>
Date: Thu, 29 Oct 2026 07:49:54 +0100
Subject: [PATCH 06/12] fix: CausalCNNEmbedding checks its arguments like
 CNNEmbedding

A wrong-length out_channels_per_layer or dilation list, or num_conv_layers=0, failed with an IndexError or Python's 'max() arg is an empty sequence'. A dilation list longer than num_conv_layers was silently cut. These now raise a ValueError that names the argument, and the existing pool size and dilation errors name the arguments and values too.
---
 sbi/neural_nets/embedding_nets/causal_cnn.py | 30 ++++++++++++++++----
 tests/embedding_net_test.py                  | 16 +++++++++++
 2 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/sbi/neural_nets/embedding_nets/causal_cnn.py b/sbi/neural_nets/embedding_nets/causal_cnn.py
index 468fc48c5..27b1aa29c 100644
--- a/sbi/neural_nets/embedding_nets/causal_cnn.py
+++ b/sbi/neural_nets/embedding_nets/causal_cnn.py
@@ -203,8 +203,12 @@ def __init__(
         total_timepoints = input_shape[0]
         if total_timepoints < pool_kernel_size:
             raise ValueError(
-                "Please ensure that the pool kernel size is not "
-                "larger than the number of observed timepoints."
+                f"pool_kernel_size ({pool_kernel_size}) must not be larger than the "
+                f"number of timepoints ({total_timepoints})."
+            )
+        if num_conv_layers < 1:
+            raise ValueError(
+                f"num_conv_layers must be at least 1, but got {num_conv_layers}."
             )
         if isinstance(dilation, str):
             match dilation.lower():
@@ -227,15 +231,27 @@ def __init__(
         else:
             if not isinstance(dilation, list):
                 raise TypeError("Please pass dilation size as list or a string option.")
+            if len(dilation) != num_conv_layers:
+                raise ValueError(
+                    f"dilation has {len(dilation)} entries, but num_conv_layers is "
+                    f"{num_conv_layers}. Pass one dilation per convolutional layer."
+                )
             dilation_per_layer = dilation
 
         if max(dilation_per_layer) >= total_timepoints:
             raise ValueError(
-                "The largest dilation is not smaller than the number of timepoints "
-                "in your input. Please provide a list with smaller dilations."
+                f"The largest dilation ({max(dilation_per_layer)}) must be smaller "
+                f"than the number of timepoints ({total_timepoints}). Pass smaller "
+                "dilations or fewer num_conv_layers."
             )
         if out_channels_per_layer is None:
             out_channels_per_layer = [16] * num_conv_layers
+        if len(out_channels_per_layer) != num_conv_layers:
+            raise ValueError(
+                f"out_channels_per_layer has {len(out_channels_per_layer)} entries, "
+                f"but num_conv_layers is {num_conv_layers}. Pass one entry per "
+                "convolutional layer."
+            )
 
         causal_layers = []
         for ll in range(num_conv_layers):
@@ -261,8 +277,10 @@ def __init__(
             )
             if aggregator_out_shape[-1] <= 1:
                 raise ValueError(
-                    "Your dimensionality is already small. "
-                    "Please ensure a larger input size or use a custom aggregator."
+                    f"After pooling with pool_kernel_size={pool_kernel_size}, "
+                    f"{aggregator_out_shape[-1]} timepoint(s) remain, but the default "
+                    "aggregator needs at least 2. Use a smaller pool_kernel_size, a "
+                    "longer input, or pass your own aggregator."
                 )
             aggregator = WaveNetSRLikeAggregator(
                 aggregator_out_shape[0],
diff --git a/tests/embedding_net_test.py b/tests/embedding_net_test.py
index 51acad96e..78c4c22de 100644
--- a/tests/embedding_net_test.py
+++ b/tests/embedding_net_test.py
@@ -222,6 +222,22 @@ def test_cnn_embedding_errors_name_arguments(input_shape, kwargs, match):
         CNNEmbedding(input_shape, **kwargs)
 
 
+@pytest.mark.parametrize(
+    ("kwargs", "match"),
+    [
+        ({"num_conv_layers": 0}, "num_conv_layers must be at least 1"),
+        ({"out_channels_per_layer": [16, 16]}, "num_conv_layers is 5"),
+        ({"dilation": [1, 2]}, "dilation has 2 entries"),
+        ({"pool_kernel_size": 128}, "pool_kernel_size"),
+        ({"pool_kernel_size": 64}, "pool_kernel_size"),
+    ],
+)
+def test_causal_cnn_embedding_errors_name_arguments(kwargs, match):
+    """CausalCNNEmbedding errors name the constructor arguments that fix them."""
+    with pytest.raises(ValueError, match=match):
+        CausalCNNEmbedding((64,), **{"pool_kernel_size": 2, **kwargs})
+
+
 def test_cnn_embedding_default_channels_follow_num_conv_layers():
     """The default channels double per layer, starting at 6."""
     embedding_net = CNNEmbedding((32,), num_conv_layers=3)

From 9ecb2d1f009ce39196186094557f4f222b3cc413 Mon Sep 17 00:00:00 2001
From: Jan <jan.boelts@mailbox.org>
Date: Thu, 29 Oct 2026 07:51:06 +0100
Subject: [PATCH 07/12] fix: WaveNet aggregator no longer prints or changes its
 arguments

WaveNetSRLikeAggregator, the default aggregator of CausalCNNEmbedding, printed its intermediate shapes on every build, so the time-series how-to guide showed them in its output. It also appended output_dim to the caller's intermediate_channel_sizes list. Both are gone, the stored notebook output no longer shows the shapes, and its docstring names its arguments correctly.
---
 docs/how_to_guide/20_time_series_embedding.ipynb |  3 ---
 sbi/neural_nets/embedding_nets/causal_cnn.py     | 15 ++++++---------
 tests/embedding_net_test.py                      | 12 ++++++++++++
 3 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/docs/how_to_guide/20_time_series_embedding.ipynb b/docs/how_to_guide/20_time_series_embedding.ipynb
index 9d219d84a..b8a820e66 100644
--- a/docs/how_to_guide/20_time_series_embedding.ipynb
+++ b/docs/how_to_guide/20_time_series_embedding.ipynb
@@ -80,9 +80,6 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "(16, 10)\n",
-      "(64, 5)\n",
-      "(16, 2)\n",
       " Neural network successfully converged after 31 epochs."
      ]
     }
diff --git a/sbi/neural_nets/embedding_nets/causal_cnn.py b/sbi/neural_nets/embedding_nets/causal_cnn.py
index 27b1aa29c..856bf2747 100644
--- a/sbi/neural_nets/embedding_nets/causal_cnn.py
+++ b/sbi/neural_nets/embedding_nets/causal_cnn.py
@@ -56,8 +56,7 @@ def WaveNetSRLikeAggregator(
     intermediate_channel_sizes: Optional[List] = None,
     stride_sizes: Union[int, List] = 1,
 ) -> nn.Module:
-    """
-    Creates a non-causal 1D CNN aggregator based on the WaveNet speach recognition task
+    """Non-causal 1D CNN aggregator based on WaveNet's speech recognition setup.
 
     By default this function creates an aggregator with two CNN layers,
     after every convolution a maxpooling operation halves the number of timepoints.
@@ -68,14 +67,14 @@ def WaveNetSRLikeAggregator(
     Args:
         in_channels: number of channels at input.
         num_timepoints: length of the input.
-        output dim: wanted number of features as output.
+        output_dim: wanted number of features as output.
         activation: activation to apply after the convolution.
         kernel_sizes: (optional) alter the kernel size used and the number of CNN layers
             (through the length of the kernel size vector).
         intermediate_channel_sizes: (optional) alter the intermediate channel sizes
             used, should have length = len(kernel_sizes) - 1.
-        stride_sizes = Optional alter the stride used, either a vector of
-            len = len(kernel_sizes) or a single integer, in which case the same stride
+        stride_sizes: (optional) alter the stride used, either a list of length
+            len(kernel_sizes) or a single integer, in which case the same stride
             is used in every convolution.
 
     Returns:
@@ -97,8 +96,8 @@ def WaveNetSRLikeAggregator(
         "Provided kernel size list should be exactly one element longer "
         "than channel size list."
     )
-    intermediate_channel_sizes += [output_dim]
-    if isinstance(stride_sizes, List):
+    intermediate_channel_sizes = [*intermediate_channel_sizes, output_dim]
+    if isinstance(stride_sizes, list):
         assert len(stride_sizes) == len(kernel_sizes), (
             "Provided stride size list should be have the same size as"
             "the kernel size list."
@@ -108,7 +107,6 @@ def WaveNetSRLikeAggregator(
 
     non_causal_layers = []
     for ll in range(len(kernel_sizes)):
-        print(aggregator_out_shape)
         conv_layer = nn.Conv1d(
             in_channels=in_channels if ll == 0 else intermediate_channel_sizes[ll - 1],
             out_channels=intermediate_channel_sizes[ll],
@@ -131,7 +129,6 @@ def WaveNetSRLikeAggregator(
                 / 2
             ),
         )
-    print(aggregator_out_shape)
     aggregator = nn.Sequential(*non_causal_layers, nn.AdaptiveAvgPool1d(1))
     return aggregator
 
diff --git a/tests/embedding_net_test.py b/tests/embedding_net_test.py
index 78c4c22de..b118f5da3 100644
--- a/tests/embedding_net_test.py
+++ b/tests/embedding_net_test.py
@@ -30,6 +30,7 @@
     SpectralConvEmbedding,
     TransformerEmbedding,
 )
+from sbi.neural_nets.embedding_nets.causal_cnn import WaveNetSRLikeAggregator
 from sbi.neural_nets.embedding_nets.lru import LRU, LRUBlock
 from sbi.simulators.linear_gaussian import (
     linear_gaussian,
@@ -238,6 +239,17 @@ def test_causal_cnn_embedding_errors_name_arguments(kwargs, match):
         CausalCNNEmbedding((64,), **{"pool_kernel_size": 2, **kwargs})
 
 
+def test_wavenet_aggregator_keeps_arguments_and_prints_nothing(capsys):
+    """Building the aggregator neither changes the caller's list nor prints."""
+    intermediate_channel_sizes = [8]
+    WaveNetSRLikeAggregator(
+        4, 32, output_dim=5, intermediate_channel_sizes=intermediate_channel_sizes
+    )
+
+    assert intermediate_channel_sizes == [8]
+    assert capsys.readouterr().out == ""
+
+
 def test_cnn_embedding_default_channels_follow_num_conv_layers():
     """The default channels double per layer, starting at 6."""
     embedding_net = CNNEmbedding((32,), num_conv_layers=3)

From 0db5088a4f120a672ff001ecfa6a9ee20a51f0a8 Mon Sep 17 00:00:00 2001
From: Jan <jan.boelts@mailbox.org>
Date: Thu, 29 Oct 2026 07:52:14 +0100
Subject: [PATCH 08/12] fix: mean aggregation handles observations without real
 trials

PermutationInvariantEmbedding warns that it treats an observation with only padded trials as one with zero trials. With aggregation_fn="sum" that held, but "mean" divided by zero and failed with 'AssertionError: NaNs in embedding.'. The mean now divides by at least one, so such observations get a zero embedding with either aggregation.

The remaining NaN check fires only when the trial net returns NaN or inf, and now raises a ValueError that says so.
---
 .../embedding_nets/permutation_invariant.py       | 15 ++++++++++-----
 tests/embedding_net_test.py                       |  7 +++++--
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/sbi/neural_nets/embedding_nets/permutation_invariant.py b/sbi/neural_nets/embedding_nets/permutation_invariant.py
index 02b45cf91..66e0d91f7 100644
--- a/sbi/neural_nets/embedding_nets/permutation_invariant.py
+++ b/sbi/neural_nets/embedding_nets/permutation_invariant.py
@@ -142,15 +142,20 @@ def forward(self, x: Tensor) -> Tensor:
         trial_embeddings = trial_embeddings * is_real_trial.unsqueeze(-1).float()
 
         # Take mean over permutation dimension divide by number of trials
-        # (instead of just taking torch.mean) to account for masking.
+        # (instead of just taking torch.mean) to account for masking. Observations
+        # without real trials get a zero embedding, as with "sum".
         if self.aggregation_fn == "mean":
-            combined_embedding = (
-                trial_embeddings.sum(dim=self.aggregation_dim) / trial_counts
-            )
+            combined_embedding = trial_embeddings.sum(
+                dim=self.aggregation_dim
+            ) / trial_counts.clamp(min=1)
         else:
             combined_embedding = trial_embeddings.sum(dim=self.aggregation_dim)
 
-        assert not torch.isnan(combined_embedding).any(), "NaNs in embedding."
+        if torch.isnan(combined_embedding).any():
+            raise ValueError(
+                "NaNs in embedding: the trial_net returned NaN or inf, e.g., because "
+                "the input contains inf."
+            )
 
         # add number of trials as additional input
         return self.fc_subnet(torch.cat([combined_embedding, trial_counts], dim=1))
diff --git a/tests/embedding_net_test.py b/tests/embedding_net_test.py
index b118f5da3..1d24ad779 100644
--- a/tests/embedding_net_test.py
+++ b/tests/embedding_net_test.py
@@ -630,15 +630,18 @@ def test_permutation_invariant_embedding_ignores_padding_per_observation(
     assert torch.allclose(batched[1], embedding(x[1:])[0])
 
 
+@pytest.mark.parametrize("aggregation_fn", ("mean", "sum"))
 @pytest.mark.parametrize(
     "pattern", ("padded_trial", "partially_nan_trial", "no_real_trial")
 )
-def test_permutation_invariant_embedding_warns_on_invalid_padding(pattern):
+def test_permutation_invariant_embedding_warns_on_invalid_padding(
+    pattern, aggregation_fn
+):
     """NaN outside whole padded trials must not silently become data; warn once."""
     embedding = PermutationInvariantEmbedding(
         FCEmbedding(input_dim=2, output_dim=4),
         trial_net_output_dim=4,
-        aggregation_fn="sum",
+        aggregation_fn=aggregation_fn,
     )
     x = torch.randn(2, 3, 2)
     x[0, 1:] = float("nan")

From b02f7cfb4bc080849bc6c9bc0bc8983f217db96c Mon Sep 17 00:00:00 2001
From: Jan <jan.boelts@mailbox.org>
Date: Thu, 29 Oct 2026 10:42:16 +0100
Subject: [PATCH 09/12] fix: dilation and WaveNet aggregator errors name their
 arguments

WaveNetSRLikeAggregator checked its intermediate_channel_sizes and stride_sizes lists with assert, so under python -O a list of the wrong length built a net with the wrong output size. Both checks now raise ValueError with the lengths. The stride message also read "same size asthe kernel size list".

The CausalCNNEmbedding errors for an unknown dilation string and a dilation of the wrong type now name dilation and what was passed.

In CNNEmbedding, the hint variable of the size error is renamed to fewer_layers_hint, since it holds message text.
---
 sbi/neural_nets/embedding_nets/causal_cnn.py | 30 ++++++++++++--------
 sbi/neural_nets/embedding_nets/cnn.py        |  4 +--
 2 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/sbi/neural_nets/embedding_nets/causal_cnn.py b/sbi/neural_nets/embedding_nets/causal_cnn.py
index 856bf2747..18f38cc85 100644
--- a/sbi/neural_nets/embedding_nets/causal_cnn.py
+++ b/sbi/neural_nets/embedding_nets/causal_cnn.py
@@ -92,16 +92,19 @@ def WaveNetSRLikeAggregator(
         ]
     if intermediate_channel_sizes is None:
         intermediate_channel_sizes = [64]
-    assert len(intermediate_channel_sizes) == len(kernel_sizes) - 1, (
-        "Provided kernel size list should be exactly one element longer "
-        "than channel size list."
-    )
+    if len(intermediate_channel_sizes) != len(kernel_sizes) - 1:
+        raise ValueError(
+            "intermediate_channel_sizes needs one entry less than kernel_sizes, but "
+            f"has {len(intermediate_channel_sizes)} entries for {len(kernel_sizes)} "
+            "kernel sizes."
+        )
     intermediate_channel_sizes = [*intermediate_channel_sizes, output_dim]
     if isinstance(stride_sizes, list):
-        assert len(stride_sizes) == len(kernel_sizes), (
-            "Provided stride size list should be have the same size as"
-            "the kernel size list."
-        )
+        if len(stride_sizes) != len(kernel_sizes):
+            raise ValueError(
+                f"stride_sizes has {len(stride_sizes)} entries, but kernel_sizes has "
+                f"{len(kernel_sizes)}. Pass one stride per kernel size."
+            )
     else:
         stride_sizes = [stride_sizes] * len(kernel_sizes)
 
@@ -221,13 +224,16 @@ def __init__(
                     dilation_per_layer = [1] * num_conv_layers
                 case _:
                     raise ValueError(
-                        f"{dilation} is not a valid option, please use \"none\", "
-                        "\"exponential\", or \"exponential_cyclic\", or pass a list "
-                        "of dilation sizes."
+                        f"dilation must be \"none\", \"exponential\", "
+                        f"\"exponential_cyclic\" or a list of dilation sizes, but got "
+                        f"\"{dilation}\"."
                     )
         else:
             if not isinstance(dilation, list):
-                raise TypeError("Please pass dilation size as list or a string option.")
+                raise TypeError(
+                    "dilation must be a list of dilation sizes or a string option, "
+                    f"but got {type(dilation).__name__}."
+                )
             if len(dilation) != num_conv_layers:
                 raise ValueError(
                     f"dilation has {len(dilation)} entries, but num_conv_layers is "
diff --git a/sbi/neural_nets/embedding_nets/cnn.py b/sbi/neural_nets/embedding_nets/cnn.py
index 08bff4a2c..ce7163469 100644
--- a/sbi/neural_nets/embedding_nets/cnn.py
+++ b/sbi/neural_nets/embedding_nets/cnn.py
@@ -182,13 +182,13 @@ def __init__(
             cnn_output_size = get_new_cnn_output_size(cnn_output_size, conv_layer, pool)
 
             if not all(size > 0 for size in cnn_output_size):
-                fewer_layers = (
+                fewer_layers_hint = (
                     f", or set num_conv_layers to at most {ii}" if ii > 0 else ""
                 )
                 raise ValueError(
                     f"CNN output size {cnn_output_size} is not positive at layer "
                     f"{ii + 1}. Use a smaller kernel_size or pool_kernel_size"
-                    f"{fewer_layers}."
+                    f"{fewer_layers_hint}."
                 )
 
         self.cnn_subnet = nn.Sequential(*cnn_layers)

From d14c4e3dc25a5481e09556475fd0b6b6a63507fd Mon Sep 17 00:00:00 2001
From: Jan <jan.boelts@mailbox.org>
Date: Thu, 29 Oct 2026 10:42:16 +0100
Subject: [PATCH 10/12] test: observations without real trials get the zero
 embedding

The padding test with an observation of only padded trials checked the warning, but not the embedding. It now checks that both "mean" and "sum" pass a zero aggregate and a zero trial count to the output net.
---
 tests/embedding_net_test.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tests/embedding_net_test.py b/tests/embedding_net_test.py
index 1d24ad779..9dfe9badf 100644
--- a/tests/embedding_net_test.py
+++ b/tests/embedding_net_test.py
@@ -652,11 +652,14 @@ def test_permutation_invariant_embedding_warns_on_invalid_padding(
 
     with warnings.catch_warnings(record=True) as caught:
         warnings.simplefilter("always")
-        embedding(x)
+        output = embedding(x)
         embedding(x)
 
     issued = [w for w in caught if "should only mark padded trials" in str(w.message)]
     assert len(issued) == (0 if pattern == "padded_trial" else 1)
+    if pattern == "no_real_trial":
+        # Zero aggregate and zero trial count, for "mean" as for "sum".
+        assert torch.allclose(output[1], embedding.fc_subnet(torch.zeros(1, 5))[0])
 
 
 @pytest.mark.parametrize("input_shape", [(32, 32), (32, 64), (111, 111)])

From 03f0e40a19bd9566fb71ff6761e0ff0d86008114 Mon Sep 17 00:00:00 2001
From: Jan <jan.boelts@mailbox.org>
Date: Wed, 11 Nov 2026 16:34:45 +0100
Subject: [PATCH 11/12] fix: permutation-invariant embedding rejects inf input

nan_to_num(x, nan=0.0) also turns inf into the float maximum before the trial_net runs, so the NaN check after aggregation never saw it: one inf trial reached the density estimator as an embedding of about 1e37. An isfinite check on the aggregated embedding does not catch it either, because the trial_net shrinks the clamped value to a large finite number.

forward now raises on inf in the input, before nan_to_num. NaN padding is unaffected, since isinf(NaN) is False. The check after aggregation uses isfinite, so a trial_net that returns inf for finite input also raises.

Training with NaN padding needs exclude_invalid_x=False, which also stops dropping inf simulations, so this path is reachable in normal use.

Reported by Manuel Gloeckler in the review of #2113.
---
 .../embedding_nets/permutation_invariant.py   | 14 +++++++--
 tests/embedding_net_test.py                   | 29 +++++++++++++++++++
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/sbi/neural_nets/embedding_nets/permutation_invariant.py b/sbi/neural_nets/embedding_nets/permutation_invariant.py
index 66e0d91f7..8db570d9a 100644
--- a/sbi/neural_nets/embedding_nets/permutation_invariant.py
+++ b/sbi/neural_nets/embedding_nets/permutation_invariant.py
@@ -119,6 +119,8 @@ def forward(self, x: Tensor) -> Tensor:
             x: Input tensor (batch_size, permutation_dim, input_dim)
         Returns:
             Network output (batch_size, output_dim).
+        Raises:
+            ValueError: If `x` contains inf, or the trial_net returns NaN or inf.
         """
 
         if not getattr(self, "_warned_nan_padding", False):
@@ -135,6 +137,12 @@ def forward(self, x: Tensor) -> Tensor:
         is_real_trial = ~torch.isnan(x).all(-1)
         trial_counts = is_real_trial.sum(dim=self.aggregation_dim).unsqueeze(-1)
 
+        # nan_to_num below would turn inf into a finite number.
+        if torch.isinf(x).any():
+            raise ValueError(
+                "Input contains inf. Only NaN is allowed, to mark padded trials."
+            )
+
         # apply trial net with nan entries replaced with 0
         masked_x = torch.nan_to_num(x, nan=0.0)
         trial_embeddings = self.trial_net(masked_x)
@@ -151,10 +159,10 @@ def forward(self, x: Tensor) -> Tensor:
         else:
             combined_embedding = trial_embeddings.sum(dim=self.aggregation_dim)
 
-        if torch.isnan(combined_embedding).any():
+        if not torch.isfinite(combined_embedding).all():
             raise ValueError(
-                "NaNs in embedding: the trial_net returned NaN or inf, e.g., because "
-                "the input contains inf."
+                "Non-finite value in the aggregated trial embedding: the trial_net "
+                "returned NaN or inf."
             )
 
         # add number of trials as additional input
diff --git a/tests/embedding_net_test.py b/tests/embedding_net_test.py
index 9dfe9badf..4e5b0c15f 100644
--- a/tests/embedding_net_test.py
+++ b/tests/embedding_net_test.py
@@ -662,6 +662,35 @@ def test_permutation_invariant_embedding_warns_on_invalid_padding(
         assert torch.allclose(output[1], embedding.fc_subnet(torch.zeros(1, 5))[0])
 
 
+@pytest.mark.parametrize("aggregation_fn", ("mean", "sum"))
+@pytest.mark.parametrize("with_padding", (False, True))
+@pytest.mark.parametrize("inf", (float("inf"), float("-inf")))
+def test_permutation_invariant_embedding_rejects_inf(inf, with_padding, aggregation_fn):
+    """Inf in a real trial raises instead of passing as a huge finite embedding."""
+    embedding = PermutationInvariantEmbedding(
+        FCEmbedding(input_dim=2, output_dim=4),
+        trial_net_output_dim=4,
+        aggregation_fn=aggregation_fn,
+    )
+    x = torch.randn(2, 3, 2)
+    x[0, 1, 0] = inf
+    if with_padding:
+        x[:, 2] = float("nan")
+
+    with pytest.raises(ValueError, match="Input contains inf"):
+        embedding(x)
+
+
+def test_permutation_invariant_embedding_rejects_non_finite_trial_net_output():
+    """A trial_net that returns inf for finite input raises."""
+    embedding = PermutationInvariantEmbedding(nn.Linear(2, 4), trial_net_output_dim=4)
+    with torch.no_grad():
+        embedding.trial_net.bias.fill_(float("inf"))
+
+    with pytest.raises(ValueError, match="trial_net returned NaN or inf"):
+        embedding(torch.randn(2, 3, 2))
+
+
 @pytest.mark.parametrize("input_shape", [(32, 32), (32, 64), (111, 111)])
 @pytest.mark.parametrize("num_channels", (1, 2, 3))
 @pytest.mark.parametrize("change_c_mode", ["conv", "zeros"])

From 26e622c6033079a57dc0424670b044e36dacc624 Mon Sep 17 00:00:00 2001
From: Jan <jan.boelts@mailbox.org>
Date: Wed, 11 Nov 2026 16:35:04 +0100
Subject: [PATCH 12/12] test: pin the two pool_kernel_size errors of
 CausalCNNEmbedding

Both error messages contain the bare word pool_kernel_size, so either case passed if the other branch raised. Each case now matches the part of its message that the other lacks.

Reported by Manuel Gloeckler in the review of #2113.
---
 tests/embedding_net_test.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/embedding_net_test.py b/tests/embedding_net_test.py
index 4e5b0c15f..754f18b32 100644
--- a/tests/embedding_net_test.py
+++ b/tests/embedding_net_test.py
@@ -229,8 +229,8 @@ def test_cnn_embedding_errors_name_arguments(input_shape, kwargs, match):
         ({"num_conv_layers": 0}, "num_conv_layers must be at least 1"),
         ({"out_channels_per_layer": [16, 16]}, "num_conv_layers is 5"),
         ({"dilation": [1, 2]}, "dilation has 2 entries"),
-        ({"pool_kernel_size": 128}, "pool_kernel_size"),
-        ({"pool_kernel_size": 64}, "pool_kernel_size"),
+        ({"pool_kernel_size": 128}, r"pool_kernel_size \(128\) must not be larger"),
+        ({"pool_kernel_size": 64}, "the default aggregator needs at least 2"),
     ],
 )
 def test_causal_cnn_embedding_errors_name_arguments(kwargs, match):