Repository navigation
Scatter gradient implementations #282
Description
Activity
- changed the title
[-]Best way to interleave two tensors and propagate gradients (can you propagate gradients through assignments?)[/-][+]Scatter gradient implementations[/+]on Dec 25, 2023 Scatter vjp is indeed missing from
mainat the moment. I started an implementation in the branchscatter-vjpso that you can be unblocked. The VJP for assignment is the only one implemented so far and not particularly well tested.Depending on the number of assignments and the size of the assigned blocks concatenating slices might not be too bad of an implementation however.
On that branch the following worked fine:
import mlx.core as mx def fn(x): output = mx.zeros([4, 128]) output[mx.array([0, 1])] = x return output.square().sum() x = mx.random.normal(shape=(2, 128)) print(fn(x)) # array(274.507, dtype=float32) grad_fn = mx.grad(fn) print(grad_fn(x)) # array([[1.08186, -1.9959, 1.18484, ..., 2.52877, 0.80407, 0.529424], # [-0.74177, -1.64652, 0.113564, ..., -1.66786, 1.31119, -0.58756]], dtype=float32)
Reacted by c0gClosing as this is now merged. Scatter multiply, max, min are still missing but assign and add are there.
Not sure if this is an issue or I'm just holding it wrong :-) I've got an alternative approach that seems to work with gradients by building up my output up by catting slices into my tensors, which does seem to work for the grads.
I have two tensors (token_embeddings and image_features), I want to interleave them in some way. It's for a llava like model, so my input tokens are like
where
<image>is a sentinel token. What I've got to so far is making my output a flattened tensor[batch * max_seq, dim]my token/image tensors[num_tokens, dim]and[num_images, dim]and forming indices to assign to, ending up looking like this:It doesn't seem like I can get gradients for this though:
gives me
Like I said, it seems like there's an alternative in catting things since:
seems to work.