Unify

I joined Unify, a Y Combinator company (W23), in March 2023 as a machine learning research engineer, after contributing to Ivy as an outside open-source contributor.

The problem

In 2023 a model or library written for one machine learning framework could not be used from another without being rewritten by hand.

Research code was mostly PyTorch, production stacks were often TensorFlow, and JAX was where teams went for TPUs and compiler speed. Unify was building Ivy to close that gap.

The framework

Ivy defined each operation once and implemented it for NumPy, PyTorch, TensorFlow, and JAX, then reimplemented each framework's own API on top, so code written for any of them could be expressed in Ivy and then in any other.

I extended its PyTorch, TensorFlow, and JAX frontends, the backends beneath them, and its stateful API, where I made converted stateful modules behave as native ones, including PyTorch modules turned into trainable Keras models with tied weights and LSTMs.

Tracing and transpiling

I led the team building Ivy's graph tracer and transpiler, which ran a model on example inputs, recorded every operation as a computation graph, and rebuilt that graph in another framework.

The hard part was what a trace does not see. A plain Python value such as x.shape[0] was frozen into a constant, so the graph broke on the next input shape. I made traces carry Python values as well as tensors, through proxy objects that behave like integers, shapes, and lists while recording what is done to them, and adapted a fork of TensorFlow's AutoGraph so built-ins such as len appeared in the trace too.

The tracer also gained jax.jit support for transpiled Hugging Face models, and transpile time on very large graphs, such as FasterRCNN converted to JAX, fell from about 65 minutes to about 25.

What tracing could not do

A trace resolves branches and loops for the inputs it saw, so a converted model could not take a different path for a different input.

Late in 2023 I restructured an early prototype of a different approach, rewriting the Python source instead of tracing it, into the staged architecture that the source-to-source transpiler at Ivy was later rebuilt on.

Beyond the code

Enterprise teams evaluated the transpiler during this period, and I prepared the builds and demonstration notebooks they worked from and supported them as they tried it on their own models.

I also interviewed candidates for Unify's hiring team, and co-mentored a Google Summer of Code project to make TensorLy, a library for tensor decompositions, usable with every framework Ivy supported.