closure_convert for Pytree callables which hold reference to Python closures

[femtomc]

Hi all!

I've been working on a library for a bit -- and recently I've been considering how to inform users about usage of closures within the library.

One common idiom in my library is to wrap a function into a Pytree:

@lib.lift
def some_fn():
    x = ...

The lifting here raises the function into a Pytree datatype (let's call it Lifted) which supports a set of interfaces which I use during transformations. These interfaces often do some computation, and then return out a stored version of the Lifted instance in a datatype.

Sometimes, users try and close over values inside of these lifted functions -- when the values are JAX tracers, they get stored in the Python closure for the object -- and eventually I get tracer leaks. I think the reason for this is the returning of the Lifted instance -- note that Lifted is defined as a Pytree and normally works fine (I've been careful with flatten and unflatten so tracers are always treated as dynamic data).

However, I'd really like to allow users to close over arrays -- so I tried to alleviate this problem by using jax.closure_convert -- e.g. when someone wants to use the interfaces on Lifted -- I closure convert under the hood, use the converted function -- and pass in the captured tracer arrays as arguments.

However, I'm still getting tracer leaks.

Was closure_convert meant to be used in this situation? If yes, any guesses as to why I'm still encountering leaks?

Sorry I'm being quite vague wrt the actual implementation (source code is closed, for now). If necessary, I could hop on a call with someone and explain / show code.

Edit: re -- can you use closure_convert to define Pytree-compat environments for closures? E.g. where the "static" information is the Python callable -- and the dynamic info is the closed over Pytree environment?

If yes -- has someone done this somewhere? I'd love to inspect it.

Edit 2: I inspected the .__closure__ dunder to understand what sort of data is being held by my original closure, then the transformed version, compared to the auxiliary arguments that come out of closure_convert:

# Aux args
[Traced<ShapedArray(float32[5])>with<DynamicJaxprTrace(level=1/1)>]
# Original closure .__closure__
(<cell at 0x2800e5de0: DynamicJaxprTracer object at 0x2801587c0>,)
# Transformed closure .__closure__
(<cell at 0x2800e6170: list object at 0x280152f40>, <cell at 0x2800e6200: jaxlib.xla_extension.pytree.PyTreeDef object at 0x280152630>, <cell at 0x2800e6350: Jaxpr object at 0x280159620>, <cell at 0x2800e6230: function object at 0x28012eca0>, <cell at 0x2800e62f0: int object at 0x104ece830>, <cell at 0x2800e63e0: jaxlib.xla_extension.pytree.PyTreeDef object at 0x280149d70>)

So e.g. -- it actually seems like the transformed variant doesn't hold any DynamicJaxprTracer objects -- but I can't truly be sure I think.

Edit 3: I'm fascinated by the prospect of lifting closures to a Pytree compat representation. I tried the following thing:

@dataclasses.dataclass
class PytreeClosure(Pytree):
    callable: Any
    traced: Any

    def flatten(self):
        return (self.traced, ), (self.callable, )

def closure_convert(callable):
    captured = []
    for (ind, cell) in enumerate(callable.__closure__):
        captured.append(cell.cell_contents)
        cell.cell_contents = None
    return PytreeClosure(callable, captured)

def some_func(x):
    @closure_convert
    def _inner():
        return x

    return _inner

x = jax.jit(some_func)(5)
print(x)

where Pytree is a metaclass which registers the dataclass as a Pytree -- surprisingly, this works.

This seems like an awful kludge -- but honestly, I'd be surprised if one of the maintainers hasn't tried this before -- what are the sharp edges here?

E.g. I was considering defining __call__ for the closure by mutating the callable.__closure__ cells before invoking it. So then calling the closure would basically "put the arrays back in" before running the code. Then perhaps you'd also need to reset the environment (because Python doesn't support a native closure conversion transform on its closures, as far as I can tell).