Skip to content

testing

PotentialTests

A mixin class providing testing utilities for validating Potential implementations.

This class provides methods to verify the mathematical consistency and correctness of Potential implementations through various assertions:

  • logw_next consistency: Verifies that token-level log weights are consistent with prefix and complete scores.
  • Autoregressive factorization: Validates that complete scores factor correctly as a sum of log token weights (with an additional correction term corresponding to the prefix weight of the empty sequence).
  • Batch consistency: Ensures batch operations produce identical results to their non-batch counterparts.

All Potential instances inherit from this class and thus automatically gain access to these testing utilities.

Source code in genlm/control/potential/testing.py 
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 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
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
class PotentialTests:
    """A mixin class providing testing utilities for validating Potential implementations.

    This class provides methods to verify the mathematical consistency and correctness
    of Potential implementations through various assertions:

    - logw_next consistency: Verifies that token-level log weights are consistent with
      prefix and complete scores.
    - Autoregressive factorization: Validates that complete scores factor correctly as
      a sum of log token weights (with an additional correction term corresponding to the
      prefix weight of the empty sequence).
    - Batch consistency: Ensures batch operations produce identical results to
      their non-batch counterparts.

    All Potential instances inherit from this class and thus automatically gain access to these
    testing utilities.
    """

    colors = {
        "green": "\033[92m",
        "yellow": "\033[93m",
        "blue": "\033[94m",
        "magenta": "\033[95m",
        "cyan": "\033[96m",
        "red": "\033[91m",
        "reset": "\033[0m",
    }

    async def assert_logw_next_consistency(
        self, context, rtol=1e-3, atol=1e-5, top=None, verbosity=0, method_args=()
    ):
        """
        Assert that `logw_next` is consistent with `prefix` and `complete`.

        For a `context` of tokens $x_1, \\ldots, x_{n-1}$, this checks (in log space) whether:

        $$
        \\textsf{logw\\_next}(x_n | x_1, \\ldots, x_{n-1}) = \\textsf{score}(x_1, \\ldots, x_n) - \\textsf{prefix}(x_1, \\ldots, x_{n-1})
        $$
        for $x_n \\in \\textsf{vocab_eos}$, i.e., the potential's vocabulary and end-of-sequence token.

        Args:
            context (list): Context to test.
            rtol (float): Relative tolerance for floating point comparison.
            atol (float): Absolute tolerance for floating point comparison.
            top (int):If specified, only test the top-k tokens by log weight. If None, test all tokens.
            verbosity (int): Verbosity level.
            method_args (tuple): Positional arguments to pass to `logw_next`, `prefix`, and `batch_score`.
                Defaults to empty tuple.

        Raises:
            AssertionError: If `logw_next` is not consistent with `prefix` and `complete`.
        """
        top_logw_next = (await self.logw_next(context, *method_args)).materialize(
            top=top
        )
        tokens = list(top_logw_next.keys())
        extended = [[*context, x] for x in tokens]

        context_w = await self.prefix(context, *method_args)
        extended_ws = np.array(
            await asyncio.gather(*[self.score(e, *method_args) for e in extended])
        )

        wants = np.array([top_logw_next[x] for x in tokens])
        haves = extended_ws - context_w

        errors, valids = [], []
        for i, (want, have) in enumerate(zip(wants, haves)):
            abs_diff, rel_diff = self._compute_diff(want, have)
            info = (want, have, abs_diff, rel_diff, tokens[i])
            (valids if abs_diff <= atol and rel_diff <= rtol else errors).append(info)

        if valids and verbosity > 0:
            print(
                f"{self.colors['green']}logw_next consistency with context={context!r} satisfied for tokens:{self.colors['reset']}\n"
            )
            for valid in valids:
                want, have, abs_diff, rel_diff, token = valid
                print(
                    self._format_diff(want, have, abs_diff, rel_diff, atol, rtol, token)
                )

        if errors:
            error_msg = f"{self.colors['red']}logw_next consistency with context={context!r} not satisfied for tokens:{self.colors['reset']}\n\n"
            for error in errors:
                want, have, abs_diff, rel_diff, token = error
                error_msg += self._format_diff(
                    want, have, abs_diff, rel_diff, atol, rtol, token
                )
            raise AssertionError(error_msg)

    async def assert_autoreg_fact(
        self, context, rtol=1e-3, atol=1e-5, verbosity=0, method_args=()
    ):
        """
        Assert that `complete` factors as an autoregressive sum of `logw_next`s.

        For a `context` of tokens $x_1, \\ldots, x_n$, this checks (in log space) whether:

        $$
        \\textsf{complete}(x_1, \\ldots, x_n) - \\textsf{prefix}(\\epsilon) = \\textsf{logw\\_next}(\\textsf{eos} \\mid x_1, \\ldots, x_{n}) + \\sum_{i=1}^{n} \\textsf{logw_next}(x_i \\mid x_1, \\ldots, x_{i-1})
        $$
        where $\\epsilon$ is the empty sequence.

        Args:
            context (list): Context to test.
            rtol (float): Relative tolerance for floating point comparison.
            atol (float): Absolute tolerance for floating point comparison.
            verbosity (int): Verbosity level.
            method_args (tuple): Positional arguments to pass to `complete`, `prefix`, and `logw_next`.
                Defaults to empty tuple.

        Raises:
            AssertionError: If the autoregressive factorization is not satisfied.
        """
        want = (await self.complete(context, *method_args)) - (
            await self.prefix([], *method_args)
        )

        logw_next_results = await asyncio.gather(
            *[self.logw_next(context[:i], *method_args) for i in range(len(context))],
            self.logw_next(context, *method_args),
        )

        have = (
            sum(logw_next_results[i][context[i]] for i in range(len(context)))
            + logw_next_results[-1][self.eos]
        )

        abs_diff, rel_diff = self._compute_diff(want, have)
        if abs_diff > atol or rel_diff > rtol:
            error_msg = (
                f"{self.colors['red']}Factorization not satisfied for context {context!r}:{self.colors['reset']}\n"
                + self._format_diff(want, have, abs_diff, rel_diff, atol, rtol)
            )
            raise AssertionError(error_msg)

        if verbosity > 0:
            print(
                f"{self.colors['green']}Factorization property satisfied for context {context}:{self.colors['reset']}\n"
            )
            print(self._format_diff(want, have, abs_diff, rel_diff, atol, rtol))

    async def assert_batch_consistency(
        self,
        contexts,
        rtol=1e-3,
        atol=1e-5,
        verbosity=0,
        batch_method_args=(),
        method_args=(),
    ):
        """
        Assert that batch results are equal to non-batch results.

        Args:
            contexts (list[list]): Contexts to test.
            rtol (float): Relative tolerance for floating point comparison.
            atol (float): Absolute tolerance for floating point comparison.
            verbosity (int): Verbosity level.
            batch_method_args (tuple): Positional arguments to pass to batch methods.
                Defaults to empty tuple.
            method_args (tuple): Positional arguments to pass to underlying potential methods.
                Defaults to empty tuple.

        Raises:
            AssertionError: If the batch results are not equal to the non-batch results.
        """
        batch_logw_nexts = await self.batch_logw_next(contexts, *batch_method_args)
        batch_scores = await self.batch_score(contexts, *batch_method_args)

        for i, context in enumerate(contexts):
            logw_next = await self.logw_next(context, *method_args)
            try:
                np.testing.assert_allclose(
                    batch_logw_nexts[i].weights, logw_next.weights, rtol=rtol, atol=atol
                )
                if verbosity > 0:
                    print(
                        f"{self.colors['green']}Batch logw_next consistency satisfied for context {context}:{self.colors['reset']}"
                    )
                    print(
                        f"{self.colors['green']}Non-batched: {logw_next.weights}\n"
                        + f"{self.colors['green']}Batched:     {batch_logw_nexts[i].weights}{self.colors['reset']}\n"
                    )
            except AssertionError:
                raise AssertionError(
                    f"{self.colors['red']}Batch logw_next mismatch for context {context}:{self.colors['reset']}\n"
                    + f"{self.colors['green']}Non-batched: {logw_next.weights}\n"
                    + f"{self.colors['red']}Batched:     {batch_logw_nexts[i].weights}{self.colors['reset']}"
                )

            score = await self.score(context, *method_args)
            abs_diff, rel_diff = self._compute_diff(score, batch_scores[i])
            if abs_diff > atol or rel_diff > rtol:
                raise AssertionError(
                    f"{self.colors['red']}Batch score mismatch for context {context}:{self.colors['reset']}\n"
                    + f"{self.colors['green']}Non-batched: {score}\n"
                    + f"{self.colors['red']}Batched:     {batch_scores[i]}{self.colors['reset']}"
                )
            elif verbosity > 0:
                print(
                    f"{self.colors['green']}Batch score consistency satisfied for context {context}:{self.colors['reset']}"
                )
                print(
                    f"{self.colors['green']}Non-batched: {score}\n"
                    + f"{self.colors['green']}Batched:     {batch_scores[i]}{self.colors['reset']}\n"
                )

    def _compute_diff(self, want, have):
        is_inf = want == float("-inf") and have == float("-inf")
        abs_diff = 0 if is_inf else abs(want - have)
        if want == 0:
            rel_diff = 0 if have == 0 else float("inf")
        else:
            rel_diff = 0 if is_inf else abs((want - have) / want)
        return abs_diff, rel_diff

    def _format_diff(self, want, have, abs_diff, rel_diff, atol, rtol, token=None):
        abs_diff_str = (
            f"{self.colors['cyan']}Abs Diff: {abs_diff:.6f} <= {atol=}\033[0m"
        )
        rel_diff_str = (
            f"{self.colors['magenta']}Rel Diff: {rel_diff:.6f} <= {rtol=}\033[0m"
        )

        want_str = f"{self.colors['green']}Expected: {want:.6f}{self.colors['reset']}"
        have_clr = (
            self.colors["yellow"]
            if abs_diff <= atol and rel_diff <= rtol
            else self.colors["red"]
        )
        have_str = f"{have_clr}Actual:   {have:.6f}{self.colors['reset']}"

        if abs_diff <= atol:
            abs_diff_str = f"{self.colors['green']}Abs Diff: {abs_diff:.6f} <= {atol=}{self.colors['reset']}"
        else:
            abs_diff_str = f"{self.colors['red']}Abs Diff: {abs_diff:.6f} > {atol=}{self.colors['reset']}"

        if rel_diff <= rtol:
            rel_diff_str = f"{self.colors['green']}Rel Diff: {rel_diff:.6f} <= {rtol=}{self.colors['reset']}"
        else:
            rel_diff_str = f"{self.colors['red']}Rel Diff: {rel_diff:.6f} > {rtol=}{self.colors['reset']}"

        token_str = (
            f"{self.colors['blue']}Token:    {token}{self.colors['reset']}\n"
            if token
            else ""
        )
        return f"{token_str}{want_str}\n{have_str}\n{abs_diff_str}\n{rel_diff_str}\n\n"

