Can I define a deep neural network as an Awkward array of matrices of different sizes? It looks very promising to compute a backpropagation step without ping-ponging with the Python interpreter.
I hadn't considered an Awkward array being used as a neural network, and I would expect specialized neural network frameworks to optimize it better: don't they fuse the matrices representing each layer into a single outside-of-Python computation? If they don't, I wonder why not.
Nevertheless, my ears perked up at "matrices of different sizes," since that was something I implemented a while back thinking it would find a use-case if it were well advertised.
You can matrix-multiply two arrays of matrices, in which all the matrices have different shapes, as long as matrix "i" in the first array is compatible with matrix "i" in the second array. Like this:
The left[0]'s 3×2 times right[0]'s 2×3 results in the output's 3×3, etc for i=1 and i=2. This is the kind of operation you wouldn't even be able to talk about if you didn't have jagged arrays.
I meant to add that autodiff is a feature we'll be working on this upcoming spring, by using JAX arrays as the buffers. Until we have autodiff, backpropagation won't be possible.
Nevertheless, my ears perked up at "matrices of different sizes," since that was something I implemented a while back thinking it would find a use-case if it were well advertised.
You can matrix-multiply two arrays of matrices, in which all the matrices have different shapes, as long as matrix "i" in the first array is compatible with matrix "i" in the second array. Like this:
Matrix-multiplying them results in this: The left[0]'s 3×2 times right[0]'s 2×3 results in the output's 3×3, etc for i=1 and i=2. This is the kind of operation you wouldn't even be able to talk about if you didn't have jagged arrays.