Making scan() and cond() functions more usable

[MRiabov]

Sup Jax,
I was writing a function to segregate values based on predicate, like this:

def _segregate_parts_by_type(parts, hitBy:PositionHitBy, specific_hit_object): # would likely return only a handful of parts
    def segretate(lists:tuple[list,list],part):
        def _segregation(lists:tuple[list,list],part):
            index=int(CreatedBy[part.properties.creation].value)
            lists[index].append(part)
            return lists, part
        return jax.lax.cond(
            jax.numpy.isin(part.id,specific_hit_object), # I know it could be one-lined, but there are more functions to come.
            _segregation(lists, part)
        )

Which has three defs in a row, which looks like abomination, and...

What if we make scan and maybe even cond something like this:

 with jax.lax.scan(init=jax.numpy.zeros(1000,1000), xs=inputArray):
     do_func()

Additionally to having jax.lax.scan maybe there would be a wscan which is identical, except it can be used in such a with statement?
Same goes for cond() which could be somehow reused to work?

Think of how much easier it would be to read and write:

with jax.wvmap(x=input): # wvmap for "with"
    s,a,r,s1 = run_ml_agent(x.at[0])
    r_stats=jnp.mean(r)
    s1_stats=jnp_mean(s1)

We won't have to write terse code with defs, then make explicit calls to them, which would cleanup code well.
I propose this as a highly useful feature, everyone would benefit if jax team implements it.

[jakevdp]

Interesting idea! Do you have a sketch of how the implementation of these context managers might look?

[MRiabov]

Hello, sorry, I'm not into much compiler writing - but one library I use (for 3d engineering) uses a lot of them as BuilderAPI, so they might have something for it.
See first "tea cup" example here.

As they chain dozens of shapes and commands with it, I would think Jax can do that too.

(I also have no idea how Jax optimizes commands under the hood, so maybe it's something with that?)

Edit: Could context managers like this simply create anon-like functions and simply automatically substitute it as a function?

So, a context manager would effectively do a def and pass it to the vmap. Something like that.

[jakevdp]

Fortunately I don't think the compiler would have to be involved here, any context manager API would be entirely implemented in Python, as a wrapper that would output a standard scan call. Did you have ideas on how that could be implemented?

[MRiabov]

No, sorry, I was thinking about it only as a user. "I want it to work this way, so they probably would know an approach." I'm from mech and ML eng so you probably know better.

…

On Mon, Sep 9, 2024 at 6:01 PM Jake Vanderplas ***@***.***> wrote: Fortunately I don't think the compiler would have to be involved here, any context manager API would be entirely implemented in Python, as a wrapper that would output a standard scan call. Did you have ideas on how that could be implemented? — Reply to this email directly, view it on GitHub <#23487 (reply in thread)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AZZOTDZAVCEPWNL2A6EZKWLZVXH7PAVCNFSM6AAAAABNZZ6VOCVHI2DSMVQWIX3LMV43URDJONRXK43TNFXW4Q3PNVWWK3TUHMYTANJZGM4DANQ> . You are receiving this because you authored the thread.Message ID: ***@***.***>

[jakevdp]

OK - in that case, could you say more about what exactly you have in mind? It's not clear to me how the context manager you wrote might translate to a scan or vmap operation. Can you write a scan operation along with the equivalent context manager under the API you have in mind?

[MRiabov]

Well, to be fair I don't know, but I know that be a nice syntax to work with. Maybe convert it to an array of lambda operations under the hood? If that's possible at all. Cheers.

…

On Mon, 9 Sept 2024, 18:33 Jake Vanderplas, ***@***.***> wrote: OK - in that case, could you say more about what exactly you have in mind? It's not clear to me how the context manager you wrote might map to a scan operation. — Reply to this email directly, view it on GitHub <#23487 (reply in thread)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AZZOTD5Q75HOZGZLTN4ZE3TZVXLYLAVCNFSM6AAAAABNZZ6VOCVHI2DSMVQWIX3LMV43URDJONRXK43TNFXW4Q3PNVWWK3TUHMYTANJZGQYTENA> . You are receiving this because you authored the thread.Message ID: ***@***.***>

[jakevdp]

but I know that be a nice syntax to work with.

I'd like clarification on this, because I don't entirely understand the proposed syntax. When you write this:

 with jax.lax.scan(init=jax.numpy.zeros(1000,1000), xs=inputArray):
     do_func()

what is your intended output? i.e., what arrays exist in the namespace after this is executed, and what values do they contain? If you could clarify that, we might be able to start talking about implementation and whether it's feasible.

[MRiabov]

Hello, I actually (think) I have an idea on how this can be implemented.

Fortunately I don't think the compiler would have to be involved here, any context manager API would be entirely implemented in Python, as a wrapper that would output a standard scan call. Did you have ideas on how that could be implemented?

During JIT parsing, jax does store the parse results somewhere, right?
So, if there is a context present that says "this is a piece of code that must be treated as a separate function", you would be able to treat it as a separate function, and later assign it to scan, or vmap, or whatever.

The intended output is:
This:

with jax.lax.scan(init=jax.numpy.ones(1000,1000), xs=inputArray):
     jnp.add(carry, 1),
     x=jnp.dot(x) // concise

Equals to this:

def func(carry, x):
     carry=jnp.add(carry, 1),
    x=jnp.dot(x)
    return x, count
result, _ =scan(func, init=jax.numpy.ones(1000,1000), xs=inputArray)  //two lines longer

Now I do see that jit would be slightly less efficient in a sense that we don't know which values are passed in and returned out, but tracing can probably deal with it, right?

So, you'd essentially say "we want code which runs with this context manager to be treated as if it was in a separate function, which we can optimize for scan later."
I understand that Jax does not care which lines to optimize, it always computes jit (where marked), so a context which says "please use this function in scan" would be reasonable.

Same for vmap and other loops.

(Get it? I can clarify.)

[jakevdp]

That looks interesting, but I'm not sure whether it would be possible given the constraints of the Python language. As you've written it, carry and x are undefined. I don't think there's any way to implicitly inject variable names into a block of code via a context manager syntax, so we'll need other ideas on how to express this kind of API.