jax_privacy.experimental.compilation_utils

Experimental utilities for handling variable batch sizes.

Functions

optimal_physical_batch_sizes(batch_sizes, ...)

Find a set of of compiled batch sizes that minimizes wasted compute.