My best case scenario for the next two weeks is that we get basic SGD up and running for simple dense networks (using XLA). The first major hurdle for this is not just differentiating the whole compute graph but identifying which parameter nodes the user wants gradients for. After that we add bindings to more XLA operations like matrix multiplication and activation functions + specify their differentiation rules, that is simple. Finally we should talk a bit about optimizers, but that is for another issue.

xla-rs defines "constants" and "parameters" in a pretty intuitive way: constants are baked into the executable whereas parameters are specified at every execution. Our current setup with the compute graphs reflects this design, which I'm happy with. The thing is, not everything which is specified at every execution should be differentiated. I saw Braden's email to ro, and this is highlighted by the question of how we distinguish "input", which is not differentiated but does change at every XLA call.

I would personally appreciate the following sort of API design (but haven't thought it through entirely):

let mut ctx = Context::new();
let learning_rate = ctx.scalar(-0.001);
let x = ctx.parameter("x", SmallVec::new()); // think of this as neural network weights
let y = ctx.parameter("y", SmallVec::new()); // think of this as neural network inputs
let loss = ctx.add(ctx.mul(x, x), y); // meaningless computation but that's not the point
let grad = ctx.gradient(x); // not yet implemented, this just gets the derivative with respect to x
let updates = ctx.mul(learning_rate, grad); // loss does not depend on this
// figure out how to apply updates when we talk about optimizers
let xla_executable = ctx.compile();

My idea is that ctx.gradient creates a node in the compute graph which is the placeholder for the gradient of a parameter. I think there's an issue with how the backend is currently set up relative to this: compile basically expects the loss node of the compute graph, and certainly this makes sense for autodiff, but overall we want the computation to potentially encapsulate computations on which the loss does not depend (note that updates depend on loss, not the other way around, so passing ctx.compile(loss) as we have it now would miss the multiplication of learning rate by the gradient). In summary, compile should take a loss node but not view it as the endpoint of the computation, instead it should use autodiff to connect placeholder Gradient nodes to actual computations which can then be converted to XLA.

This opens up the question of how the user specifies what tensors the XLA executable should return (usually you are interested in network predictions, loss, parameter updates, and maybe some auxiliary data), I will sleep on this. Interested in you guys' thoughts.

We have a lot of useless build warnings. At some point someone should go through and tidy this up.

This issue seems to stem from the fact that we are using data from a previous layer to update the weights and biases of the current layer. When changing the gradient we should likely be using an activation function that we are not, but I cannot find out which one. I have tried using the current layer and the previous layer, so perhaps we have to use the next layers activation function to move forward, which would make the most sense

we need a way to have a computation with multiple outputs. add an operation type which is like Parameter​ but also holds a NodeIdentifier​ for the computation it outputs.

This should be a pretty simple one, xla-rs has an enum

pub enum ElementType {
    Pred,
    S8,
    S16,
    S32,
    S64,
    U8,
    U16,
    U32,
    U64,
    F16,
    F32,
    Bf16,
    F64,
    C64,
    C128,
}

Pred is bool and C stands for complex. This should mesh easily if the numpy importing is already working. This can also be solved in the same branch as #28 and #29

Currently, all of our tests are littered with .expect, which is really ugly. I would vote to panic from within all user-facing functions. I'm not really sure what the API would gain from doing anything else.

We should have utility functions for constructing dense/convolutional layers (eventually more complex layers like LSTM or multihead attention), which take a context, input node identifiers, and initialization instructions, and return node identifiers for the layer output and the parameters.

A basic example of this is visible in the mnist_xla example.

Intializers should use XLA RNGs.

comb through dimensionality propagation and ensure its correct

Need to alter the fit method for creating mini batches based on batch size, then taking each batch in parallel and generating the gradients for each input/output pair in parallel again. Once every thread has a full vector of gradients, generate the average gradient and finally update weights and biases according to this new vector of average gradients among batches.

asynchronously generating data from every layer of forward propagation and asynchronously gathering the gradient is currently implemented, I just need to look into the best way to execute the batches in parallel given the layers are Boxed dyn traits, which tokio and rayon don't like very much

Something somewhere is making the bf16 conversion fail. I'll be checking it out

The MNIST example is too nice insofar as both the train and test set have a number of samples divisible by the batch size of 100. In general we should not assume this is the case. Our abstract model API should support some form of dynamic batch size.

I am not sure what the best approach for this is yet. However, if we assume that the model is executing on a fixed batch size and that only once per epoch will it receive a differing batch size, then when we average the loss we could multiply by a mask along the batch axis (1s where there are samples, 0s where there are not) and divide by the sum of the mask. If we do it like this, maybe it should be opt-in, since it does introduce a few extra floating point operations.

We still need to bind and write differentiation rules for the following operations to train dense networks:
matmul
transpose
log
div
pow
reduce_sum
reduce_mean
In order to differentiate the last two we need a repeat or tile operation, which is apparently not yet bound in xla-rs. We also need a scatter operation in order to differentiate reduce_max and slice_in_dim, so I might look into adding bindings to xla-rs.

Issue is either coming from back prop or compilation.

I only implemented scalar, vector, and matrix for Context. Arbitrary dimension constants are supported, they just arent exposed by any API right now.

traverse the graph building a hashmap of nodes to node identifiers. upon insertion, if there's already an entry, replace references to the current node with the entry. make sure to update entry for the modified node, as the hash will change. do not include callsite when calculating the hash. start at the leaf nodes, work your way up.

add an operation type to ensure that a computation is already present in the compute graph. useful for referencing intermediate derivatives generated by autodiff with peace of mind that no additional work is created. remove these nodes if the constraint is satisified during common term extraction, and error if they are still present when compiling to XLA

write tests for every reduction. build an input graph and an expected output graph, run a reduction, and compare the output. test single steps, and total reductions. test that adding a vector to a scalar errors. write tests that dont pass yet but will pass, for example test that adding constants 2 + 1 folds to constant 3.

reverse Context::to_string using nom and strum. This will enable save/load and easier unit testing.

short term we should probably do this via serde and save out the whole context for simplicity.

https://github.com/LaurentMazare/xla-rs

This might be a trivial one, but it's important from a usability perspective.

PjRt executables expect a slice of buffers or literals as input to run on. However, if you have a model with dozens of parameter tensors organizing them all into a slice manually becomes tedious (already visible in the mnist_xla example), so I think our abstract model API should basically be callable on any user-defined structure implementing Into<Vec<Literal>> or Into<Vec<PjRtBuffer>>. The gradient engine should also return gradients with the same desired structure if the parameter struct implements From<Vec<T>>.

This is basically the equivalent of JAX tree flattening/unflattening

I don't think we have to do anything very thoughtful for this one, just write the type abstract signatures and call into and from appropriately.

Models and optimizers are generally thought of as separate objects, although currently they are executed in the same context.

This might be appropriate as a second step after forward/backward pass separation.

The critical reasons why we want the optimizer separate are 1) if updated weights are returned by the same context returning training loss which we want to print then weights are being bussed too and from the GPU at every step and 2) XLA supports Send and Recv operations which would allow us to compute gradient updates while simultaneously bussing the next model inputs and labels to the GPU.

