assemble_tt_variation_jets#

t3toolbox.backend.sampling_derivatives.assemble_tt_variation_jets(sigma_tildes, tau_tildes, deta_tildes, xi_jets, mu_jets, nu_jets, trs, n_probe, sum_over_probes, chunk_size=100)#
def assemble_tt_variation_jets(
        sigma_tildes:   typ.Sequence[NDArray],
        tau_tildes:     typ.Sequence[NDArray],
        deta_tildes:    typ.Sequence[NDArray],
        xi_jets:        typ.Sequence[NDArray],
        mu_jets:        typ.Sequence[NDArray],
        nu_jets:        typ.Sequence[NDArray],
        trs:            NDArray,
        n_probe:        int,
        sum_over_probes: bool,
        chunk_size:     typ.Optional[int] = 100,    # W-chunk size; None (or >= W) -> dense. See docs/chunking.md
) -> NDArray:                                        # dG_tildes supercore, [W+]C+(rLi,nUi,rRi)

Assemble the TT variation gradient (standard W-chunked form; dense reference assemble_tt_variation_jets_trs()).

The dense assembly’s peak is exactly LINEAR in the sample stack W (measured: ~5.3 MB / W-row at r=128), so it is chunked along W: the dense assembly runs per W-chunk and the partials are combined. Peak ~ chunk_size * (per-W-row) instead of W * (per-W-row).

The reducer is the seam that keeps this from locking into W-only (Nick, 2026-07-16). A chunked batch axis is combined by ADD if it is summed (sum_over_probes -> the gradient) or by CONCAT if it is kept (the sum_over_probes=False per-probe output, and – later – a chunked frame stack C, which is always kept). The per-chunk assembler is axis-agnostic (the dense assembly), so extending to C-chunking is a new slice front + the same reducer, not a rewrite.

Chunking runs on the uniform path with a single W axis; ragged / multi-W / chunk_size unset / W <= chunk_size fall back to the dense assembly. The chunk map is a real lax.map (sequential), so only one chunk’s intermediate is resident.

Parameters:

  • sigma_tildes (t3toolbox.backend.common.typ.Sequence[NDArray])

  • tau_tildes (t3toolbox.backend.common.typ.Sequence[NDArray])

  • deta_tildes (t3toolbox.backend.common.typ.Sequence[NDArray])

  • xi_jets (t3toolbox.backend.common.typ.Sequence[NDArray])

  • mu_jets (t3toolbox.backend.common.typ.Sequence[NDArray])

  • nu_jets (t3toolbox.backend.common.typ.Sequence[NDArray])

  • trs (NDArray)

  • n_probe (int)

  • sum_over_probes (bool)

  • chunk_size (t3toolbox.backend.common.typ.Optional[int])

Return type:

NDArray