The ML Engineer's Path #2intermediate

TensorFlow, JAX & the Google AI Stack

Google builds more ML infrastructure than anyone. Learn TensorFlow + Keras for production, JAX for research-speed training, TPUs for scale, and Vertex AI for deployment. Plus the Gemini API for multi-modal AI.

What You Build

An end-to-end app using Google AI: Keras model, Vertex AI deployment, TF.js frontend, Gemini integration

Sections

Lessons

Estimated

Concepts

Prerequisites

PyTorch from Zero to Production

tensorflowkerasjaxtpuvertex-aigeminigoogleml

Dev Jordan

walkthrough · casual

First 4 lessons are free (29% of the course). No login required.

The Google AI Landscape

Map the ecosystem: TensorFlow, Keras, JAX, TPUs, Vertex AI, Gemini.

The Google AI Ecosystem Map

20 min Build in Public

Tweet: Tweet: 'Mapping the Google AI stack: TensorFlow, Keras, JAX, TPUs, Vertex AI. Here's when to use each.'

TensorFlow vs JAX: When to Use Which

20 min

TensorFlow & Keras

Build models with the most deployed ML framework in the world.

Keras 3: The Universal API

35 min Ship

Ship: A Keras model trained on tabular data with evaluation metrics

TensorFlow Data Pipelines (tf.data)

30 min Ship

Ship: Efficient data pipeline processing 100K+ samples

Preview — content blurs halfway through

TensorFlow Hub & Pre-trained Models

25 min

TensorBoard: Visualize Everything

25 min

Unlock the rest of this course

$5 for full course or $25/mo for everything

See pricing

JAX: Research-Grade Speed

Google's next-gen framework: JIT compilation, auto-diff, TPU-native.

JAX Fundamentals: Numpy on Steroids

35 min

Automatic Differentiation with grad()

30 min

Flax & Optax: JAX Neural Networks

35 min

TPU Training: Free Colab to Cloud TPUs

35 min

Deploy with Google Cloud

TFLite, TF.js, Vertex AI, and the Gemini API.

TFLite & TF.js: Models Everywhere

30 min

Vertex AI: Google's MLOps Platform

35 min

Gemini API & Google AI Studio

30 min

Final Ship

Build and deploy an app using the full Google AI stack.

Ship a Google AI-Powered App

50 min

Next Course

Meta AI & Open Source ML

Meta has open-sourced more AI than any company: Llama, FAISS, SAM, Detectron2. Learn to run LLMs locally, fine-tune with LoRA, build vector search, and deploy RAG pipelines — all with zero API costs. Plus Microsoft ONNX, Hugging Face, and MLflow.

TensorFlow, JAX & the Google AI Stack

What You Build

An end-to-end app using Google AI: Keras model, Vertex AI deployment, TF.js frontend, Gemini integration