estimate_chunk_size
===================

.. py:function:: t3toolbox.backend.sampling_derivatives.estimate_chunk_size(mode_shapes, tucker_ranks, tt_ranks, order, n_probes, *, n_tangent = 1, n_shards = 1, dtype = None)

   .. code-block:: python

      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
   :py:func:`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.