assert_logw_next_consistency(context, rtol=0.001, atol=1e-05, top=None, verbosity=0, method_args=()) async

Assert that logw_next is consistent with prefix and complete.

For a context of tokens \(x_1, \ldots, x_{n-1}\), this checks (in log space) whether:

$$ \textsf{logw_next}(x_n | x_1, \ldots, x_{n-1}) = \textsf{score}(x_1, \ldots, x_n) - \textsf{prefix}(x_1, \ldots, x_{n-1}) $$ for \(x_n \in \textsf{vocab_eos}\), i.e., the potential's vocabulary and end-of-sequence token.

Parameters:

Name Type Description Default
context list

Context to test.

 required
rtol float

Relative tolerance for floating point comparison.

 0.001
atol float

Absolute tolerance for floating point comparison.

 1e-05
top int

If specified, only test the top-k tokens by log weight. If None, test all tokens.

 None
verbosity int

Verbosity level.

 0
method_args tuple

Positional arguments to pass to logw_next, prefix, and batch_score. Defaults to empty tuple.

 ()

Raises:

Type Description
AssertionError

If logw_next is not consistent with prefix and complete.
Source code in genlm/control/potential/testing.py 
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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
async def assert_logw_next_consistency(
    self, context, rtol=1e-3, atol=1e-5, top=None, verbosity=0, method_args=()
):
    """
    Assert that `logw_next` is consistent with `prefix` and `complete`.

    For a `context` of tokens $x_1, \\ldots, x_{n-1}$, this checks (in log space) whether:

    $$
    \\textsf{logw\\_next}(x_n | x_1, \\ldots, x_{n-1}) = \\textsf{score}(x_1, \\ldots, x_n) - \\textsf{prefix}(x_1, \\ldots, x_{n-1})
    $$
    for $x_n \\in \\textsf{vocab_eos}$, i.e., the potential's vocabulary and end-of-sequence token.

    Args:
        context (list): Context to test.
        rtol (float): Relative tolerance for floating point comparison.
        atol (float): Absolute tolerance for floating point comparison.
        top (int):If specified, only test the top-k tokens by log weight. If None, test all tokens.
        verbosity (int): Verbosity level.
        method_args (tuple): Positional arguments to pass to `logw_next`, `prefix`, and `batch_score`.
            Defaults to empty tuple.

    Raises:
        AssertionError: If `logw_next` is not consistent with `prefix` and `complete`.
    """
    top_logw_next = (await self.logw_next(context, *method_args)).materialize(
        top=top
    )
    tokens = list(top_logw_next.keys())
    extended = [[*context, x] for x in tokens]

    context_w = await self.prefix(context, *method_args)
    extended_ws = np.array(
        await asyncio.gather(*[self.score(e, *method_args) for e in extended])
    )

    wants = np.array([top_logw_next[x] for x in tokens])
    haves = extended_ws - context_w

    errors, valids = [], []
    for i, (want, have) in enumerate(zip(wants, haves)):
        abs_diff, rel_diff = self._compute_diff(want, have)
        info = (want, have, abs_diff, rel_diff, tokens[i])
        (valids if abs_diff <= atol and rel_diff <= rtol else errors).append(info)

    if valids and verbosity > 0:
        print(
            f"{self.colors['green']}logw_next consistency with context={context!r} satisfied for tokens:{self.colors['reset']}\n"
        )
        for valid in valids:
            want, have, abs_diff, rel_diff, token = valid
            print(
                self._format_diff(want, have, abs_diff, rel_diff, atol, rtol, token)
            )

    if errors:
        error_msg = f"{self.colors['red']}logw_next consistency with context={context!r} not satisfied for tokens:{self.colors['reset']}\n\n"
        for error in errors:
            want, have, abs_diff, rel_diff, token = error
            error_msg += self._format_diff(
                want, have, abs_diff, rel_diff, atol, rtol, token
            )
        raise AssertionError(error_msg)

