Add a way to get the tree size of a jaxpr

[carlosgmartin]

Feature request: Add a way to get the size of a Jaxpr and ClosedJaxpr, in terms of the size of the tree itself (not the length of its string representation).

This could be done by adding a __len__ dunder to these classes, or another method.

[superbobry]

Can you add a few details about what you are trying to do, please?

[jakevdp]

Does len(jaxpr.eqns) do what you want?

[carlosgmartin]

@superbobry I was looking for some rough measure of jaxpr "size" or "complexity"/"simplicity" in order to compare different potential implementations of a JAX function, alongside cost analysis. (See this comment, for example.)

@jakevdp No.

import jax
from jax import numpy as jnp
args = jnp.zeros([2, 3, 5]), jnp.zeros([2, 7, 5], int), 1., 1, False
jaxpr = jax.make_jaxpr(jnp.put_along_axis, [3, 4])(*args)
print(len(jaxpr.eqns))  # 1

[jakevdp]

@jakevdp No.

Maybe if you disable JIT it will give you something closer to what you're after?

from jax import numpy as jnp
args = jnp.zeros([2, 3, 5]), jnp.zeros([2, 7, 5], int), 1., 1, False
with jax.disable_jit():
  jaxpr = jax.make_jaxpr(jnp.put_along_axis, [3, 4])(*args)
print(len(jaxpr.eqns))  # 24

[jakevdp]

Side note though: jaxpr size is not a very good proxy for complexity in general: for example, jaxprs do not do any dead code elimination. Given this, compiler-based measures like cost analysis will probably be more robust.

[carlosgmartin]

@jakevdp That's closer to what I'm looking for!

The reason I brought this up is because of your comment:

I think we should optimize for a simple jaxpr here, not optimize for runtime on a particular device (which may or may not reflect runtime on other device types, or with other array sizes).

which could also be good advice for other functions beyond put_along_axis.

[jakevdp]

Oh OK – I don't think "simple jaxpr" is necessarily the right metric in general, but in that specific case we were talking about implementing put_along_axis, which I believe could be implemented in most cases via a single scatter call. In that particular case, we should aim to write that single scatter call.