We should also support SGD, RMSProp, and Adam optimizers. It would also make sense to make batch normalization a part of this issue (as it is its own sort of custom optimizer).

Write in an intuitive way that doesn't require a bunch of unnecessary gradient calculation

When trying to implement clone for a network, for some reason adding + Clone to the Layer trait ends up creating a cycle in minibatch? No idea how this is being caused

traverse the graph applying rules like 2 + 1 = 3, 1 * x = x, 0 * x = 0, x + 0 = x, etc.

highest priority is 0 * x = 0 rule because autodiff leaves behind many Const 0

node traversal should look a lot like the autodiff implementation

expose a way to save out a computation graph, accompanied by all constants parameters and model weights as relative-path hlo format files, xla-rs has utils for this and its the expected format for model/optimizer saving in XLA

We need to carefully design the API such that the user has access to both an executable that returns the desired diff calls (for training) and an executable that returns everything except that (for testing).

This is part of a larger array of issues that will emerge from the need to embed contexts in other contexts (for example, separating the optimizer step, or designing recurrent architectures). In this case it might make sense to allow the user to design a forward pass context which doesn't take labels or output gradients, then allow them to clone that context and recover all desired node identifiers in order to create another context that takes both labels and inputs and outputs both predictions and loss and gradients. Then both executables can be used separately.

Question mark because it might be most reasonable to depend on another crate for this.

At the most abstract level, I imagine us having a Training struct which takes of course a model context and optimizer context but also very importantly an example generator. The generator should have a specified number of threads which it uses to load and preprocess and specified number of samples before they are required by the training loop. It should be known whether the generator is finite or infinite (counting epochs on an infinite training set would of course be silly).

Here is the simplest implementation in python: https://github.com/justheuristic/prefetch_generator. Note that all loading and preprocessing is user-defined, it basically just implements the multiprocessing aspect. Is there already a crate we can use for this?

allow loading a numpy file as a constant node.

• MatMul: Dimensionality is not good when one of the Nodes is a 3d matrix(stack of matrices), make it more generalized like the dimension checking function
• Softmax: Problem during diff

As it is currently written, autodiff expects a scalar output node and the tests which present a vector output fail. We should extend the logic to be able to cope with the case of higher-dimensional output.

The XLA conversion certainly has a lot of boilerplate and it would be nice if there is a way to reduce that. Someone should also check if there are easy optimizations that can be done.