assert_autoreg_fact(context, rtol=0.001, atol=1e-05, verbosity=0, method_args=()) async

Assert that complete factors as an autoregressive sum of logw_nexts.

For a context of tokens \(x_1, \ldots, x_n\), this checks (in log space) whether:

$$ \textsf{complete}(x_1, \ldots, x_n) - \textsf{prefix}(\epsilon) = \textsf{logw_next}(\textsf{eos} \mid x_1, \ldots, x_{n}) + \sum_{i=1}^{n} \textsf{logw_next}(x_i \mid x_1, \ldots, x_{i-1}) $$ where \(\epsilon\) is the empty sequence.

Parameters:

Name Type Description Default
context list

Context to test.

 required
rtol float

Relative tolerance for floating point comparison.

 0.001
atol float

Absolute tolerance for floating point comparison.

 1e-05
verbosity int

Verbosity level.

 0
method_args tuple

Positional arguments to pass to complete, prefix, and logw_next. Defaults to empty tuple.

 ()

Raises:

Type Description
AssertionError

If the autoregressive factorization is not satisfied.
Source code in genlm/control/potential/testing.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
async def assert_autoreg_fact(
    self, context, rtol=1e-3, atol=1e-5, verbosity=0, method_args=()
):
    """
    Assert that `complete` factors as an autoregressive sum of `logw_next`s.

    For a `context` of tokens $x_1, \\ldots, x_n$, this checks (in log space) whether:

    $$
    \\textsf{complete}(x_1, \\ldots, x_n) - \\textsf{prefix}(\\epsilon) = \\textsf{logw\\_next}(\\textsf{eos} \\mid x_1, \\ldots, x_{n}) + \\sum_{i=1}^{n} \\textsf{logw_next}(x_i \\mid x_1, \\ldots, x_{i-1})
    $$
    where $\\epsilon$ is the empty sequence.

    Args:
        context (list): Context to test.
        rtol (float): Relative tolerance for floating point comparison.
        atol (float): Absolute tolerance for floating point comparison.
        verbosity (int): Verbosity level.
        method_args (tuple): Positional arguments to pass to `complete`, `prefix`, and `logw_next`.
            Defaults to empty tuple.

    Raises:
        AssertionError: If the autoregressive factorization is not satisfied.
    """
    want = (await self.complete(context, *method_args)) - (
        await self.prefix([], *method_args)
    )

    logw_next_results = await asyncio.gather(
        *[self.logw_next(context[:i], *method_args) for i in range(len(context))],
        self.logw_next(context, *method_args),
    )

    have = (
        sum(logw_next_results[i][context[i]] for i in range(len(context)))
        + logw_next_results[-1][self.eos]
    )

    abs_diff, rel_diff = self._compute_diff(want, have)
    if abs_diff > atol or rel_diff > rtol:
        error_msg = (
            f"{self.colors['red']}Factorization not satisfied for context {context!r}:{self.colors['reset']}\n"
            + self._format_diff(want, have, abs_diff, rel_diff, atol, rtol)
        )
        raise AssertionError(error_msg)

    if verbosity > 0:
        print(
            f"{self.colors['green']}Factorization property satisfied for context {context}:{self.colors['reset']}\n"
        )
        print(self._format_diff(want, have, abs_diff, rel_diff, atol, rtol))

