max_chunk_size_within

t3toolbox.backend.sampling_derivatives.max_chunk_size_within(mode_shapes, tucker_ranks, tt_ranks, order, n_probes, target_bytes, *, n_tangent=1, n_shards=1, dtype=None)

def max_chunk_size_within(
        mode_shapes:    typ.Sequence[int],  # (N_1..N_d)   ambient dims
        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| (global; divided by n_shards)
        target_bytes:   float,              # absolute peak-memory cap for the assembly (per device)
        *,
        n_tangent:      int = 1,
        n_shards:       int = 1,
        dtype:          typ.Any = None,
) -> int:                                   # the largest chunk_size whose assembly peak fits target_bytes

The largest chunk_size whose gradient-assembly peak stays within target_bytes (per device) – the “use my whole device” policy: pass an absolute byte budget (e.g. a fraction of device memory). Contrast the device-agnostic estimate_chunk_size(). Same measured per-row cost.

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)

• target_bytes (float)

• n_tangent (int)

• n_shards (int)

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

Return type:

int