Mojo for compiling ML models

[waveworks-ai]

I’m curious about thoughts on how to use Mojo to define and compile machine learning models. What I mean here is not using the Modular inference engine, nor using Mojo to interface libraries like TensorFlow or PyTorch, but to define the actual model data flow and computations in Mojo itself and have them compiled to efficient binary code.

For example:

fn model(x: ndarray) -> ndarray:
    let h1 = conv1d[filters=32, kernel_size=5](x)
    let h2 = lstm[units=16](h1)
    let y = dense[units=1](h2)
    return y

And then there would be efficient implementations of conv1d(), lstm() and dense() also written in Mojo. I’m particularly interested in how one would deal with model parameters (or maybe I should say trainable variables or weights, since “parameters” seems to have another meaning in Mojo - a bit unfortunate if you ask me). Ideally there could be a syntax to annotate trainable variables, for example:

fn dense[units: int](x: ndarray) -> ndarray:
    # Trainable variables
    @trainable
    let kernel = glorot_uniform[len(x), units]
    @trainable
    let bias = np.zeros[units]
    # Assuming relu activation
    return relu(np.transpose(kernel) @ x + bias)

I think such a syntax would be useful for two purposes:

• The code can be compiled to a computation graph for training in TensorFlow or Pytorch. There could be an official Mojo implementation of Keras layers and ops, and in addition it would be very easy to write custom layers and runtime data preprocessing.
• Then the same code can be compiled to efficient binary code for inference, and the compiler can take care of injecting the weights. For edge devices like microcontrollers, the compiler can inject the weights statically at compile time and potentially optimize further.

For context, I designed (and wrote a compiler for) a small language called FL with exactly this feature (https://github.com/waveworks-ai/FL). This workflow has proven very efficient at my company, and now I’m hoping that we will be able to do the same with Mojo in one way or the other.

Thanks in advance for any thoughts or insights!