Coverage for src/bartz/mcmcloop/_callback.py: 93%
258 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 11:03 +0000
1# bartz/src/bartz/mcmcloop/_callback.py
2#
3# Copyright (c) 2024-2026, The Bartz Contributors
4#
5# This file is part of bartz.
6#
7# Permission is hereby granted, free of charge, to any person obtaining a copy
8# of this software and associated documentation files (the "Software"), to deal
9# in the Software without restriction, including without limitation the rights
10# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
11# copies of the Software, and to permit persons to whom the Software is
12# furnished to do so, subject to the following conditions:
13#
14# The above copyright notice and this permission notice shall be included in all
15# copies or substantial portions of the Software.
16#
17# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
20# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
22# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
23# SOFTWARE.
25"""Progress-reporting callbacks for `run_mcmc`."""
27import itertools
28from collections.abc import Callable
29from dataclasses import dataclass, replace
30from functools import partial, wraps
31from typing import Any, Literal, TypeVar
33import numpy
34from equinox import Module, field
35from jax import debug, eval_shape, lax, random, tree
36from jax import numpy as jnp
37from jax.scipy.special import logsumexp
38from jaxtyping import Array, ArrayLike, Bool, Float32, Int32, Integer, PyTree, Shaped
39from tqdm.auto import tqdm
41from bartz._typing import kwdict
42from bartz.grove import forest_mean_leaves
43from bartz.mcmcloop._loop import Callback, _replicate
44from bartz.mcmcstep import State
45from bartz.mcmcstep._axes import chain_to_axis, chain_vmap_axes, chainful_axis
48class CallbackTuple(Callback):
49 """Callback that runs a sequence of callbacks in order.
51 Each callback receives the state as (possibly) updated by the preceding
52 ones, so their effects compose left to right.
53 """
55 callbacks: tuple[Callback, ...]
56 """The callbacks to run, in order."""
58 def __call__(self, *, state: State, **kwargs: Any) -> tuple[State, 'CallbackTuple']:
59 """Invoke each callback in turn, threading the state through them."""
60 new_callbacks = []
61 for i, callback in enumerate(self.callbacks):
62 # decorrelate the keys by position, so a callback keeps its key if
63 # other callbacks are appended after it
64 kw = dict(kwargs, key=random.fold_in(kwargs['key'], i))
65 rt = callback(state=state, **kw)
66 if rt is None:
67 new_callbacks.append(callback)
68 else:
69 state, new_callback = rt
70 new_callbacks.append(new_callback)
71 return state, replace(self, callbacks=tuple(new_callbacks))
74class StatsReport(Module):
75 """Forest diagnostics produced by `StatsAccumulator.report` for one report."""
77 grow_prop: Float32[Array, '']
78 """Fraction of trees proposed for a grow move."""
80 move_acc: Float32[Array, '']
81 """Fraction of trees on which a grow or prune move was accepted."""
83 mean_leaves: Float32[Array, '']
84 """Mean number of leaves per tree."""
86 peff: Float32[Array, ''] | None
87 """Effective number of predictors, or `None` when variable selection is off."""
89 n_samples: Int32[Array, ''] | None
90 """Number of iterations averaged over, or `None` when not averaging."""
92 num_chains: int | None = field(static=True)
93 """Number of chains averaged over, or `None` when single-chain."""
95 max_leaves: int = field(static=True)
96 """Maximum possible number of leaves per tree."""
98 p: int | None = field(static=True)
99 """Number of predictors, or `None` when variable selection is off."""
102class StatsAccumulator(Module):
103 """Running average of the forest diagnostics shown during the MCMC.
105 When enabled, it sums the per-iteration statistics so a report shows their
106 average over the iterations since the previous report. When disabled it
107 carries no running state and a report shows the latest iteration only.
108 """
110 sums: dict[str, Float32[Array, '']] | None
111 """Running sums of the averaged statistics, or `None` when disabled."""
113 count: Int32[Array, '']
114 """Number of iterations accumulated since the last reset."""
116 @classmethod
117 def initial(cls, state: State, *, enabled: bool) -> 'StatsAccumulator':
118 """Create a zeroed accumulator, inert unless `enabled`."""
119 if enabled:
120 # only the structure is needed, so avoid computing the statistics
121 shapes = eval_shape(cls._avg_stats, state)
122 sums = tree.map(lambda s: jnp.zeros(s.shape, s.dtype), shapes)
123 else:
124 sums = None
125 return cls(sums=sums, count=jnp.int32(0))
127 def update(self, state: State) -> 'StatsAccumulator':
128 """Add the latest iteration's statistics; no-op when disabled."""
129 if self.sums is None:
130 return self
131 sums = tree.map(jnp.add, self.sums, self._avg_stats(state))
132 return replace(self, sums=sums, count=self.count + 1)
134 def reset_if(self, cond: bool | Bool[Array, '']) -> 'StatsAccumulator':
135 """Zero the running sums where `cond` holds; no-op when disabled."""
136 if self.sums is None:
137 return self
138 sums = tree.map(lambda s: jnp.where(cond, 0, s), self.sums)
139 return replace(self, sums=sums, count=jnp.where(cond, 0, self.count))
141 def report(self, state: State) -> StatsReport:
142 """Statistics to display: the windowed average if enabled, else the latest."""
143 if self.sums is None:
144 averaged: kwdict = self._avg_stats(state)
145 n_samples = None
146 else:
147 averaged = tree.map(lambda s: s / self.count, self.sums)
148 n_samples = self.count
149 return StatsReport(**averaged, **self._static_stats(state), n_samples=n_samples)
151 @staticmethod
152 def _avg_stats(state: State) -> dict[str, Float32[Array, ''] | None]:
153 """Per-iteration diagnostics that are averaged over the report window."""
154 forest = state.forest
155 chain_axis = chain_vmap_axes(forest).split_tree
156 num_trees_axis = chainful_axis(0, chain_axis) # (num_trees, hts)
157 split_tree = chain_to_axis(forest.split_tree, chain_axis)
158 prop_total = forest.split_tree.shape[num_trees_axis]
160 log_s = forest.log_s
161 if log_s is None:
162 peff = None
163 else:
164 log_s = chain_to_axis(log_s, chain_vmap_axes(forest).log_s)
165 peff = StatsAccumulator._effective_predictors(log_s)
167 return dict(
168 grow_prop=forest.grow_prop_count.mean() / prop_total,
169 move_acc=(forest.grow_acc_count.mean() + forest.prune_acc_count.mean())
170 / prop_total,
171 mean_leaves=forest_mean_leaves(split_tree),
172 peff=peff,
173 )
175 @staticmethod
176 def _static_stats(state: State) -> dict[str, int | None]:
177 """Per-iteration diagnostics shown as-is, constant over the run."""
178 forest = state.forest
179 split_tree = chain_to_axis(
180 forest.split_tree, chain_vmap_axes(forest).split_tree
181 )
182 log_s = forest.log_s
183 if log_s is None:
184 p = None
185 else:
186 *_, p = chain_to_axis(log_s, chain_vmap_axes(forest).log_s).shape
187 return dict(num_chains=state.num_chains(), max_leaves=split_tree.shape[-1], p=p)
189 @staticmethod
190 def _effective_predictors(log_s: Float32[Array, '*chains p']) -> Float32[Array, '']:
191 """Effective number of predictors used for splitting across all chains.
193 Perplexity (exponential of the Shannon entropy) of the split-variable
194 distribution ``s = softmax(log_s)`` pooled (averaged) over chains. It is
195 1 when all chains concentrate on a single shared predictor and ``p`` when
196 the pooled distribution is uniform; in general a pooled distribution
197 spread evenly over ``k`` predictors gives ``k``. Chains are pooled before
198 taking the entropy because predictions average over all chains, so a
199 predictor used by any chain counts as used.
200 """
201 *_, p = log_s.shape
202 # normalize each chain
203 log_prob = log_s - logsumexp(log_s, axis=-1, keepdims=True)
204 per_chain = log_prob.reshape(-1, p)
205 num_chains, _ = per_chain.shape
206 # mix over chains. WORKAROUND(jax<0.7.1): once we bump jax to v0.7.1
207 # this is `jax.nn.logmeanexp(per_chain, axis=0)`
208 log_pool = logsumexp(per_chain, axis=0) - jnp.log(num_chains)
209 prob = jnp.exp(log_pool)
210 # the where avoids the 0 * -inf = nan term where a probability is 0, the
211 # same guard `jax.scipy.special.entr` uses, but reusing the log we have
212 entropy = -jnp.sum(prob * jnp.where(prob, log_pool, 1.0))
213 return jnp.exp(entropy)
216def make_print_callback(
217 state: State,
218 *,
219 dot_every: int | Integer[Array, ''] | None = 1,
220 report_every: int | Integer[Array, ''] | None = 100,
221 average: bool = True,
222) -> 'PrintCallback':
223 """
224 Prepare a progress-printing callback for `run_mcmc`.
226 The callback prints a dot on every iteration, and a longer report
227 periodically.
229 Parameters
230 ----------
231 state
232 The MCMC state to use the callback with, used to determine device
233 sharding.
234 dot_every
235 A dot is printed every `dot_every` MCMC iterations, `None` to disable.
236 report_every
237 A one line report is printed every `report_every` MCMC iterations,
238 `None` to disable.
239 average
240 If `True`, the reported statistics are averaged over the iterations
241 since the previous report; if `False`, they reflect the current
242 iteration only. Ignored when `report_every` is `None`.
244 Returns
245 -------
246 A `PrintCallback` to pass as the `callback` argument of `run_mcmc`.
248 Examples
249 --------
250 >>> run_mcmc(key, state, ..., callback=make_print_callback(state, ...))
251 """
253 def as_replicated_array_or_none(
254 val: Shaped[ArrayLike, '*shape'] | None,
255 ) -> None | Shaped[Array, '*shape']:
256 return None if val is None else _replicate(jnp.asarray(val), state.config.mesh)
258 accumulator = tree.map(
259 partial(_replicate, mesh=state.config.mesh),
260 StatsAccumulator.initial(state, enabled=average and report_every is not None),
261 )
263 return PrintCallback(
264 as_replicated_array_or_none(dot_every),
265 as_replicated_array_or_none(report_every),
266 accumulator,
267 )
270class PrintCallback(Callback):
271 """Progress-printing callback for `run_mcmc`, see `make_print_callback`."""
273 dot_every: Int32[Array, ''] | None
274 """A dot is printed every `dot_every` MCMC iterations, `None` to disable."""
276 report_every: Int32[Array, ''] | None
277 """A one line report is printed every `report_every` MCMC iterations,
278 `None` to disable."""
280 accumulator: StatsAccumulator
281 """Running average of the reported statistics, inert unless averaging."""
283 def __call__(
284 self,
285 *,
286 state: State,
287 burnin: Bool[Array, ''],
288 i_total: Int32[Array, ''],
289 n_burn: Int32[Array, ''],
290 n_save: Int32[Array, ''],
291 n_skip: Int32[Array, ''],
292 **_: Any,
293 ) -> tuple[State, 'PrintCallback']:
294 """Print a dot and/or a report periodically during the MCMC."""
295 report_every = self.report_every
296 dot_every = self.dot_every
297 it = i_total + 1
299 accumulator = self.accumulator.update(state)
301 def get_cond(every: Int32[Array, ''] | None) -> bool | Bool[Array, '']:
302 return False if every is None else it % every == 0
304 report_cond = get_cond(report_every)
305 dot_cond = get_cond(dot_every)
307 def line_report_branch() -> None:
308 if report_every is None: 308 ↛ 309line 308 didn't jump to line 309 because the condition on line 308 was never true
309 return
310 if dot_every is None: 310 ↛ 311line 310 didn't jump to line 311 because the condition on line 310 was never true
311 print_newline = False
312 else:
313 print_newline = it % report_every > it % dot_every
314 debug.callback(
315 _print_report,
316 accumulator.report(state),
317 print_dot=dot_cond,
318 print_newline=print_newline,
319 burnin=burnin,
320 it=it,
321 n_iters=n_burn + n_save * n_skip,
322 )
324 def just_dot_branch() -> None:
325 if dot_every is None: 325 ↛ 326line 325 didn't jump to line 326 because the condition on line 325 was never true
326 return
327 # terminate the dot line on the final iteration so subsequent output
328 # doesn't continue on the same line as the dots
329 last_iter = it == n_burn + n_save * n_skip
330 lax.cond(
331 last_iter,
332 lambda: debug.callback(lambda: print('.', flush=True)), # noqa: T201
333 lambda: debug.callback(lambda: print('.', end='', flush=True)), # noqa: T201
334 )
335 # logging can't do in-line printing so we use print
337 lax.cond(
338 report_cond,
339 line_report_branch,
340 lambda: lax.cond(dot_cond, just_dot_branch, lambda: None),
341 )
343 accumulator = accumulator.reset_if(report_cond)
344 return state, replace(self, accumulator=accumulator)
347def make_tqdm_callback(
348 state: State,
349 *,
350 update_every: int = 1,
351 report_every: int | None = 100,
352 average: bool = True,
353 **tqdm_kwargs: Any,
354) -> 'TqdmCallback':
355 """
356 Prepare a `tqdm` progress-bar callback for `run_mcmc`.
358 The callback shows a progress bar that advances with the MCMC iterations,
359 optionally annotated with the proposal acceptance statistics.
361 Parameters
362 ----------
363 state
364 The MCMC state to use the callback with, used to determine device
365 sharding.
366 update_every
367 The bar position is refreshed every `update_every` MCMC iterations
368 (`tqdm` further throttles the actual redraw rate on its own).
369 report_every
370 The acceptance statistics shown next to the bar are refreshed every
371 `report_every` MCMC iterations, `None` to omit them.
372 average
373 If `True`, the statistics shown are averaged over the iterations since
374 the previous refresh; if `False`, they reflect the current iteration
375 only. Ignored when `report_every` is `None`.
376 **tqdm_kwargs
377 Additional keyword arguments forwarded to the `tqdm.tqdm` constructor,
378 e.g., ``desc``, ``file``, or ``disable``.
380 Returns
381 -------
382 A `TqdmCallback` to pass as the `callback` argument of `run_mcmc`.
384 Notes
385 -----
386 Works with chains sharded across multiple devices. If the run is interrupted
387 (e.g. with ^C), the bar is left as-is; the next `make_tqdm_callback` call
388 closes it, so a subsequent run starts from a clean line.
390 Examples
391 --------
392 >>> run_mcmc(key, state, ..., callback=make_tqdm_callback(state, ...))
393 """
394 _close_stale_bars() # clean up after any previous run that was interrupted
395 bar_id = next(_TQDM_BAR_COUNTER)
396 _TQDM_REGISTRY[bar_id] = _TqdmEntry(tqdm_kwargs)
398 def as_replicated_array(
399 val: Shaped[ArrayLike, '*shape'],
400 ) -> Shaped[Array, '*shape']:
401 return _replicate(jnp.asarray(val), state.config.mesh)
403 return TqdmCallback(
404 bar_id=as_replicated_array(jnp.int32(bar_id)),
405 update_every=as_replicated_array(jnp.int32(update_every)),
406 report_every=None
407 if report_every is None
408 else as_replicated_array(jnp.int32(report_every)),
409 accumulator=tree.map(
410 partial(_replicate, mesh=state.config.mesh),
411 StatsAccumulator.initial(
412 state, enabled=average and report_every is not None
413 ),
414 ),
415 )
418class TqdmCallback(Callback):
419 """`tqdm` progress-bar callback for `run_mcmc`, see `make_tqdm_callback`."""
421 bar_id: Int32[Array, '']
422 """Handle identifying the bar in the module-level `tqdm` bar registry."""
424 update_every: Int32[Array, '']
425 """The bar position is refreshed every `update_every` MCMC iterations."""
427 report_every: Int32[Array, ''] | None
428 """The acceptance statistics are refreshed every `report_every` MCMC
429 iterations, `None` to omit them."""
431 accumulator: StatsAccumulator
432 """Running average of the reported statistics, inert unless averaging."""
434 def __call__(
435 self,
436 *,
437 state: State,
438 i_total: Int32[Array, ''],
439 n_burn: Int32[Array, ''],
440 n_save: Int32[Array, ''],
441 n_skip: Int32[Array, ''],
442 **_: Any,
443 ) -> tuple[State, 'TqdmCallback']:
444 """Advance a `tqdm` progress bar during the MCMC."""
445 it = i_total + 1
446 n_iters = n_burn + n_save * n_skip
447 bar_id = self.bar_id
448 last = it == n_iters
450 accumulator = self.accumulator.update(state)
452 # The callbacks are unordered: `ordered=True` is unsupported with more
453 # than one device, and we need this to work with chains sharded across
454 # devices. `_tqdm_advance` is therefore robust to out-of-order
455 # invocations.
457 # refresh the statistics first so they tend to be visible by the time
458 # the bar is advanced
459 report_every = self.report_every
460 if report_every is not None: 460 ↛ 469line 460 didn't jump to line 469 because the condition on line 460 was always true
461 report_cond = (it % report_every == 0) | last
463 def report_branch() -> None:
464 debug.callback(_tqdm_report, accumulator.report(state), bar_id, n_iters)
466 lax.cond(report_cond, report_branch, lambda: None)
467 accumulator = accumulator.reset_if(report_cond)
469 lax.cond(
470 (it % self.update_every == 0) | last,
471 lambda: debug.callback(_tqdm_advance, bar_id, it, n_iters),
472 lambda: None,
473 )
475 return state, replace(self, accumulator=accumulator)
478T = TypeVar('T')
481def _convert_jax_arrays_in_args(func: Callable[..., T]) -> Callable[..., T]:
482 """Remove jax arrays from a function arguments.
484 Converts all `jax.Array` instances in the arguments to either Python scalars
485 or numpy arrays.
486 """
488 def convert_jax_arrays(pytree: PyTree) -> PyTree:
489 def convert_jax_array(val: object) -> object:
490 if not isinstance(val, Array):
491 return val
492 elif val.shape: 492 ↛ 493line 492 didn't jump to line 493 because the condition on line 492 was never true
493 return numpy.array(val)
494 else:
495 return val.item()
497 return tree.map(convert_jax_array, pytree)
499 @wraps(func)
500 def new_func(*args: Any, **kw: Any) -> T:
501 args = convert_jax_arrays(args)
502 kw = convert_jax_arrays(kw)
503 return func(*args, **kw)
505 return new_func
508@_convert_jax_arrays_in_args
509# convert all jax arrays in arguments because operations on them could lead to
510# deadlock with the main thread
511def _print_report(
512 report: StatsReport,
513 *,
514 print_dot: bool,
515 print_newline: bool,
516 burnin: bool,
517 it: int,
518 n_iters: int,
519) -> None:
520 """Print the report for `PrintCallback`."""
521 # determine prefix
522 if print_dot: 522 ↛ 524line 522 didn't jump to line 524 because the condition on line 522 was always true
523 prefix = '.\n'
524 elif print_newline:
525 prefix = '\n'
526 else:
527 prefix = ''
529 # determine suffix in parentheses: what the statistics are averaged over
530 avg_over = []
531 if report.num_chains is not None:
532 avg_over.append(f'{report.num_chains} chains')
533 if report.n_samples is not None: 533 ↛ 535line 533 didn't jump to line 535 because the condition on line 533 was always true
534 avg_over.append(f'{report.n_samples} samples')
535 msgs = []
536 if avg_over: 536 ↛ 538line 536 didn't jump to line 538 because the condition on line 536 was always true
537 msgs.append('avg. ' + ' x '.join(avg_over))
538 if burnin: 538 ↛ 539line 538 didn't jump to line 539 because the condition on line 538 was never true
539 msgs.append('burnin')
540 suffix = f' ({", ".join(msgs)})' if msgs else ''
542 # variable-selection concentration, only shown when it is enabled
543 if report.peff is None:
544 var_msg = ''
545 else:
546 var_msg = f'var: {report.peff:.1f}/{report.p}, '
548 print( # noqa: T201, see PrintCallback for why not logging
549 f'{prefix}Iteration {it}/{n_iters}, '
550 f'grow prob: {report.grow_prop:.0%}, '
551 f'move acc: {report.move_acc:.0%}, '
552 f'{var_msg}'
553 f'leaves: {report.mean_leaves:.1f}/{report.max_leaves}{suffix}'
554 )
557@dataclass(frozen=True)
558class _TqdmEntry:
559 """An entry in the `tqdm` bar registry."""
561 kwargs: dict[str, Any]
562 """Keyword arguments to construct the bar with, from `make_tqdm_callback`."""
564 bar: tqdm | None = None
565 """The bar, created lazily on the first callback invocation, `None` until then."""
568# tqdm carries Python state that cannot live in a jax pytree, so the bars are
569# kept here and referenced from the jax loop through the integer handle stored
570# in `TqdmCallback.bar_id` (a traceable scalar, so the loop pytree stays
571# stable across runs and is not recompiled).
572_TQDM_REGISTRY: dict[int, _TqdmEntry] = {}
573_TQDM_BAR_COUNTER = itertools.count()
575# tqdm's default layout, but without the ': ' that `format_meter` forces after a
576# non-empty description; the label is set as a `{desc}` ending in a space instead
577_TQDM_BAR_FORMAT = (
578 '{desc}{percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} '
579 '[{elapsed}<{remaining}, {rate_fmt}{postfix}]'
580)
583def _close_stale_bars() -> None:
584 """Close and drop any bars left over from a previous (e.g. interrupted) run."""
585 for entry in _TQDM_REGISTRY.values():
586 if entry.bar is not None:
587 entry.bar.close()
588 _TQDM_REGISTRY.clear()
591def _get_or_create_bar(bar_id: int, n_iters: int) -> tqdm | None:
592 """Return the bar for `bar_id`, creating it on first use, `None` if finished."""
593 entry = _TQDM_REGISTRY.get(bar_id)
594 if entry is None:
595 # the bar was already closed (the loop finished, possibly out of order)
596 return None
597 if entry.bar is None:
598 bar = tqdm(**{'total': n_iters, 'bar_format': _TQDM_BAR_FORMAT, **entry.kwargs})
599 _TQDM_REGISTRY[bar_id] = replace(entry, bar=bar)
600 return bar
601 return entry.bar
604@_convert_jax_arrays_in_args
605# convert all jax arrays in arguments, see _print_report for why
606def _tqdm_advance(bar_id: int, it: int, n_iters: int) -> None:
607 """Advance the bar towards absolute position `it`, closing it at the end."""
608 bar = _get_or_create_bar(bar_id, n_iters)
609 if bar is None: 609 ↛ 610line 609 didn't jump to line 610 because the condition on line 609 was never true
610 return
611 bar.update(max(0, it - bar.n)) # forward-only: callbacks may arrive out of order
612 if it >= n_iters:
613 bar.close()
614 del _TQDM_REGISTRY[bar_id]
617@_convert_jax_arrays_in_args
618# convert all jax arrays in arguments, see _print_report for why
619def _tqdm_report(report: StatsReport, bar_id: int, n_iters: int) -> None:
620 """Set the bar description and acceptance-statistics postfix."""
621 bar = _get_or_create_bar(bar_id, n_iters)
622 if bar is None:
623 return
624 # set_description_str (not set_description) to avoid tqdm's ': ' suffix; the
625 # trailing space separates the label from the bar
626 bar.set_description_str('train ', refresh=False)
627 # keep this terse so the bar stays narrow, e.g. '4ch 100sa acc 25% leaves 3.4/32'
628 msgs = []
629 if report.num_chains is not None:
630 msgs.append(f'{report.num_chains}ch')
631 if report.n_samples is not None:
632 msgs.append(f'{report.n_samples}sa')
633 msgs.append(f'acc {report.move_acc:.0%}')
634 if report.peff is not None:
635 msgs.append(f'var {report.peff:.1f}/{report.p}')
636 msgs.append(f'leaves {report.mean_leaves:.1f}/{report.max_leaves}')
637 bar.set_postfix_str(' '.join(msgs))
640class CheckPlatformCallback(Callback):
641 """Check the given platform matches the actual platform at runtime."""
643 platform: Literal['cpu', 'gpu'] = field(static=True)
644 """The expected platform."""
646 def __call__(self, *, i_total: Int32[Array, ''], **_: Any) -> None:
647 """Check the platform on the first iteration, raise on mismatch."""
649 def check() -> None:
650 lax.platform_dependent(
651 cpu=partial(_check_platform, 'cpu', self.platform),
652 cuda=partial(_check_platform, 'gpu', self.platform),
653 )
655 lax.cond(i_total == 0, check, lambda: None)
658def _check_platform(actual_platform: str, expected_platform: str) -> None:
659 """Raise from a debug callback if the platform differs from expected."""
661 def raise_if_mismatch() -> None:
662 if actual_platform != expected_platform:
663 msg = (
664 f'`Bart` deduced the platform as {expected_platform!r}, '
665 'but the MCMC is running on '
666 f'{actual_platform!r}, so the automatic deduction was wrong; please '
667 'tell `Bart` what the correct platform is by setting the `devices` '
668 'argument.'
669 )
670 raise RuntimeError(msg)
672 debug.callback(raise_if_mismatch)