Behaviour in-place update of dataclasses

[GJBoth]

Hi,

I'm working on a model which requires a lot of memory, so I'm trying to make a few memory optimizations. The code is structured as functions updating a field, which I've implemented as a simple flax dataclass consisting of the field and some other stuff, i.e. like this:

class Field(struct.PyTreeNode):
    u: jnp.ndarray

When I modify the field I use the replace method to update the field, as its immutable, i.e.

field = Field(u)
u = field.u + f          # calculate new field
field.replace(u=u)  # update field

If I understood correctly, this code makes a copy of u, while the whole operation could be done in-place using the new .at syntax,

u = u.at[:].add(f) # in-place.

This is guaranteed to by in-place after jitting, and actually works with flax dataclasses too:

@jit
def update(field):
    field = field.replace(u=field.u.at[:].add(1))
    return field

but I suspect this isn't in-place, since when I run something like this

new_field = update(old_field)

old_field does not get updated.

What's the recommended way of dealing with in-place updates and dataclasses? Make the dataclass unfrozen? Or is there something possible with updating the replace method with the new .at syntax?

Thanks!

[marcvanzee]

When you jit the at operator, the XLA compiler will optimize an assignment like x = x.at[idx].set(y) so that no copy of the original array is needed. This saves memory. However, when you jit any function, the references to the passed arguments change, so you can't expect an outside variable to be modified in-place (this isn't allowed in a functional API).

In your example (where it seems you aren't jitting at all), in the update function you lose the reference to the object old_field, because you are rebinding the reference when you do field = .... So the outer reference will still point to the original object, meaning old_field isn't updated.

[GJBoth]

Thanks for your answer - I forgot to add the jit operator but added it now. I understand with the references why it doesn't work, but if I do something like this

field = update(field)

does the field get updated in-place, i.e. without copying field.u?

[jheek]

The dataclass itself is only something that exists at the Python level. JAX only has immutable arrays but XLA optimizes memory allocation and can do in-place updates.

There as some things you should keep in mind though:

• XLA by default cannot reduce memory of input arguments of jit/pmap unless you "donate them" using the donate_argnums kwarg. So for example even x.at[idx].add(1) will not work if x is an argument to jit(f) unless you use donate_argnums
• The dataclasses themselves don't have significant memory overhead but holding on to arrays unnecessarily (only outside of jit) can of course cause extra memory consumption. XLA doesn't care about when the Python code drops the reference for compiled code it will know when an array can be deallocated by itself.

[GJBoth]

Thanks! Are you sure about your first point with donate argnums? I can't find anything about it in the Jax documentation... Am I understanding you correctly when you say there's no "big" memory optimizations to be made? I.e. that code like this is as efficient as it'll be:

@jit
def forward(x)
field = create_field(x)
field = update1(field)
field = update2(field)
...
return field

[jheek]

Yes I would only expect an improvement from using @jit(forward, donate_argnums=(0,)).
The only other case I can think of is that update1/2 are sparse in some way so that they only update a subset of x. In that case the .at() syntax can help.
There are some differences between XLA backend btw. I'm not sure how well donate_argnums works outside of TPU

[GJBoth]

Alright, I'll see what that gives me. I'm pmapping this whole thing so I also have to be mindful of cpu->gpu transfers, so I'm not expecting too much improvement but let's see. I was hoping XLA wasn't as good as it is and I could make more improvements but well :-) Thanks for your help!