Gather operation does not handle most cases

[rakadam]

Gather operation is basically broken, for example x[:, 0] works correctly, but x[:, 1] is broken (although this exact bug is fixable with a few lines of extra code).
Also, most arguments passed to the gather operation are ignored, which results in assuming the simplest most general version, always.
So most of the slicing in JAX which gets turned into lax.gather does not work correctly.

Other examples:
This is working: x[:, 0, :]
this one x[:, :, 0] this does not work, and actually gets compiled into the same as x[:, 0, :]

We need jax2onnx for our production release soon(ish), so if my time allows I will make the correct implementation into gather.py myself. I am reporting this issue so you know that I am starting to work on this. If somewhere there is already a good implementation please warn me so I don't work needlessly.

[enpasos]

Welcome, and thanks a lot for reporting this — and even more for offering to work on a fix! Please don’t hesitate to ask if you need any pointers or help along the way. Just to note: the upcoming 0.9 release will be based on the onnx_ir_and_jax_072 branch, which I’ll merge into main in the next few days.

[rakadam]

This proved to be a very complex endeavour, because lax.gather can represent many operations. So I wrote a compiler just for the gather operation. The front-end and middle-end are done, and are mostly working, and the back-end, the ONNX-IR generator is still TODO. The good news is that I figured out basically how to map any possible lax.gather operation provided it is in mode=promise_in_bounds and already implemented the hard parts.

a single lax.gather can be compiled for example into a series of ONNX operations like this:
Slice, Transpose, GatherND, Transpose

Or in other cases:
Transpose(on input tensor), Transpose(on the index tensor), GatherND, Transpose(on output tensor)

I figured out how to detect and convert the weird gather-slice hybrid operation too, which lax.gather can do. It is not always efficient though.

Also, important: I had to constant fold the jaxpr for the index tensor, because by far the most often used lax.gather-s use a constant tensor, and they may contain slicing that can be decomposed from the gather, this can only be sanely detected from the constant tensor. My constant folding code is very simple, but I had to modify the interface a little and allow passing the literals and jaxpr to the lower function.

I am constantly tracking the onnx_ir_and_jax_072 branch.

[enpasos]

Thanks a lot for diving into this, @rakadam — that’s serious work and a big step forward. I’m excited to see your solution, and it’s great that you’ve been tracking the onnx_ir_and_jax_072 changes. If anything blocks you, just ping me.

[enpasos]

Quick update: onnx_ir_and_jax_072 is now merged into main.

[justinchuby]

If you are looking to support advanced indexing, feel free to look at https://github.com/microsoft/onnxscript/blob/59c3d32ea0cf18fbd348d8b4e23fdb8dad6427ea/onnxscript/function_libs/torch_lib/ops/core.py#L4055

[rakadam]

I am handling almost all cases, and tested for the most usual cases. So I am pretty sure my code is working for the simple cases.

Few things:

1. I had to add extra arguments to the lower, because I needed early (but simple) constant prop: def lower(self, ctx: "IRContext", eqn, jpr, literals): # type: ignore[name-defined]
2. I wrote a new class named LowerDimExpr for lowering complex symbolic shapes into ONNX IR. The most usual complex shape is B-1 but my code should work for some pretty crazy expressions too. It has some internal caching, so it will need to be moved to the IR builder to be persistent. (the caching makes the ONNX graph way less crazy looking)
3. The case I am not handling is merging in the dependent complex slicing into the gather indices for parametric (not a constant index) for lax.gather

Questions:

1. Should I move my LowerDimExpr into the IR Builder myself before sending a PR? Should be pretty simple to do.
2. Should I keep the few lines of modification in the code which lets the extra arguments jpr, literals work before sending a PR?

[benmacadam64]

So, I wonder if at this point it makes more sense to map to the stablehlo gather operation, and then use the StableHLO's compiler passes to simplify the gather operation? The same trick would work for operations like dot_general, etc.

Of course, that turns into a shlo2onnx layer pretty quickly, but maybe it's worth it, since it's only ~100 symbols to cover and onnx-ir's tape thing handles SSA.

[rakadam]

I think I have done all the difficult simplification and mapping in my code for gather, that just need to be marked in.
I am only not handling a very weird case, which is very rarely used.

Does dot_general need that complex case handling?

[enpasos]

Thanks for driving this, @rakadam — really appreciate it.

