Question: Constructing a sparse matrix

How can I make a function like this work with autograd, given that assignment is forbidden?

def func(a, b):   # a and b are scalars
    m = np.zeros((10, 10))
    m[5,7] = a
    m[3,6] = b
    return m

[mattjj]

Here are two ways to do it. Because func is a linear map, you can implement it in a sparse matrix form treating a and b as coefficients on fixed basis elements:

def func(a, b):
    def one_hot(i, j):
        out = np.zeros((10, 10))
        out[i, j] = 1
        return out
    return a * one_hot(5, 7) + b * one_hot(3, 6)

If the indices (5, 7) and (3, 6) are really constants (or you’re going to call func many times with the same values for those indices) you can make that more efficient with staged evaluation, so that the basis elements aren’t re-allocated on every invocation.

Another way to do it is to write your own primitive, which is probably both more straightforward and more performant:

@primitive
def func(a, b):
    m = np.zeros((10, 10))
    m[5, 7] = a
    m[3, 6] = b
    return m
func3.defvjp(lambda g, ans, vs, gvs, a, b: g[5, 7], argnum=0)
func3.defvjp(lambda g, ans, vs, gvs, a, b: g[3, 6], argnum=1)

Inside a primitive, there are no autograd boxed values; everything is raw numpy, so you’re back in the world you’re used to. You have to define any vjp’s you need for your primitives, though.

When writing custom primitives, check_grads is your friend!

Side notes: of course, you only need to do this if you’re going to differentiate through func. Assignment into autograd’s ArrayNodes is not supported, but when you’re using autograd.numpy, ArrayNodes are only created as necessary during tracing, i.e. when an autograd.numpy function gets an ArrayNode as one of its arguments; otherwise, it’s just raw numpy as usual. In fact, the kind of assignment you wrote might work anyway, since m is being created as a zeros matrix and won’t be an ArrayNode, though when a or b is an ArrayNode, the assignment would have to turn it into an array with dtype object and things might get weird.

[mattjj]

If you are differentiating with respect to a and b directly, rather than them being some intermediates in the computation, you can also get them both at once (without redoing the forward pass or backward pass) using multigrad, which is in convenience_wrappers.py. Check out that file, since there is a lot of functionality in there.

You can also do the same thing manually by writing a function that takes an (a, b) tuple argument instead of separate arguments (since autograd will happily differentiate with respect to a tuple argument).