Robert Dyro

I work on JAX at Google. I completed my PhD in Autonomous Systems Lab (ASL) at Stanford University, on differentiable optimization applied to behavior inference in autonomous driving.

Email  /  Resume  /  Bio  /  Google Scholar  /  Github

Here are some articles I (co)authored showcasing some of the projects I worked on.

Open-source Deepseek R1 inference using JAX, minimal yet performant.

An attempt to improve the inference speed of HuggingFace's implementation of Mistral-7B in JAX for RTX 3090 via Pallas kernel tuning.

This set is puzzles is meant to teach you how to use Pallas from first principles in an interactive fashion. Adapted from Sasha Rush's Triton-Puzzles.

Post-training reduction of quantized NN memory footprint via SVD.

A package for efficiently wrapping existing PyTorch code as JAX functions in a JIT-compatible way. AutoDiff gradients can be automatically defined, fully integrating any PyTorch code with JAX differentiation.

Automatic short answer grading with pretrained language models.

Controlling an airplane using optimal control with tuned objective parameterization using ray.tune.

An introduction to GraphGym and an extension of GraphGym to perform Neural Architecture Search (NAS).

This tutorial is at attempt to produce an educational reference for JAX techniques we found particularly interesting in our work and research. As such, it is a collection of a few disjoint topics.

Optimizing Models for Fairness and Explainability via Shapley value and local linear model penalization and differentiation.

Optimization sensitivity analysis with PyTorch and JAX.

A QP solver fitting into a single CUDA block.

Reverse-mode automatic differentiation with full sparse Jacobians and Hessians.

Model Predictive Control - general and consensus optimization in Julia & Python.

I'm interested in optimization, robotics, optimization, and machine learning. My research is about creating new models and algorithms for robotic systems that exploit known structure of the problem.

Extracting learned behavior distribution for realistic counterfactual generation via efficient and scalable Hessian sketching.

Optimal planning & control in new environments can be done effectively by learning deep signed distance function (SDF) maps from LiDAR data.

Differentiating through optimization can be done twice to obtain bilevel program Hessians.

Sampling arbitrary dynamical uncertainty and then optimization over the possibilities <em>jointly</em> gives safe real-time optimal control.