Pageshift is seeking a specialized engineer to profile and optimize long-context JAX workloads on TPUs, focusing on XLA behavior and Pallas kernel development. You will implement and tune kernels, profile memory and sharding, and benchmark changes to validate improvements with data, contributing to the efficiency of large-scale AI storytelling workloads.You will collaborate with researchers and engineers to iterate on performance optimizations, and bring an experimental mindset, strong #J-18808-Ljbffr
TPU/JAX ML Performance Engineer Arbeitgeber: Pageshift
Pageshift ist ein hervorragender Arbeitgeber, der innovative Talente sucht, um an der Spitze der KI-Entwicklung zu arbeiten. Mit einem dynamischen Arbeitsumfeld, das Kreativität und Zusammenarbeit fördert, bieten wir unseren Mitarbeitern die Möglichkeit, sich kontinuierlich weiterzuentwickeln und an spannenden Projekten zu arbeiten, die die Zukunft des Geschichtenerzählens prägen. Unsere Mitarbeiter profitieren von flexiblen Arbeitszeiten, einer unterstützenden Kultur und der Chance, direkt mit führenden Forschern und Ingenieuren zusammenzuarbeiten.
