estimate_chunk_size#

t3toolbox.backend.sampling_derivatives.estimate_chunk_size(mode_shapes, tucker_ranks, tt_ranks, order, n_probes, *, n_tangent=1, n_shards=1, dtype=None)#
def estimate_chunk_size(
        mode_shapes:    typ.Sequence[int],  # (N_1..N_d)   ambient dims (as passed to TuckerTensorTrain.randn)
        tucker_ranks:   typ.Sequence[int],  # (nU_1..nU_d) Tucker ranks
        tt_ranks:       typ.Sequence[int],  # (r_0..r_d)   the d+1 TT bonds
        order:          int,                # highest derivative order K
        n_probes:       int,                # |W|, the number of probes (global; divided by n_shards)
        *,
        n_tangent:      int = 1,            # tangent stack K (a batch of tangents sharing the frame)
        n_shards:       int = 1,            # W split across this many devices -> chunk sizes the LOCAL shard
        dtype:          typ.Any = None,     # array dtype; default float32 (jax default). float64 under x64
) -> int:                                   # a memory-balanced chunk_size for tv_probe_derivatives_transpose

A memory-balanced chunk_size for the probe-derivative transpose assembly, from the problem shapes – call it once (eagerly, outside jit) and pass the int as chunk_size.

Picks the largest chunk whose assembly peak is comparable to the edge-variable jets already resident (the necessary floor), so the assembly is never the tallest pole – if the rest of the pipeline fits, so does the assembly (total peak ~2x the floor). Device-agnostic: no device-memory query, only the shapes. The assembly per-row cost is measured via XLA’s own scratch accounting (needs jax; a ~2 s one-time compile, cached by shape). For an absolute byte cap instead, see max_chunk_size_within(); the memory model is in the chunking design note.

Sized for the larger of the TT-core and Tucker gradients (Tucker wins when N >> n). n_shards divides W so the chunk sizes each device’s shard (use shard_map; see the chunking note). Only meaningful on the uniform+jax path – numpy/ragged never chunk.

Parameters:

  • mode_shapes (t3toolbox.backend.common.typ.Sequence[int])

  • tucker_ranks (t3toolbox.backend.common.typ.Sequence[int])

  • tt_ranks (t3toolbox.backend.common.typ.Sequence[int])

  • order (int)

  • n_probes (int)

  • n_tangent (int)

  • n_shards (int)

  • dtype (t3toolbox.backend.common.typ.Any)

Return type:

int