jax

bfd9e1fb - Add support for parallelizing lapack invocations on the batch dimension.

Add support for parallelizing lapack invocations on the batch dimension. The rough idea is that we divide batched kernels into chunks of batch elements, where we have two constraints: * we want to divide work evenly amongst the thread pool, if there is enough. * we want a minimum work unit size so we don't put tiny units of work on each thread. For this we use an analytic flop estimate for each kernel. PiperOrigin-RevId: 895873630