Check number of context sets and check aux variable

Author: wesselb

neuralprocesses/architectures/convgnp.py

@@ -1,6 +1,4 @@
-import lab as B
 import neuralprocesses as nps  # This fixes inspection below.
-import wbml.out as out
 from plum import convert
 
 from ..util import register_model
@@ -30,7 +28,7 @@ def _convgnp_resolve_architecture(
     elif "conv" in conv_arch:
         conv_out_channels = conv_channels
         if conv_receptive_field is None:
-            raise ValueError(f"Must specify `conv_receptive_field`.")
+            raise ValueError("Must specify `conv_receptive_field`.")
     else:
         raise ValueError(f'Architecture "{conv_arch}" invalid.')
     return conv_out_channels
@@ -62,6 +60,16 @@ def _convgnp_construct_encoder_setconvs(
     )
 
 
+def _convgnp_assert_form_contexts(nps, dim_yc):
+    if len(dim_yc) == 1:
+        return nps.Chain(
+            nps.SqueezeParallel(),
+            nps.AssertNoParallel(),
+        )
+    else:
+        return nps.AssertParallel(len(dim_yc))
+
+
 def _convgnp_construct_decoder_setconv(
     nps,
     decoder_scale,
@@ -159,8 +167,8 @@ def construct_convgnp(
         conv_receptive_field (float, optional): Receptive field of the standard
             architecture. Must be specified if `conv_arch` is set to `"conv"`.
         conv_layers (int, optional): Layers of the standard architecture. Defaults to 8.
-        conv_channels (int, optional): Channels of the standard architecture. Defaults to
-            64.
+        conv_channels (int, optional): Channels of the standard architecture. Defaults
+            to 64.
         num_basis_functions (int, optional): Number of basis functions for the
             low-rank likelihood. Defaults to `512`.
         dim_lv (int, optional): Dimensionality of the latent variable. Defaults to 0.
@@ -255,6 +263,8 @@ def construct_convgnp(
             # Not necessary. Just let the CNN produce the right number of channels.
             conv_out_channels = likelihood_in_channels
             linear_after_set_conv = lambda x: x
+        # Also assert that there is no augmentation given.
+        likelihood = nps.Chain(nps.AssertNoAugmentation(), likelihood)
 
     # Construct the core CNN architectures for the encoder, which is only necessary
     # if we're using a latent variable, and for the decoder. First, we determine
@@ -348,6 +358,7 @@ def construct_convgnp(
         nps.FunctionalCoder(
             disc,
             nps.Chain(
+                _convgnp_assert_form_contexts(nps, dim_yc),
                 nps.PrependDensityChannel(),
                 _convgnp_construct_encoder_setconvs(
                     nps,

neuralprocesses/architectures/fullconvgnp.py

@@ -1,15 +1,16 @@
+import neuralprocesses as nps  # This fixes inspection below.
 import wbml.out as out
 
-import neuralprocesses as nps  # This fixes inspection below.
+from ..util import register_model
 from .convgnp import (
-    _convgnp_init_dims,
-    _convgnp_resolve_architecture,
+    _convgnp_assert_form_contexts,
+    _convgnp_construct_decoder_setconv,
     _convgnp_construct_encoder_setconvs,
+    _convgnp_init_dims,
     _convgnp_optional_division_by_density,
-    _convgnp_construct_decoder_setconv,
+    _convgnp_resolve_architecture,
 )
 from .util import parse_transform
-from ..util import register_model
 
 __all__ = ["construct_fullconvgnp"]
 
@@ -211,6 +212,7 @@ def construct_fullconvgnp(
                 nps.FunctionalCoder(
                     disc_mean,
                     nps.Chain(
+                        _convgnp_assert_form_contexts(nps, dim_yc),
                         nps.PrependDensityChannel(),
                         _convgnp_construct_encoder_setconvs(
                             nps,
@@ -232,6 +234,7 @@ def construct_fullconvgnp(
                     disc_kernel,
                     nps.MapDiagonal(  # Map to diagonal of squared space.
                         nps.Chain(
+                            _convgnp_assert_form_contexts(nps, dim_yc),
                             nps.PrependDensityChannel(),
                             _convgnp_construct_encoder_setconvs(
                                 nps,

neuralprocesses/architectures/gnp.py

@@ -2,6 +2,7 @@
 from plum import convert
 
 from ..util import register_model
+from .convgnp import _convgnp_assert_form_contexts
 from .util import construct_likelihood, parse_transform
 
 __all__ = ["construct_gnp"]
@@ -185,6 +186,7 @@ def construct_mlp(dim_yci):
     encoder = nps.Chain(
         # We need to explicitly copy, because there will be multiple context sets in
         # parallel, which will otherwise dispatch to the wrong method.
+        _convgnp_assert_form_contexts(nps, dim_yc),
         nps.Copy(2 + (dim_lv > 0)),
         nps.Parallel(
             nps.Chain(

neuralprocesses/coders/augment.py

@@ -4,9 +4,9 @@
 from ..augment import AugmentedInput
 from ..datadims import data_dims
 from ..materialise import _repeat_concat
-from ..util import register_module, register_composite_coder
+from ..util import register_composite_coder, register_module
 
-__all__ = ["Augment"]
+__all__ = ["Augment", "AssertNoAugmentation"]
 
 
 @register_composite_coder
@@ -47,3 +47,18 @@ def _augment(xz: AugmentedInput, z: B.Numeric):
 @_dispatch
 def _augment(xz: AugmentedInput):
     return xz.x
+
+
+@register_module
+class AssertNoAugmentation:
+    """Assert no augmentation of the target inputs."""
+
+
+@_dispatch
+def code(coder: AssertNoAugmentation, xz, z, x, **kw_args):
+    return xz, z
+
+
+@_dispatch
+def code(coder: AssertNoAugmentation, xz, z, x: AugmentedInput, **kw_args):
+    raise AssertionError("Did not expect augmentation of the target inputs.")

neuralprocesses/coders/shaping.py

@@ -5,7 +5,14 @@
 from ..parallel import Parallel
 from ..util import register_module, split
 
-__all__ = ["Identity", "Splitter", "RestructureParallel"]
+__all__ = [
+    "Identity",
+    "Splitter",
+    "RestructureParallel",
+    "AssertParallel",
+    "SqueezeParallel",
+    "AssertNoParallel",
+]
 
 
 @register_module
@@ -93,3 +100,82 @@ def _restructure_create(element_map, i):
 @_dispatch
 def _restructure_create(element_map, x: tuple):
     return Parallel(*(_restructure_create(element_map, xi) for xi in x))
+
+
+@register_module
+class AssertParallel:
+    """Assert a parallel of `n` elements.
+
+    Args:
+        n (int): Number of elements asserted in parallel.
+
+    Attributes:
+        n (int): Number of elements asserted in parallel.
+    """
+
+    def __init__(self, n):
+        self.n = n
+
+
+@_dispatch
+def code(p: AssertParallel, xz, z, x, **kw_args):
+    raise AssertionError(f"Expected a parallel of elements, but got `{xz}` and `{z}`.")
+
+
+@_dispatch
+def code(p: AssertParallel, xz: Parallel, z: Parallel, x, **kw_args):
+    if not len(xz) == len(z) == p.n:
+        raise AssertionError(
+            f"Expected a parallel of {p.n} elements, "
+            f"but got {len(xz)} inputs and {len(z)} outputs."
+        )
+    return xz, z
+
+
+@register_module
+class SqueezeParallel:
+    """If there is a parallel of exactly one element, remove the parallel."""
+
+
+@_dispatch
+def code(p: SqueezeParallel, xz, z, x, **kw_args):
+    return xz, z
+
+
+@_dispatch
+def code(p: SqueezeParallel, xz: Parallel, z: Parallel, x, **kw_args):
+    if len(xz) == len(z) == 1:
+        return xz[0], z[0]
+    else:
+        return xz, z
+
+
+@register_module
+class AssertNoParallel:
+    """Assert exactly one element in parallel or not a parallel of elements."""
+
+
+@_dispatch
+def code(p: AssertNoParallel, xz, z, x, **kw_args):
+    return xz, z
+
+
+@_dispatch
+def code(p: AssertNoParallel, xz: Parallel, z, x, **kw_args):
+    raise AssertionError(
+        "Expected not a parallel of elements, but got inputs in parallel."
+    )
+
+
+@_dispatch
+def code(p: AssertNoParallel, xz, z: Parallel, x, **kw_args):
+    raise AssertionError(
+        "Expected not a parallel of elements, but got outputs in parallel."
+    )
+
+
+@_dispatch
+def code(p: AssertNoParallel, xz: Parallel, z: Parallel, x, **kw_args):
+    raise AssertionError(
+        "Expected not a parallel of elements, but got inputs and outputs in parallel."
+    )

neuralprocesses/mask.py

@@ -12,7 +12,8 @@ class Masked:
     """A masked output.
 
     Args:
-        y (tensor): Output to mask.
+        y (tensor): Output to mask. The masked values can have any non-NaN value, but
+            they cannot be NaN!
         mask (tensor): A mask consisting of zeros and ones and just one channel.
 
     Attributes:

neuralprocesses/parallel.py

@@ -2,9 +2,12 @@
 from matrix.util import indent
 
 from . import _dispatch
-from .util import register_module, is_framework_module
+from .util import is_framework_module, register_module
 
-__all__ = ["Parallel", "broadcast_coder_over_parallel"]
+__all__ = [
+    "Parallel",
+    "broadcast_coder_over_parallel",
+]
 
 
 @register_module

tests/test_architectures.py

@@ -4,6 +4,7 @@
 import numpy as np
 import pytest
 from neuralprocesses.aggregate import Aggregate, AggregateInput
+from neuralprocesses.parallel import Parallel
 from plum import isinstance
 
 from .util import approx, generate_data
@@ -420,3 +421,42 @@ def test_data_eq(nps, dim_x, dim_y, constructor, config, dim_lv):
     batch = gen.generate_batch()
     pred = model(batch["contexts"], batch["xt"])
     check_prediction(nps, pred, batch["yt"])
+
+
+@pytest.mark.parametrize(
+    "constructor_kw_args",
+    [
+        ("construct_gnp", {}),
+        ("construct_agnp", {}),
+        ("construct_convgnp", {"points_per_unit": 4}),
+        ("construct_fullconvgnp", {"points_per_unit": 4}),
+    ],
+)
+def test_context_verification(nps, constructor_kw_args):
+    constructor, kw_args = constructor_kw_args
+
+    c = (B.randn(nps.dtype, 4, 1, 5), B.randn(nps.dtype, 4, 1, 5))
+    xt = B.randn(nps.dtype, 4, 1, 15)
+
+    # Test one context set.
+    model = getattr(nps, constructor)(dim_yc=1, dtype=nps.dtype, **kw_args)
+    pred1 = model(*c, xt)
+    pred2 = model([c], xt)
+    approx(pred1.mean, pred2.mean)
+    with pytest.raises(AssertionError, match="(?i)got inputs and outputs in parallel"):
+        model([c, c], xt)
+    with pytest.raises(AssertionError, match="(?i) got inputs in parallel"):
+        model(Parallel(c[0], c[0]), c[1], xt)
+    with pytest.raises(AssertionError, match="(?i) got outputs in parallel"):
+        model(c[0], Parallel(c[1], c[1]), xt)
+
+    # Test two context sets.
+    model = getattr(nps, constructor)(dim_yc=(1, 1), dtype=nps.dtype, **kw_args)
+    model([c, c], xt)
+    with pytest.raises(AssertionError, match="(?i)expected a parallel of elements"):
+        model([c], xt)
+    with pytest.raises(
+        AssertionError,
+        match="(?i)expected a parallel of 2 elements, but got 4 inputs and 4 outputs",
+    ):
+        model([c, c, c, c], xt)

tests/test_augment.py

@@ -1,5 +1,6 @@
 import lab as B
 import pytest
+from plum import NotFoundLookupError
 
 from .test_architectures import check_prediction
 from .util import nps  # noqa
@@ -37,3 +38,23 @@ def test_convgnp_auxiliary_variable(nps):
     )
 
     check_prediction(nps, pred, B.randn(nps.dtype, 16, 3, 15))
+
+    # Check that the model cannot be run forward without the auxiliary variable.
+    with pytest.raises(NotFoundLookupError):
+        model(
+            [observed_data, aux_var1, aux_var2],
+            B.randn(nps.dtype, 16, 2, 15),
+        )
+
+
+def test_convgnp_auxiliary_variable_given_but_not_specified(nps):
+    """Test that giving the auxiliary variable without specifying `dim_aux_t` raises
+    an error."""
+    model = nps.construct_convgnp(points_per_unit=4)
+    with pytest.raises(AssertionError, match="(?i)did not expect augmentation"):
+        model(
+            B.randn(nps.dtype, 4, 1, 15),
+            B.randn(nps.dtype, 4, 1, 15),
+            B.randn(nps.dtype, 4, 1, 10),
+            aux_t=B.randn(nps.dtype, 4, 2, 10),
+        )