Skip to content

Abstract Sampler

popinn.AbstractSampler ¤

Abstract base class for collocation re-samplers.

A sampler returns a new data container with some fields redrawn from the given PRNG key.

Subclasses must implement __call__. To keep the jitted training step from recompiling, a sampler MUST return batches of constant shape across calls.

Note: the trainer calls samplers structurally (sample_fn(key)) and does not require AbstractSampler specifically -- any callable with the same signature is accepted.

Source code in src/popinn/sampling.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
class AbstractSampler(eqx.Module):
    """Abstract base class for collocation re-samplers.

    A sampler returns a new data container with some fields redrawn from the given PRNG key.

    Subclasses must implement `__call__`. To keep the jitted training step
    from recompiling, a sampler MUST return batches of constant shape across
    calls.

    Note: the trainer calls samplers structurally (`sample_fn(key)`) and does
    not require AbstractSampler specifically -- any callable with the same
    signature is accepted.
    """

    @abc.abstractmethod
    def __call__(self, key: PRNGKeyArray) -> eqx.Module:
        """Draw a new batch from `key`.

        Args:
            key (PRNGKeyArray): PRNG key supplied by the trainer (freshly
                split each resample).

        Returns:
            (equinox.Module): A new data container with the resampled
                fields replaced and all others carried through unchanged.
        """
        raise NotImplementedError

    # The named `field` is assumed to be a tuple of 1-D coordinate arrays, one
    # per coordinate axis. Each
    # axis is independently redrawn uniformly within its bounds; every other
    # field of the reference batch is carried through unchanged. Each call
    # constructs a new `equinox.Module` data container.

    # *Fields* :
    # - `data` (`equinox.Module`): the reference training-data container. Its
    #     non-resampled fields are reused as-is in every returned batch.
    #     Held as an array-bearing (non-static) field, so this sampler is
    #     itself a pytree containing the reference batch.
    # - `bounds` (`tuple`): one (min, max) pair per coordinate axis of `field`.
    # - `counts` (`tuple`): points per coordinate axis (static).
    # - `field` (`str`): name of the coordinate field to resample in `data` (static).
    # """
__call__(key) ¤

Draw a new batch from key.

Parameters:

Name Type Description Default
key typing.Union

PRNG key supplied by the trainer (freshly split each resample).

 required

Returns:

Type Description
equinox.Module

A new data container with the resampled fields replaced and all others carried through unchanged.
Source code in src/popinn/sampling.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
@abc.abstractmethod
def __call__(self, key: PRNGKeyArray) -> eqx.Module:
    """Draw a new batch from `key`.

    Args:
        key (PRNGKeyArray): PRNG key supplied by the trainer (freshly
            split each resample).

    Returns:
        (equinox.Module): A new data container with the resampled
            fields replaced and all others carried through unchanged.
    """
    raise NotImplementedError