JIT-compiled first- and second-order optimization algorithms in JAX, blazingly fast and scalable.

Implementation of MAML algorithm for RL for continuous control tasks using TorchRL.

Exploring the 'grokking' phenomenon for modular addition using a one-layer transformer with bidirectional attention.

Training word embeddings using negative sampling.

Proof and simulation of convergence of (scaled) stochastic SIR epidemic model to solution of the deterministic SIR ODE system.