Understanding jnp.bincount jaxpr

[limarta]

Hi,

I'm trying to write an analogous ONNX computation for a jax function. One crucial step is computing jnp.bincount for a set of integer assignments, and to understand the underlying computation I am looking at the jaxpr:

{ lambda ; a:u32[2] b:i32[3] c:f32[3]. let
    d:i32[3] = broadcast_in_dim[
      broadcast_dimensions=()
      shape=(3,)
      sharding=None
    ] 0
    e:i32[3] = pjit[
      name=clip
      jaxpr={ lambda ; f:i32[3] g:i32[]. let
          h:i32[] = convert_element_type[new_dtype=int32 weak_type=False] g
          i:i32[3] = max h f
        in (i,) }
    ] b 0
    j:i32[] = device_put[
      copy_semantics=[<CopySemantics.ALIAS: 1>]
      devices=[None]
      srcs=[None]
    ] 1
    k:bool[3] = lt e 0
    l:i32[3] = add e 3
    m:i32[3] = select_n k e l
    n:i32[3,1] = broadcast_in_dim[
      broadcast_dimensions=(0,)
      shape=(3, 1)
      sharding=None
    ] m
    o:i32[3] = broadcast_in_dim[
      broadcast_dimensions=()
      shape=(3,)
      sharding=None
    ] j
    p:i32[3] = scatter-add[
      dimension_numbers=ScatterDimensionNumbers(update_window_dims=(), inserted_window_dims=(0,), scatter_dims_to_operand_dims=(0,), operand_batching_dims=(), scatter_indices_batching_dims=())
      indices_are_sorted=False
      mode=GatherScatterMode.FILL_OR_DROP
      unique_indices=False
      update_consts=()
      update_jaxpr={ lambda ; q:i32[] r:i32[]. let s:i32[] = add q r in (s,) }
    ] d n o
  in (p,) }

Some pieces of the jaxpr make sense, but steps l and p seem bizarre. What is the motivation for adding 3 to e and how does this help in the scatter-add step? The mode is set to FILL_OR_DROP which probably responsible for getting the right answer.

Could someone walk me through how an array such as [0,0,1] with length=3 would work? This is what I get:

d=[0,0,0]
e=[3,3,4]
o=[1,1,1]
p=??? # should be [2,1,0]

[jakevdp]

The l step comes from supporting negative indices in array indexing expressions: along with the surrounding operations, it comes from something like i = jnp.where(i < 0, i + x.shape[0], i)

The p step is a translation of array indexing into the semantics of the scatter operation, which is how indexing in JAX is implemented.

To be honest, I would not suggest going down the path of trying to reverse-engineer scatter. It's not really a human-consumable operation.

Can you say more about what your ultimate goal is here? There may be a better path forward than trying to parse the meaning of scatter.

[limarta]

I have a small modeling DSL which produces jax functions to the end user. For pedagogical reasons, I'm looking to host these functions function as ONNX models on a webpage so that users can run the code on their local GPU, relieving me of computational burden.

I know jax2tf.py translates a jax function into a set of tensorflow operations which we can then call tensorflow-onnx on. I tried using it ran into ModuleNotFoundError: No module named 'distutils' from jax2tf.py's tensorflow depedency. I ended just hand-rolling my own jaxpr interpreter for this and it works correctly for my use case (no guarantees on anything out of scope). I just wish I didn't have to do this though. Maybe if I can get jax2tf.py running, that would be a more reliable solution right?

[limarta]

The l step comes from supporting negative indices in array indexing expressions: along with the surrounding operations, it comes from something like i = jnp.where(i < 0, i + x.shape[0], i)

Yea I think I realized what I did wrong. I flipped the true and false args in m:i32[3] = select_n k e l.

[jakevdp]

I see, thanks! This is currently a known issue with higher-level computations that are constructed from multiple primitives: it's hard to know from the jaxpr (or more generally the HLO) whether a particular sequence of operations implements something particular like bincount or softmax.

Currently the only good approach for this is pattern matching: e.g. in XLA-GPU, the compiler pattern-matches the operations output by jax.nn.softmax in order to lower to a device-specific softmax kernel. This is not the most robust solution, unfortunately. There's some ongoing thought about how to add metadata in situations like this to preserve the higher-level information of compound operations: the new lax.composite primitive is one example. @dfm may be able to say more about plans in this area.

[limarta]

Interesting! Thanks for the info.

Currently the only good approach for this is pattern matching: e.g. in XLA-GPU, the compiler pattern-matches the operations output by jax.nn.softmax in order to lower to a device-specific softmax kernel.

Would just translating every single primitive and recursively entering any subj-jaxprs work? Even if the emitted code may not be as well optimized as a jitted-version - which is ok for an in-browser demo, would it in principle be correct? Or are there pitfalls to this approach?