assert_batch_consistency(contexts, rtol=0.001, atol=1e-05, verbosity=0, batch_method_args=(), method_args=()) async

Assert that batch results are equal to non-batch results.

Parameters:

Name Type Description Default
contexts list[list]

Contexts to test.

 required
rtol float

Relative tolerance for floating point comparison.

 0.001
atol float

Absolute tolerance for floating point comparison.

 1e-05
verbosity int

Verbosity level.

 0
batch_method_args tuple

Positional arguments to pass to batch methods. Defaults to empty tuple.

 ()
method_args tuple

Positional arguments to pass to underlying potential methods. Defaults to empty tuple.

 ()

Raises:

Type Description
AssertionError

If the batch results are not equal to the non-batch results.
Source code in genlm/control/potential/testing.py 
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
async def assert_batch_consistency(
    self,
    contexts,
    rtol=1e-3,
    atol=1e-5,
    verbosity=0,
    batch_method_args=(),
    method_args=(),
):
    """
    Assert that batch results are equal to non-batch results.

    Args:
        contexts (list[list]): Contexts to test.
        rtol (float): Relative tolerance for floating point comparison.
        atol (float): Absolute tolerance for floating point comparison.
        verbosity (int): Verbosity level.
        batch_method_args (tuple): Positional arguments to pass to batch methods.
            Defaults to empty tuple.
        method_args (tuple): Positional arguments to pass to underlying potential methods.
            Defaults to empty tuple.

    Raises:
        AssertionError: If the batch results are not equal to the non-batch results.
    """
    batch_logw_nexts = await self.batch_logw_next(contexts, *batch_method_args)
    batch_scores = await self.batch_score(contexts, *batch_method_args)

    for i, context in enumerate(contexts):
        logw_next = await self.logw_next(context, *method_args)
        try:
            np.testing.assert_allclose(
                batch_logw_nexts[i].weights, logw_next.weights, rtol=rtol, atol=atol
            )
            if verbosity > 0:
                print(
                    f"{self.colors['green']}Batch logw_next consistency satisfied for context {context}:{self.colors['reset']}"
                )
                print(
                    f"{self.colors['green']}Non-batched: {logw_next.weights}\n"
                    + f"{self.colors['green']}Batched:     {batch_logw_nexts[i].weights}{self.colors['reset']}\n"
                )
        except AssertionError:
            raise AssertionError(
                f"{self.colors['red']}Batch logw_next mismatch for context {context}:{self.colors['reset']}\n"
                + f"{self.colors['green']}Non-batched: {logw_next.weights}\n"
                + f"{self.colors['red']}Batched:     {batch_logw_nexts[i].weights}{self.colors['reset']}"
            )

        score = await self.score(context, *method_args)
        abs_diff, rel_diff = self._compute_diff(score, batch_scores[i])
        if abs_diff > atol or rel_diff > rtol:
            raise AssertionError(
                f"{self.colors['red']}Batch score mismatch for context {context}:{self.colors['reset']}\n"
                + f"{self.colors['green']}Non-batched: {score}\n"
                + f"{self.colors['red']}Batched:     {batch_scores[i]}{self.colors['reset']}"
            )
        elif verbosity > 0:
            print(
                f"{self.colors['green']}Batch score consistency satisfied for context {context}:{self.colors['reset']}"
            )
            print(
                f"{self.colors['green']}Non-batched: {score}\n"
                + f"{self.colors['green']}Batched:     {batch_scores[i]}{self.colors['reset']}\n"
            )