Disk-Backed vs. On-Batch Methods

This page explains and compares the two complementary execution styles provided by ImpactModel:

• Disk-backed (default) methods iterate over the input in chunks, materialize results incrementally, and persist structured artifacts (Zarr-backed xarray.DataTree plus metadata) to a temporary or user-specified output directory.

• On-batch (*_on_batch suffix) methods execute a single, fully in-memory pass and can optionally return a plain dict instead of a xarray.DataTree. The naming mirrors the Keras convention to signal an immediate, single-batch, memory-resident operation.

Why Disk-Backed by Default

The non-*_on_batch methods default to a disk-backed (chunked) execution model for several reasons:

• Posterior predictive and prior predictive tensors can scale as (#samples x #dims x #posterior_samples x ...). Even moderate increases in any axis (time, spatial units, parameter samples) can exceed host or accelerator RAM.

• Using batch_size with chunked iteration limits peak memory and prevents out-of-memory errors.

• Persisted Zarr arrays with metadata (coords, dims, attrs) create an artifact you can reopen without rerunning inference.

• The xarray.DataTree + Zarr format integrates with scientific Python tools such as Dask and ArviZ.

• Summaries (means, HDIs, residual PPC stats) can be computed lazily over chunked storage without first materializing dense arrays.

• One API works for both small experiments and large-scale use cases.

Comparison

Disk-backed variants target larger datasets, enable chunked processing, multi-device parallelism, and stable artifact generation. These methods build internal data loaders, iterate in chunks, and decouple sampling from file I/O, enabling concurrent execution. Outputs consolidate into a single xarray.DataTree backed by Zarr files for post-hoc analysis. On-batch variants, in contrast, favor minimal overhead, immediate return, and greater flexibility when posterior sample shapes are not shard-friendly.

Feature Summary

Feature	Disk-backed (default)	On-batch (*_on_batch)
Typical dataset size	Medium → large	Small → moderate
Supported use cases	Standard models	Broader model support
Peak memory usage	Chunk-bounded	Full batch resident
Writes to disk	Yes	No
Return type	xarray.DataTree	xarray.DataTree or dict (via return_datatree=False)
Custom batch sizing	Yes (batch_size)	No (single pass)
Device parallelism (sharding)	Yes	No
Automatic fallback	Yes (may auto‑delegate to on‑batch)	No (final mode)
Latency (small data)	Higher (I/O + orchestration)	Minimal

Capability Matrix

Capability	Disk-backed (default)	On-batch (*_on_batch)
Full dataset training	fit()	fit_on_batch()
Single training step	N/A	train_on_batch()
Prior predictive sampling	sample_prior_predictive()	sample_prior_predictive_on_batch()
Posterior sampling	sample()	N/A
Posterior predictive sampling	predict() or sample_posterior_predictive()	predict_on_batch() or sample_posterior_predictive_on_batch()
Log-likelihood computation	log_likelihood()	N/A
Effect estimation	estimate_effect()	(consumes outputs above)

Quick Recommendations

• Moderate or large data, or need persisted outputs: use disk-backed (e.g., fit(), predict()).

• Small data, rapid iteration, CI, or read-only / ephemeral filesystem: use on-batch (*_on_batch).

• If predict() issues a fallback warning, call predict_on_batch() directly. This occurs when the model or posterior sample shapes are incompatible with shard-based chunked execution.

• Custom training loop: iterate with train_on_batch().

• Need multi-device (sharding) execution: disk-backed.

• Need raw NumPy/dict outputs (no xarray.DataTree): on-batch with return_datatree=False.

Note

For MCMC inference, only fit_on_batch() or sample() is supported for training and posterior sampling, as MCMC is incompatible with epoch-based or chunked batch processing. See MCMC Support for more details.

Example: predict() with Fallback Warning

A common scenario for the fallback warning occurs when the model contains local latent variables, which can make posterior sample shapes incompatible with shard-based parallel execution. The example below illustrates this case.

import logging

import jax.numpy as jnp
import numpyro.distributions as dist
from jax import random
from jax.typing import ArrayLike
from numpyro import optim, plate, sample
from numpyro.infer import SVI, Trace_ELBO
from numpyro.infer.autoguide import AutoNormal

from aimz import ImpactModel

logging.basicConfig(level=logging.INFO, force=True)


def model(X: ArrayLike, y: ArrayLike | None = None) -> None:
    # Model includes a local latent variable
    sigma = sample("sigma", dist.Exponential().expand((X.shape[0],)))
    with plate("data", size=X.shape[0]):
        sample("y", dist.Normal(0.0, sigma), obs=y)


rng_key = random.key(42)
rng_key, rng_key_X, rng_key_y = random.split(rng_key, 3)
X = random.normal(rng_key_X, (100, 2))
y = random.normal(rng_key_y, (100,))


im = ImpactModel(
    model,
    rng_key=rng_key,
    inference=SVI(
        model,
        guide=AutoNormal(model),
        optim=optim.Adam(step_size=1e-3),
        loss=Trace_ELBO(),
    ),
# This internally calls the `.run()` method of `SVI`
).fit_on_batch(X, y)

# Calling `.predict()` triggers a fallback warning
im.predict(X)

/tmp/ipykernel_829/1252916572.py:2: UserWarning: One or more posterior sample shapes are not compatible with `.predict()` under sharded parallelism; falling back to `.predict_on_batch()`.
  im.predict(X)

<xarray.DataTree 'root'>
Group: /
├── Group: /posterior
│       Dimensions:      (chain: 1, draw: 1000, sigma_dim_0: 100)
│       Coordinates:
│         * chain        (chain) int64 8B 0
│         * draw         (draw) int64 8kB 0 1 2 3 4 5 6 ... 993 994 995 996 997 998 999
│         * sigma_dim_0  (sigma_dim_0) int64 800B 0 1 2 3 4 5 6 ... 93 94 95 96 97 98 99
│       Data variables:
│           sigma        (chain, draw, sigma_dim_0) float32 400kB 1.615 0.8619 ... 2.558
│       Attributes:
│           created_at:    2026-05-24T03:02:13.883093+00:00
│           aimz_version:  0.12.0
└── Group: /posterior_predictive
        Dimensions:  (chain: 1, draw: 1000, y_dim_0: 100)
        Coordinates:
          * chain    (chain) int64 8B 0
          * draw     (draw) int64 8kB 0 1 2 3 4 5 6 7 ... 993 994 995 996 997 998 999
          * y_dim_0  (y_dim_0) int64 800B 0 1 2 3 4 5 6 7 8 ... 92 93 94 95 96 97 98 99
        Data variables:
            y        (chain, draw, y_dim_0) float32 400kB -3.284 -0.4314 ... 3.398
        Attributes:
            created_at:    2026-05-24T03:02:13.885558+00:00
            aimz_version:  0.12.0

xarray.DataTree

'root'

• /posterior(6)

  • Dimensions:
    • chain: 1
    • draw: 1000
    • sigma_dim_0: 100
  • Coordinates: (3)
    • chain
      (chain)
      int64
      0

      array([0])

    • draw
      (draw)
      int64
      0 1 2 3 4 5 ... 995 996 997 998 999

      array([  0,   1,   2, ..., 997, 998, 999], shape=(1000,))

    • sigma_dim_0
      (sigma_dim_0)
      int64
      0 1 2 3 4 5 6 ... 94 95 96 97 98 99

      array([ 0,  1,  2,  3,  4,  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])

  • Data variables: (1)
    • sigma
      (chain, draw, sigma_dim_0)
      float32
      1.615 0.8619 0.7909 ... 1.017 2.558

      array([[[1.6151611 , 0.8618528 , 0.790872  , ..., 1.6389374 ,0.33729208, 0.42537448],[1.3211713 , 0.4304575 , 2.0427656 , ..., 0.7860633 ,0.15371482, 0.27378327],[1.2203506 , 0.5631171 , 2.225587  , ..., 0.23594743,0.08826305, 0.11016101],...,[0.5101084 , 0.20624015, 1.5212504 , ..., 0.40487286,0.16055402, 0.8135405 ],[0.1295606 , 1.2459134 , 3.3112087 , ..., 0.13780202,0.18131228, 0.21599922],[0.36301368, 0.7408936 , 1.5247383 , ..., 0.6440523 ,1.0174416 , 2.5577676 ]]], shape=(1, 1000, 100), dtype=float32)

  • Attributes: (2)

    created_at :

    2026-05-24T03:02:13.883093+00:00

    aimz_version :

    0.12.0

  /posterior_predictive(6)

  • Dimensions:
    • chain: 1
    • draw: 1000
    • y_dim_0: 100
  • Coordinates: (3)
    • chain
      (chain)
      int64
      0

      array([0])

    • draw
      (draw)
      int64
      0 1 2 3 4 5 ... 995 996 997 998 999

      array([  0,   1,   2, ..., 997, 998, 999], shape=(1000,))

    • y_dim_0
      (y_dim_0)
      int64
      0 1 2 3 4 5 6 ... 94 95 96 97 98 99

      array([ 0,  1,  2,  3,  4,  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])

  • Data variables: (1)
    • y
      (chain, draw, y_dim_0)
      float32
      -3.284 -0.4314 ... -1.696 3.398

      array([[[-3.2839305 , -0.43143615, -1.52514   , ..., -2.6352901 ,-0.08678479, -0.46478134],[ 0.15668277, -0.24548501,  0.5022501 , ..., -2.0746102 ,-0.19719376,  0.3086634 ],[-1.0833476 ,  0.31739277, -4.512692  , ..., -0.19868693,0.03921507, -0.12311906],...,[ 0.04891762,  0.18733212, -0.6157141 , ...,  0.693681  ,0.0554003 ,  0.06972526],[ 0.23430954, -0.8604826 , -1.3664719 , ...,  0.04435274,-0.08675548,  0.1391157 ],[-0.73923504,  0.34483376, -1.0438837 , ...,  0.49250746,-1.6960533 ,  3.3979347 ]]], shape=(1, 1000, 100), dtype=float32)

  • Attributes: (2)

    created_at :

    2026-05-24T03:02:13.885558+00:00

    aimz_version :

    0.12.0

Performance Tips

• Tune batch_size appropriately; it also determines the chunk size for Zarr-backed arrays.

• Monitor disk usage, as chunk sizes scale with batch_size and num_samples.

• Reduce num_samples first for faster iteration.

• Use on-batch methods in tests to minimize I/O overhead.