• Your decomposition strategy (Slice/Transpose/GatherND sequences and the gather–slice hybrid) is exactly the direction we want. We prefer a direct JAX→ONNX mapping over introducing a StableHLO layer (thanks @benmacadam64 for the idea) because it:

  • Keeps the converter IR-only and avoids another abstraction to maintain.
  • Improves debuggability and expect_graph parity (less indirection).
  • Reduces dependency surface and avoids tracking StableHLO passes/semantics drift.

• For constant-folded index tensors, let’s keep the plugin API as lower(self, ctx, eqn[, params]). Use IRContext utilities (ctx.get_value_for_var, ctx.bind_const_for_var) to fold constants. If you hit a gap, we can add a minimal helper to IRContext rather than threading ClosedJaxpr/literal maps through plugins.

• Re: LowerDimExpr: yes, please move it into the IR builder area (a small shared helper next to the builder) so caching is persistent and reusable.

• It’s fine to land without the rare “parametric dependent slicing merged into indices” case. Document the current limitation in the PR; we can handle it in a follow-up.

• Scope guideline: plugins should only receive the context they need for their operation; avoid passing broader calling-context details unless necessary.

An early PR against the main branch would be great — we can quickly iterate on any small IRContext hooks and finalize placement/naming for LowerDimExpr.

[rakadam]

If we are keeping the lower steps pure without any constant folding that makes handling some of the more usual cases way more complex, and way less efficient, later ONNX IR unlikely to optimize it.
My point is that there are cases where we can only recognize some slicing by looking into the index tensor, which is usually constant for these operations. But jaxpr makes a relatively complicated calculation to make these constant tensors, so I cannot access the constant tensor from the lower() because it is not made yet.

So the only solutions I see:

1. intentionally generate really deoptimized code, which likely won't be optimized later, because we are growing (editing) the index tensor by one or more indices. (significant size and complexity increase!)
2. do a simple constant folding in the IR Builder optionally.

I propose adding this constant folding algorithm to the ctx / IR Builder code directly, since it is so simple.

my current constant folder code is extremely simple:

class ConstantFolder:
    def __init__(self, jaxpr, const_vals):
        # Build producer map once
        self.producer_map = {}  # var_id -> equation
        for eqn in jaxpr.eqns:
            for outvar in eqn.outvars:
                self.producer_map[id(outvar)] = eqn
        
        # Initialize known constants
        self.known_constants = {}
        for var, val in zip(jaxpr.constvars, const_vals):
            print(var, id(var), val)
            self.known_constants[id(var)] = val

    def eval_jax_primitive(self, primitive, inputs, params):
        try:
            return primitive.bind(*inputs, **params)
        except:
            return None  # Not foldable

    def fold_if_constant(self, var):
        var_id = id(var)
        
        if hasattr(var, 'val'):  # It's a Literal packaged into the eqn
            return var.val
        
        # Already computed?
        if var_id in self.known_constants:
            return self.known_constants[var_id]
        
        # Find producer O(1) lookup
        eqn = self.producer_map.get(var_id)
        if not eqn:
            return None  # Input or already constant
        
        # Recursively fold inputs
        input_vals = []
        for inp in eqn.invars:
            val = self.fold_if_constant(inp)
            if val is None:
                return None
            input_vals.append(val)
        
        # Evaluate and cache
        result = self.eval_jax_primitive(eqn.primitive, input_vals, eqn.params)
        self.known_constants[var_id] = result
        return result

[enpasos]

Thanks for digging into this, @rakadam — the need to “peek” into constant index tensors definitely resonates. We can support that without breaking the “core stays generic” rule from docs/design.md by letting plugins opt in to the folding they understand. Here’s a concrete plan you can drop straight into your PR:

1. Add a small constant-folder helper that the IR context owns.
   The core ships it with an empty registry; plugins register the primitives they know how to evaluate.

2. Expose two hooks on IRContext:

def register_constant_evaluator(self, primitive, handler):
    self._const_folder.register_handler(primitive.name, handler)

def try_evaluate_const(self, var):
    return self._const_folder.try_evaluate(var)

register_handler just records primitive.name -> handler.
try_evaluate walks producer equations, checks the registry, caches any success, and returns None as soon as something isn’t foldable.

3. Seeding & producer install.
   Seed the folder with existing constvars inside bind_const_for_var, and call ctx._const_folder.install_producers(jpr) before lowering each jaxpr so it can trace back to the producing equation when needed.

4. Gather plugin registration.
   In the gather plugin (module init or a one-time guard in lower), register the handful of pure, NumPy-friendly primitives you rely on — reshape, broadcast, squeeze, slice helpers, and simple arithmetic. Lower stays pure; no extra jaxpr/literal arguments.

