Leveraging Python and JAX in R workflows

Andrés Cruz (UT Austin)

NU Statistical Computing Workshop

Apr 22, 2026

Python won

  • (in ML/AI)
  • Are we missing out?

R vs Python

  • R: statistics/data analysis focused
    • Data wrangling
    • Plotting
    • Solid stats / econometrics / pol. methodology
  • Python: general-purpose
    • Inter-operability (e.g., API calls)
    • Cutting-edge AI/ML
  • Question: are LLMs better at R or Python?

R vs Python: LLM performance

  • 🐍 “LLMs Love Python” (Twist et al. 2026)

    • In language-agnostic queries, “Python accounts for 90-97% of generated solutions” (7)

  • ®️ AutoCodeBench performance (Chou et al. 2025, 6)

A compromise

  • Integrate tidbits of Python into our R workflows

Source: Wikipedia

What we’ll cover today

  1. Using Python from R

    • Intro to reticulate

    • Example: sentence-level semantic similarity

  2. Leveraging JAX, a high-performance Python library

    • Automatic differentiation

    • Example: sensitivity analysis for error propagation

1. Using Python from R

reticulate

  • A package to interact with Python from R

    • Manages Python packages and environments

    • Translates between R and Python objects

  • Everything happens in R. You write R code!

    • e.g., (1) R data prep; (2) Python snippet; (3) R analysis

What’s (usually) better in Python?

  • Commercial APIs

    • LLMs (Anthropic, OpenAI, Google)

    • OCR (Mistral)

    • Data download (Youtube, Google Maps)

  • Pre-trained models for image/video/text

  • Cutting-edge scientific computing

Example: embeddings

  • Working with image/video/text usually involves:
    1. Wrangling
    2. Numerical representation
    3. Analysis (stats, plots, etc.)
  • For step 2, we often want to use pre-trained encoder models to generate “embeddings”
    • Wide availability in Python, both closed- and open-source

Source: Wikimedia Commons

Constitutional consultation (Cruz et al. 2023)

  • ~250k citizen submissions to the CL constitutional process

  • Compare submissions with topics and text from the world’s constitutions

    1. Wrangle data

    2. Embed text using multilanguage encoder

    3. Compute “semantic” similarities

Script: 1_embeddings.R

Translating Python to reticulate

  • The package automatically handles loops, indexing, and other tricky parts

  • R lists do a lot: e.g., package.module.function() becomes package$model$function()

  • In my experience, LLMs are good for short Python-to-reticulate translations

    • except for infrastructure (setting, exit)

2. Leveraging JAX from R

JAX

  • A high-performance Python package/ecosystem

  • By Google: powers some of their ML/AI

  • Good CPU and GPU performance

  • Key features:

    • Just-in-time-compilation

    • Automatic vectorization

    • Automatic differentiation

How to calculate derivatives?

  • Analytic differentiation

    • e.g., we know \(f(x)=x^2; f'(x)=2x\).

  • Numeric differentiation

    • Approximate derivative by guessing, guessing, guessing…

  • Automatic differentiation

    • Under the hood: boil functions down to their elementary operations; use the chain rule.

    • In JAX: differentiate (pretty much) arbitrary functions, e.g., most things in numpy (ex: marginaleffects).

Script: 2_jax.R (part one)

Example: sensitivity analysis

  • Examine the degree to which our assumptions affect our results

  • Often apply extreme pressure

  • Tools: differentiation, optimization

A sensitivity approach to measurement error

  • How much measurement error would invalidate a downstream estimate?

A constrained optimization problem

  • Minimize the quantity of interest, while respecting the pointwise error ranges

  • We have great algorithms for box-constrained optimization, e.g., L-BFGS-B (Byrd et al. 1995)

  • They require derivatives

    • Analytical differentiation is not always available (e.g., MLE)

    • Numerical differentiation is slow and prone to error

  • Solution: automatic differentiation!

Script: 2_jax.R (part two)

Summary

  • We can integrate Python tidbits into R via reticulate

  • Good candidates for Python snippets: commercial APIs, pre-trained models, scientific computing

  • LLM translation often works, but infrastructure is an always-moving target

References

Byrd, Richard H., Peihuang Lu, Jorge Nocedal, and Ciyou Zhu. 1995. “A Limited Memory Algorithm for Bound Constrained Optimization.” SIAM Journal on Scientific Computing 16 (5): 1190–1208. https://doi.org/10.1137/0916069.
Chou, Jason, Ao Liu, Yuchi Deng, Zhiying Zeng, Tao Zhang, Haotian Zhu, Jianwei Cai, et al. 2025. “AutoCodeBench: Large Language Models Are Automatic Code Benchmark Generators.” https://arxiv.org/abs/2508.09101.
Cinelli, Carlos, and Chad Hazlett. 2020. “Making Sense of Sensitivity: Extending Omitted Variable Bias.” Journal of the Royal Statistical Society Series B: Statistical Methodology 82 (1): 39–67. https://doi.org/10.1111/rssb.12348.
Coppedge, Michael, John Gerring, Carl Henrik Knutsen, Staffan I. Lindberg, Jan Teorell, David Altman, Fabio Angiolillo, et al. 2025. “V-Dem Dataset V15.” Varieties of Democracy (V-Dem) Project. https://doi.org/10.23696/vdemds25.
Cruz, Andrés, Zachary Elkins, Roy Gardner, Matthew Martin, and Ashley Moran. 2023. “Measuring Constitutional Preferences: A New Method for Analyzing Public Consultation Data.” PLOS ONE 18 (12): e0295396. https://doi.org/10.1371/journal.pone.0295396.
Twist, Lukas, Jie M. Zhang, Mark Harman, Don Syme, Joost Noppen, Helen Yannakoudakis, and Detlef Nauck. 2026. “A Study of LLMs’ Preferences for Libraries and Programming Languages.” https://arxiv.org/abs/2503.17181.
VanderWeele, Tyler J., and Peng Ding. 2017. “Sensitivity Analysis in Observational Research: Introducing the e-Value.” Annals of Internal Medicine 167 (4): 268–74. https://doi.org/10.7326/m16-2607.