Computing Jacobians with ragged arrays (Awkward Array)

[agoose77]

Hi all!

I'm a developer over at the https://github.com/scikit-hep/awkward-1.0/ project.

From our docs

Awkward Array is a library for nested, variable-sized data, including arbitrary-length lists, records, mixed types, and missing data, using NumPy-like idioms.

I recently started looking at our prototype JAX integration, and I'm looking to add the ability to compute the Jacobian of a function, e.g.

def func(array):
    return ak.concatenate([array, array])

Awkward Arrays are built by composing layout nodes, so a NumPy array with shape (4, 3, 2) is built as:

<RegularArray size='3'>
    <content><RegularArray size='2' >
        <content><NumpyArray dtype='int64' len='24'>
            [ 0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21
             22 23]
        </NumpyArray></content>
    </RegularArray></content>
</RegularArray>

Our initial approach to integration with JAX was to represent Awkward Arrays as singular atoms with a number of leaves, i.e. the example array above would be flattened into

([array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16,
       17, 18, 19, 20, 21, 22, 23])], <AUX DATA>)

This means that both the input and output arrays of e.g. func are flattened in JAX terms into flat single buffers (in this case) alongside some auxillary data.

However, to make it easier to reason about, I've mocked up a demo where our entire layout nodes can be converted into trees (not just the root node). In this demo, there is only NumpyArray and RegularArray support, even though we have a wider range of node types.

In the example, the array is not ragged, so the input array has shape (4, 3, 2), and the output of func(array) should have shape (8, 3, 2). Therefore, the Jacobian should have shape (8, 3, 2, 4, 3, 2). Using our layout system, we'd expect the following structure:

<RegularArray size='3'>                           # (8, 3, _, _, _, _) 
    <content><RegularArray size='2'>              # (_, 3, 2, _, _, _)
        <content><RegularArray size='4'>          # (_, _, 2, 4, _, _)
            <content><RegularArray size='3'>      # (_, _, _, 4, 3, _)
                <content><RegularArray size='2'>  # (_, _, _, _, 3, 2)
                    <content><NumpyArray dtype='float32'>...</NumpyArray></content> 
                </RegularArray></content>
            </RegularArray></content>
        </RegularArray></content>

    </RegularArray></content>
</RegularArray>

In the comments, a helpful key shows how the size (RHS and length LHS) are computed from each RegularArray node.

What I'd like from JAX is to ask our tree hooks to:

1. unflatten the argument node (with actual shape (4,3,2)) with the flat 8*3*2*4*3*2=1152 Jacobian values
2. unflatten the result node with (1)

I have played around with this a bit already, but it's clear to me that I'm not 100% confident on the details of what JAX is doing. Rather than guessing, I was hoping that I might get some pointers from the JAX team as to how to proceed from here.

Please feel free to follow up questions if anything needs explaining further. Thanks!

[jakevdp]

Hi - happy to help with this, but I'm not sure how to interpret your XML-like syntax for array representation. What do the fields mean? How should I interpret (8, 3, _, _, _, _) on an outer field? Is this something like a jagged array with five rows where the first row is size 8 and the second row is size 3? What do the underscores mean? And how should I understand the difference between a RegularArray and a NumpyArray?

(Sorry if this is standard stuff... if you have docs about these representations you can point to, I may be able to read up and better understand what you're asking)

[agoose77]

Oh thanks. I misread the tree_unflatten implementation as abstract in the _base module.

[jakevdp]

I think either approach (1) or (3) could be really interesting in this case... let me know if you want to do a video chat to brainstorm / pair-code. It could be fun 😀

[agoose77]

Thanks Jake. I think the final two options look like those that I was expecting. I've been poking around in the JAX source, and as I understand it, the reason that I'm seeing square 2D arrays in the unflatten routines is because JAX is doing the equivalent of jnp.eye(len(y)), where y is the output of the function i.e. in the cookbook. Then, the VJP which maps from the codomain to the domain is batched using vmap, which produces an array of (len(y), len(x)).

In your suggestion to change the representation, I am not confident that such a thing would be possible, at least as I understand it. Awkward Array is built around composition of layouts, and this system is fairly well established.

That said, I know that I do not have a strong grasp yet on what I expect JAX to do here. In particular, what does it mean for JAX to unflatten an existing node with new data of an additional dimension. If I can understand this part better, then it might be feasible to map from JAXs representation into our interpretation, in such a way that we can reverse the interpretation in tree_flatten when JAX needs to read the flattened tree again.

[agoose77]

@jakevdp hmm, GitHub didn't update with your most recent response. We're trying to get a major release out in the next couple of months, so we don't have a huge amount of time to work on this right now. However, I can tell that it would be very productive (for me) to get a sense of how to think about this, and your offer is most appreciated. Are you around next week some time before Friday? That might be a bit presumptuous...!

[jakevdp]

Sounds good - I'll send you an email at the address listed on your github profile