5. During lowering:

const_idx = ctx.try_evaluate_const(indices_var)
if const_idx is not None:
    # emit the efficient Slice/Transpose/GatherND path
else:
    # fall back to the generic lowering

Helper skeleton (e.g., jax2onnx/converter/ir_constants.py):

from __future__ import annotations
from typing import Any, Callable, Dict, Optional
import numpy as np

class ConstantFolder:
    def __init__(self) -> None:
        self._known: Dict[int, np.ndarray] = {}
        self._producer: Dict[int, Any] = {}
        self._handlers: Dict[str, Callable[..., Any]] = {}

    def register_const(self, var: Any, value: np.ndarray) -> None:
        self._known[id(var)] = np.asarray(value)

    def install_producers(self, jaxpr) -> None:
        self._producer.clear()
        for eqn in jaxpr.eqns:
            for out in eqn.outvars:
                self._producer[id(out)] = eqn

    def register_handler(self, primitive_name: str, handler: Callable[..., Any]) -> None:
        self._handlers[primitive_name] = handler

    def try_evaluate(self, var: Any) -> Optional[np.ndarray]:
        vid = id(var)
        if vid in self._known:
            return self._known[vid]

        literal = getattr(var, "val", None)
        if literal is not None:
            arr = np.asarray(literal)
            self._known[vid] = arr
            return arr

        eqn = self._producer.get(vid)
        if eqn is None:
            return None

        handler = self._handlers.get(eqn.primitive.name)
        if handler is None:
            return None

        inputs: list[np.ndarray] = []
        for invar in eqn.invars:
            val = self.try_evaluate(invar)
            if val is None:
                return None
            inputs.append(val)

        try:
            out = handler(*inputs, **eqn.params)
        except Exception:
            return None

        arr = np.asarray(out)
        self._known[vid] = arr
        return arr

Glue in IRContext:

class IRContext:
    def __init__(self, ...):
        ...
        self._const_folder = ConstantFolder()

    def register_constant_evaluator(self, primitive, handler):
        self._const_folder.register_handler(primitive.name, handler)

    def try_evaluate_const(self, var):
        return self._const_folder.try_evaluate(var)

Inside bind_const_for_var:

self._const_folder.register_const(var, array)

Before lowering each jaxpr:

ctx._const_folder.install_producers(jpr)

Example registration in the gather plugin (run once):

from jax import lax

def _ensure_constant_folders_registered(ctx):
    ctx.register_constant_evaluator(lax.broadcast_in_dim_p, lax.broadcast_in_dim)
    ctx.register_constant_evaluator(lax.reshape_p,            lax.reshape)
    ctx.register_constant_evaluator(lax.squeeze_p,            lax.squeeze)
    ctx.register_constant_evaluator(lax.dynamic_slice_p,      lax.dynamic_slice)
    ctx.register_constant_evaluator(lax.slice_in_dim_p,       lax.slice_in_dim)
    ctx.register_constant_evaluator(lax.add_p,                lax.add)
    ctx.register_constant_evaluator(lax.sub_p,                lax.sub)
    ctx.register_constant_evaluator(lax.mul_p,                lax.mul)
    ctx.register_constant_evaluator(lax.min_p,                lax.min)
    ctx.register_constant_evaluator(lax.max_p,                lax.max)

Then, inside lower, call _ensure_constant_folders_registered(ctx) and check const_idx = ctx.try_evaluate_const(indices_var) as shown above.

This keeps the converter core clean, gives gather the constant visibility it needs, and lets other plugins opt in later if they need folding too. If this direction looks good to you, go ahead and wire it into your PR — I’m happy to iterate on the minimal API surface once it’s up. Thanks again for driving this forward!

[enpasos]

Thanks for the heavy lift on gather — the LowerDimExpr and constant-folding approach looks solid. Please open a PR (draft is fine) with your current plugin as-is, including the lower(..., jpr, literals) tweaks. My earlier notes (e.g., a small constant-folding hook on IRContext and moving LowerDimExpr into the builder) were brainstorming, not blockers; I’m happy to adjust the framework during review so your approach lands smoothly. We can iterate on interface shape, expect_graph, and tests together.

[enpasos]

Thanks @rakadam! #105 landed the general gather implementation. It should cover most of what you reported; remaining rough edges might be non-constant start_indices and index clipping (current support is PROMISE_IN_BOUNDS). If things look good on your side, happy to close.