Re:Awkward: Nested, jagged, differentiable, mixed type, GPU-...
by jpivarski on Thursday December 16, 2021 @12:45PM
(#29581177)
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.