guess_batch_size

anri.fwd.guess_batch_size(entries, hkls, F2, geom, row, det_shape, window=(3, 7, 7), mesh=None, memory_fraction=0.25)

Largest batch for render_row() whose render step fits in a fraction of the free memory.

Compiles the render step for two small batches (nothing is rendered) and reads XLA’s memory analysis to get the bytes per peak, then fits as many peaks as memory_fraction of the free memory allows: on GPUs, of the least free device; on CPU, of the host memory shared by the CPU devices, counting the host copy of each batch’s output too. Precision matters: float64 needs about twice the memory of float32. Selecting the peaks (select_chunk in render_row()) needs memory too, a few hundred MB by default, which is not counted.

Parameters:

• entries (dict) – As for render_row()

• hkls (ndarray) – As for render_row()

• F2 (ndarray) – As for render_row()

• geom (dict) – As for render_row()

• row (dict) – As for render_row()

• det_shape (tuple[int, int]) – As for render_row()

• window (tuple[int, int, int], default: (3, 7, 7)) – As for render_row()

• mesh (Mesh | None, default: None) – As for render_row()

• memory_fraction (float, default: 0.25) – Fraction of the free memory to use. The default leaves room for other users of a shared machine.

Returns:

batch (int) – A power of two times the number of devices, at least 64 peaks per device