tv_probe_derivatives_transpose#

t3toolbox.backend.sampling_derivatives.tv_probe_derivatives_transpose(ztildes, ww, pp, frame, order, sum_over_probes=False, chunk_size=100)#
def tv_probe_derivatives_transpose(
        ztildes:    typ.Sequence[NDArray],  # residual jets, len=d, elm_shape=(order+1,)+W+K+C+(Ni,)
        ww:         typ.Sequence[NDArray],  # probe vectors X,        len=d, elm_shape=W+(Ni,)
        pp:         typ.Sequence[NDArray],  # perturbation vectors P, len=d, elm_shape=W+(Ni,)
        frame:       typ.Tuple[
            typ.Sequence[NDArray], typ.Sequence[NDArray],
            typ.Sequence[NDArray], typ.Sequence[NDArray],
        ],                                  # = T3Frame.data = (U, O, P, Q)
        order:      int,                    # highest derivative order K
        sum_over_probes: bool = False,      # True: sum the sample stack W (the J^T r back-projection)
        chunk_size: typ.Optional[int] = 100,   # W-chunk size for the gradient assembly; None -> dense. docs/chunking.md
) -> typ.Tuple[
    typ.Tuple[NDArray, ...],  # dU_tildes (Tucker variation gradient)
    typ.Tuple[NDArray, ...],  # dG_tildes (TT variation gradient)
]:                                          # = T3Variations.data

Transpose of tv_probe_derivatives(): back-project residual jets ztildes into a variation gradient (dU_tildes, dG_tildes). The jet-ified adjoint-state method (t4s.pdf Thm 7): every forward contraction is swapped for its trs version, then stationarity of the Lagrangian gives the adjoint sweeps (sigma/tau/dxi_tilde jets) and the order-less gradient assembly.

Full W + K + C stacking, base-inner as in plain probing: the residual jets ztildes carry the sample stack W, the tangent stack K (the forward output’s K), and the frame stack C. The tangent batch K always rides through to the variation gradient. With sum_over_probes=False the sample stack W rides through into the variation stack too; with True it is summed (the J^T r back-projection used for fitting), K/C always kept. Verified against the dense adjoint identity <r, J v> = <J^T r, v> and jax.linear_transpose.

chunk_size bounds the peak memory of the (uniform+jax) gradient assembly by processing the sample stack W in slices of that size; None (or >= W) runs the dense assembly. The default is a safe fixed value, tuned only for moderate problems (a fixed chunk_size bounds the chunk count, not the bytes); for a memory-balanced value from the problem shapes call estimate_chunk_size() once (eagerly). Chunking engages only on the uniform+jax path (numpy/ragged free eagerly). See Memory chunking of the derivative transpose (chunk_size).

Parameters:

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

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

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

  • frame (t3toolbox.backend.common.typ.Tuple[t3toolbox.backend.common.typ.Sequence[NDArray], t3toolbox.backend.common.typ.Sequence[NDArray], t3toolbox.backend.common.typ.Sequence[NDArray], t3toolbox.backend.common.typ.Sequence[NDArray]])

  • order (int)

  • sum_over_probes (bool)

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

Return type:

t3toolbox.backend.common.typ.Tuple[t3toolbox.backend.common.typ.Tuple[NDArray, Ellipsis], t3toolbox.backend.common.typ.Tuple[NDArray, Ellipsis]]