To be clear, we tested a large set of batch sizes from 1 up to ~1e8. The smaller the batch size, the less competitive JAX was. The crossover point for JAX to be faster was amortizing across a batch size in the tens of thousands I believe. The ideal batch size for some of our applications is one to five. Please see symforce/benchmarks/notebooks for some of our timing.

SymForce is not trying to compute on large data-parallelism, but it could generate an inner kernel for that. In that way it complements TensorFlow nicely.