abstracted_axes and eval_jaxpr

[erick-xanadu]

Hi,

I am looking now into abstracted_axes and the eval_jaxpr function. The function has the following signature:

def eval_jaxpr(jaxpr: Jaxpr, consts, *args, propagate_source_info=True):

I am trying to eval_jaxpr a jaxpr computation that was generated with abstracted_axes={0: "n"}. This means that the the jaxpr computation will have an implicit parameter corresponding to the value of the axis 0. But eval_jaxpr does not know about this implicit argument. I think that a possible good way to solve this is for eval_jaxpr to obtain the implicit arguments from args. However, before implementing this, I would like to ask if this could be a good way to solve this problem or if you have some other alternatives in mind.

 436 def eval_jaxpr(jaxpr: Jaxpr, consts, *args, propagate_source_info=True):
 437   def read(v: Atom) -> Any:
 438     return v.val if isinstance(v, Literal) else env[v]
 439 
 440   def write(v: Var, val: Any) -> None:
 441     if config.jax_enable_checks and not config.jax_dynamic_shapes:
 442       assert typecheck(v.aval, val), (v.aval, val)
 443     env[v] = val
 444 
 445   env: dict[Var, Any] = {}
 446   map(write, jaxpr.constvars, consts)
 447   # Here, can I add the implicit input variables from args and perhaps some other sources of information?
 #         E.g., args.insert(0, args[0].shape[0])
 448   map(write, jaxpr.invars, args)

Thanks!

[mattjj]

Thanks for the question!

This means that the the jaxpr computation will have an implicit parameter corresponding to the value of the axis 0.

Actually, the parameter isn't implicit in the jaxpr; it appears as just an ordinary parameter, and correspondingly the caller can just pass it as an ordinary argument:

import jax
import jax.numpy as jnp

jax.config.update('jax_dynamic_shapes', True)

def f(x):
  return jnp.sin(x) + jnp.cos(x)

jaxpr = jax.make_jaxpr(f, abstracted_axes={0: 'n'})(jnp.arange(3.)).jaxpr
print(jaxpr)
# { lambda ; a:i32[] b:f32[a]. let
#     c:f32[a] = sin b
#     d:f32[a] = cos b
#     e:f32[a] = add c d
#   in (e,) }

from jax._src.core import eval_jaxpr
ans, = eval_jaxpr(jaxpr, (), 3, jnp.arange(3.))
print(ans)  # [1.         1.3817732  0.49315056]

The axis size arguments (and outputs) are only implicit at the Python tracing level, but everything is made explicit at the jaxpr level. That means things like the dispatch path for jitted computations needs to infer axis sizes from arguments, as you describe, but not eval_jaxpr.

That said, we're not actively working on this dynamic shape stuff right now, so there's probably lots of other stuff that doesn't work.

WDYT?

[soraros]

Pallas looks very interesting indeed. I'm wondering if its internals will be made public (Ref etc). I'm asking because I also see JAX as a programming language project, a more production ready Dax if you will. The work on mutable state, effects and dynamic shapes are all expending JAX's type system, which make JAX more versatile and interesting as a language.

Thanks again for the reply! It's a bit off topic for this thread, it would be great if we could continue it somewhere else though.

[erick-xanadu]

Thanks @mattjj , your answer prompted me to look further and I found _extract_implicit_args(in_type, explicit_args) and this can be applied to the context where I am working on.

[mattjj]

@soraros Did you have a particular sense of 'public' in mind? Like, use within broader JAX and not just pallas_call? Or did you mean some kind of documentation?

[soraros]

@mattjj I meant use within broader JAX.

[mattjj]

Yeah! Actually those Ref and state effect features were developed before Pallas, and Pallas just adopted them. We were going to use them to offer a more flexible loop alternative than scan, which we sometimes called five_loop or bloop. And sort of as a special case we were going to provide run_state : (Ref h a -> {State h} ()) -> a -> a too. But then working on the Triton integration, we had the idea for Pallas, and then basically switched all our energies to developing Pallas.

I suspect we'll come back and finish (and document) five_loop/bloop and run_state someday. There are quite a few tests already. run_state in particular may be useful for plumbing values out of the backward pass of grad.

Anyway, all that is to say: yes I expect to see them more broadly used in JAX, but I don't have a timetable.