JAX on SageMaker
Amazon SageMaker is AWS’s ML platform that covers many of the steps in the ML platform. One of the useful
abstractions in SageMaker is the Estimator. The Estimator determines the underlying framework and gives access
to many of the SageMaker APIs such as Training Jobs and Hosted Endpoints.
I co-authored a blog post that describes how to integrate JAX with Amazon SageMaker’s estimator. You can read it here
Dr. Archis Joglekar
ML Researcher | Research Engineer | Theoretical Physicist
I like doing math with computers. I got a PhD in fusion plasma physics. It happens to be a perfect blend of applied mathematics, physics, and computing. I used to use supercomputers to do the math, now I use the cloud. I also like to do written math. I am currently working on something new at the intersection of deep learning and fusion. I am an Affiliate Researcher with the Laboratory for Laser Energetics and I am also an Adjunct Professor at the University of Michigan.