Coverage for src/bartz/mcmcloop/_callback.py: 93%
257 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-10-15 10:04 +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 and at the first one (`tqdm` further throttles the actual redraw rate
369 on its own).
370 report_every
371 The acceptance statistics shown next to the bar are refreshed every
372 `report_every` MCMC iterations, `None` to omit them.
373 average
374 If `True`, the statistics shown are averaged over the iterations since
375 the previous refresh; if `False`, they reflect the current iteration
376 only. Ignored when `report_every` is `None`.
377 **tqdm_kwargs
378 Additional keyword arguments forwarded to the `tqdm.tqdm` constructor,
379 e.g., ``desc``, ``file``, or ``disable``.
381 Returns
382 -------
383 A `TqdmCallback` to pass as the `callback` argument of `run_mcmc`.
385 Notes
386 -----
387 Works with chains sharded across multiple devices. If the run is interrupted
388 (e.g. with ^C), the bar is left as-is; the next `make_tqdm_callback` call
389 closes it, so a subsequent run starts from a clean line.
391 Examples
392 --------
393 >>> run_mcmc(key, state, ..., callback=make_tqdm_callback(state, ...))
394 """
395 _close_stale_bars() # clean up after any previous run that was interrupted
396 bar_id = next(_TQDM_BAR_COUNTER)
397 _TQDM_REGISTRY[bar_id] = _TqdmEntry(tqdm_kwargs)
399 def as_replicated_array(
400 val: Shaped[ArrayLike, '*shape'],
401 ) -> Shaped[Array, '*shape']:
402 return _replicate(jnp.asarray(val), state.config.mesh)
404 return TqdmCallback(
405 bar_id=as_replicated_array(jnp.int32(bar_id)),
406 update_every=as_replicated_array(jnp.int32(update_every)),
407 report_every=None
408 if report_every is None
409 else as_replicated_array(jnp.int32(report_every)),
410 accumulator=tree.map(
411 partial(_replicate, mesh=state.config.mesh),
412 StatsAccumulator.initial(
413 state, enabled=average and report_every is not None
414 ),
415 ),
416 )
419class TqdmCallback(Callback):
420 """`tqdm` progress-bar callback for `run_mcmc`, see `make_tqdm_callback`."""
422 bar_id: Int32[Array, '']
423 """Handle identifying the bar in the module-level `tqdm` bar registry."""
425 update_every: Int32[Array, '']
426 """The bar position is refreshed every `update_every` MCMC iterations."""
428 report_every: Int32[Array, ''] | None
429 """The acceptance statistics are refreshed every `report_every` MCMC
430 iterations, `None` to omit them."""
432 accumulator: StatsAccumulator
433 """Running average of the reported statistics, inert unless averaging."""
435 def __call__(
436 self,
437 *,
438 state: State,
439 i_total: Int32[Array, ''],
440 n_burn: Int32[Array, ''],
441 n_save: Int32[Array, ''],
442 n_skip: Int32[Array, ''],
443 **_: Any,
444 ) -> tuple[State, 'TqdmCallback']:
445 """Advance a `tqdm` progress bar during the MCMC."""
446 it = i_total + 1
447 n_iters = n_burn + n_save * n_skip
448 bar_id = self.bar_id
449 last = it == n_iters
451 accumulator = self.accumulator.update(state)
453 # The callbacks are unordered: `ordered=True` is unsupported with more
454 # than one device, and we need this to work with chains sharded across
455 # devices. `_tqdm_advance` is therefore robust to out-of-order
456 # invocations.
458 # refresh the statistics first so they tend to be visible by the time
459 # the bar is advanced
460 report_every = self.report_every
461 if report_every is not None: 461 ↛ 470line 461 didn't jump to line 470 because the condition on line 461 was always true
462 report_cond = (it % report_every == 0) | last
464 def report_branch() -> None:
465 debug.callback(_tqdm_report, accumulator.report(state), bar_id, n_iters)
467 lax.cond(report_cond, report_branch, lambda: None)
468 accumulator = accumulator.reset_if(report_cond)
470 lax.cond(
471 (it == 1) | (it % self.update_every == 0) | last,
472 lambda: debug.callback(_tqdm_advance, bar_id, it, n_iters),
473 lambda: None,
474 )
476 return state, replace(self, accumulator=accumulator)
479T = TypeVar('T')
482def _convert_jax_arrays_in_args(func: Callable[..., T]) -> Callable[..., T]:
483 """Remove jax arrays from a function arguments.
485 Converts all `jax.Array` instances in the arguments to either Python scalars
486 or numpy arrays.
487 """
489 def convert_jax_arrays(pytree: PyTree) -> PyTree:
490 def convert_jax_array(val: object) -> object:
491 if not isinstance(val, Array):
492 return val
493 elif val.shape: 493 ↛ 494line 493 didn't jump to line 494 because the condition on line 493 was never true
494 return numpy.array(val)
495 else:
496 return val.item()
498 return tree.map(convert_jax_array, pytree)
500 @wraps(func)
501 def new_func(*args: Any, **kw: Any) -> T:
502 args = convert_jax_arrays(args)
503 kw = convert_jax_arrays(kw)
504 return func(*args, **kw)
506 return new_func
509@_convert_jax_arrays_in_args
510# convert all jax arrays in arguments because operations on them could lead to
511# deadlock with the main thread
512def _print_report(
513 report: StatsReport,
514 *,
515 print_dot: bool,
516 print_newline: bool,
517 burnin: bool,
518 it: int,
519 n_iters: int,
520) -> None:
521 """Print the report for `PrintCallback`."""
522 # determine prefix
523 if print_dot: 523 ↛ 525line 523 didn't jump to line 525 because the condition on line 523 was always true
524 prefix = '.\n'
525 elif print_newline:
526 prefix = '\n'
527 else:
528 prefix = ''
530 # determine suffix in parentheses: what the statistics are averaged over
531 avg_over = []
532 if report.num_chains is not None:
533 avg_over.append(f'{report.num_chains} chains')
534 if report.n_samples is not None: 534 ↛ 536line 534 didn't jump to line 536 because the condition on line 534 was always true
535 avg_over.append(f'{report.n_samples} samples')
536 msgs = []
537 if avg_over: 537 ↛ 539line 537 didn't jump to line 539 because the condition on line 537 was always true
538 msgs.append('avg. ' + ' x '.join(avg_over))
539 if burnin: 539 ↛ 540line 539 didn't jump to line 540 because the condition on line 539 was never true
540 msgs.append('burnin')
541 suffix = f' ({", ".join(msgs)})' if msgs else ''
543 # variable-selection concentration, only shown when it is enabled
544 if report.peff is None:
545 var_msg = ''
546 else:
547 var_msg = f'var: {report.peff:.1f}/{report.p}, '
549 print( # noqa: T201, see PrintCallback for why not logging
550 f'{prefix}Iteration {it}/{n_iters}, '
551 f'grow prob: {report.grow_prop:.0%}, '
552 f'move acc: {report.move_acc:.0%}, '
553 f'{var_msg}'
554 f'leaves: {report.mean_leaves:.1f}/{report.max_leaves}{suffix}'
555 )
558@dataclass(frozen=True)
559class _TqdmEntry:
560 """An entry in the `tqdm` bar registry."""
562 kwargs: dict[str, Any]
563 """Keyword arguments to construct the bar with, from `make_tqdm_callback`."""
565 bar: tqdm | None = None
566 """The bar, created lazily on the first callback invocation, `None` until then."""
569# tqdm carries Python state that cannot live in a jax pytree, so the bars are
570# kept here and referenced from the jax loop through the integer handle stored
571# in `TqdmCallback.bar_id` (a traceable scalar, so the loop pytree stays
572# stable across runs and is not recompiled).
573_TQDM_REGISTRY: dict[int, _TqdmEntry] = {}
574_TQDM_BAR_COUNTER = itertools.count()
576# tqdm's default layout, but without the ': ' that `format_meter` forces after a
577# non-empty description; the label is set as a `{desc}` ending in a space instead
578_TQDM_BAR_FORMAT = (
579 '{desc}{percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} '
580 '[{elapsed}<{remaining}, {rate_fmt}{postfix}]'
581)
584def _close_stale_bars() -> None:
585 """Close and drop any bars left over from a previous (e.g. interrupted) run."""
586 for entry in _TQDM_REGISTRY.values():
587 if entry.bar is not None:
588 entry.bar.close()
589 _TQDM_REGISTRY.clear()
592def _get_or_create_bar(bar_id: int, n_iters: int) -> tqdm | None:
593 """Return the bar for `bar_id`, creating it on first use, `None` if finished."""
594 entry = _TQDM_REGISTRY.get(bar_id)
595 if entry is None:
596 # the bar was already closed (the loop finished, possibly out of order)
597 return None
598 if entry.bar is None:
599 bar = tqdm(
600 **{
601 'total': n_iters,
602 'bar_format': _TQDM_BAR_FORMAT,
603 'desc': 'train ',
604 **entry.kwargs,
605 }
606 )
607 _TQDM_REGISTRY[bar_id] = replace(entry, bar=bar)
608 return bar
609 return entry.bar
612@_convert_jax_arrays_in_args
613# convert all jax arrays in arguments, see _print_report for why
614def _tqdm_advance(bar_id: int, it: int, n_iters: int) -> None:
615 """Advance the bar towards absolute position `it`, closing it at the end."""
616 bar = _get_or_create_bar(bar_id, n_iters)
617 if bar is None: 617 ↛ 618line 617 didn't jump to line 618 because the condition on line 617 was never true
618 return
619 bar.update(max(0, it - bar.n)) # forward-only: callbacks may arrive out of order
620 if it >= n_iters:
621 bar.close()
622 del _TQDM_REGISTRY[bar_id]
625@_convert_jax_arrays_in_args
626# convert all jax arrays in arguments, see _print_report for why
627def _tqdm_report(report: StatsReport, bar_id: int, n_iters: int) -> None:
628 """Set the acceptance-statistics postfix of the bar."""
629 bar = _get_or_create_bar(bar_id, n_iters)
630 if bar is None:
631 return
632 # keep this terse so the bar stays narrow, e.g. '4ch 100sa acc 25% leaves 3.4/32'
633 msgs = []
634 if report.num_chains is not None:
635 msgs.append(f'{report.num_chains}ch')
636 if report.n_samples is not None:
637 msgs.append(f'{report.n_samples}sa')
638 msgs.append(f'acc {report.move_acc:.0%}')
639 if report.peff is not None:
640 msgs.append(f'var {report.peff:.1f}/{report.p}')
641 msgs.append(f'leaves {report.mean_leaves:.1f}/{report.max_leaves}')
642 bar.set_postfix_str(' '.join(msgs))
645class CheckPlatformCallback(Callback):
646 """Check the given platform matches the actual platform at runtime."""
648 platform: Literal['cpu', 'gpu'] = field(static=True)
649 """The expected platform."""
651 def __call__(self, *, i_total: Int32[Array, ''], **_: Any) -> None:
652 """Check the platform on the first iteration, raise on mismatch."""
654 def check() -> None:
655 lax.platform_dependent(
656 cpu=partial(_check_platform, 'cpu', self.platform),
657 cuda=partial(_check_platform, 'gpu', self.platform),
658 )
660 lax.cond(i_total == 0, check, lambda: None)
663def _check_platform(actual_platform: str, expected_platform: str) -> None:
664 """Raise from a debug callback if the platform differs from expected."""
666 def raise_if_mismatch() -> None:
667 if actual_platform != expected_platform:
668 msg = (
669 f'`Bart` deduced the platform as {expected_platform!r}, '
670 'but the MCMC is running on '
671 f'{actual_platform!r}, so the automatic deduction was wrong; please '
672 'tell `Bart` what the correct platform is by setting the `devices` '
673 'argument.'
674 )
675 raise RuntimeError(msg)
677 debug.callback(raise_if_mismatch)