Beam for Large-Scale, Accelerated ML Inference at Google

Bulk Inference in Machine Learning (ML) refers to the challenge of how to organize and compute model predictions for a large pool of available input data with no latency requirements. JAX is an open-source computation library commonly used by both engineers and researchers for flexible, high-performant ML development. This talk will illustrate how teams at Google are using Beam to ergonomically design, orchestrate, and scale JAX Bulk Inference workloads across various accelerator platforms.