Question: what's the best way to support symbolic execution of JAX functions?

[Z80coder]

Hi,

For reasons I need to symbolically evaluate JAX functions. E.g. Say I have a JAX function f that takes a 1-d array then I'd like to pass a list of strings that represent input values:

f(['x1', 'x2', ...])

and return a string that represents the result of the computation

relu(0.32 x1 + 1.45 x2 + ...)

In general, I want to support all JAX functions, and also support symbolic weights.

I've experimented with rolling-my-own "shadow" functions. But this won't scale, because I need to re-implement the JAX api.

Instead, I'd like to support strings and string operations at a deeper level in the JAX system, so all the existing functionality can operate on symbolic expressions.

Any pointers much appreciated!

[jakevdp]

Hi - thanks for the question. This is pretty far-afield from what JAX is designed to do, but you may be able to take advantage of JAX's transformation framework to define a custom transform that would do this kind of thing. I don't have any pointers to relevant documentation or examples, because any such solution would rely on a lot of internal APIs.

Edit: here's a project with a similar goal that you might use and/or get ideas from: https://git.informatik.uni-hamburg.de/0moin/google-research/-/tree/master/jax2tex. I don't think the code is maintained, though.

[Z80coder]

Thanks for the reply!

I had seen jax2tex and was wondering if this was the way to go. So thanks for pointing me to this!

I plan, therefore, to generate symbolic representations by interpreting the jaxprs.

[patrick-kidger]

If you're looking for a way to pass strings across JAX API boundaries, then equinox.internal.str2jax might be of interest. It returns a new PyTree whose repr is the specified string. You could then do string manipulations on them, by casting each way:

from equinox.internal import str2jax
from jax import jit

hello = str2jax("hello")
world = str2jax("world")

@jit
def concat_strings(x, y):
    return str2jax(str(x) + str(y))

concat_strings(hello, world)

All the string manipulations necessarily happen during trace time, i.e. they're "static". If you're looking for something deeper then that then I don't think that'd really be supported by any JAX code that is lowered to XLA (i.e. jit'd in the normal way).

[Z80coder]

Thanks for the pointer! Very relevant to my problem!
Would applying jax.vmap to functions that manipulate str2jax strings "just work"?

[patrick-kidger]

Yes, but the strings are always broadcast rather than vectorised (i.e. like passing in_axes=None for the string arguments).

I've not thought about a way to vectorise over strings. It might be possible using jax.interpreters.batching.is_vmappable (IIRC) and a bit of ingenuity, but you're in complicated/unknown territory there.

[Z80coder]

@jakevdp @patrick-kidger

In case this is relevant to others I implemented a symbolic evaluator for JAX expressions in this repo:

• db-nets

The relevant python files are:

• symbolic_generation.py
• symbolic_primitives.py

Note that I've implemented symbolic versions only for the primitive JAX operations that I need. But it's not difficult to add more. So the overall structure might be useful to others.

I construct neural nets that learn discrete boolean functions, and then use the symbolic evaluator to extract the boolean function that's been learnt. See the (draft) paper:

• "∂B nets: learning discrete functions by gradient descent" (April 2023).

Thanks for writing JAX!

[patrick-kidger]

Interesting!
What's the distinction between this and make_jaxpr? Which looks to me like it does something very similar, building up a symbolic graph.
(I can see you have some partial evaluation in there too -- the all_concrete_values check -- and whilst JAX only has a semi-public API for this, this does exist in JAX too, under jax.interpreters.partial_eval.trace_to_jaxpr.)

[Z80coder]

I use make_jaxpr to get the jaxpr. And then I symbolically evaluate the jaxpr with symbolic inputs to yield a string expression that is the composite function that the jaxpr computes.

Thanks for the pointer to jax.interpreters.partial_eval.trace_to_jaxpr. I'm sure I'm missing lots of other tricks due to my lack of deep knowledge of the inner workings of JAX.