Is this pattern already implemented somewhere?

[davisyoshida]

I noticed that in writing both a 4-bit matrix implementation and a LoRA implementation that there's a common pattern of "thing which represents a matrix but avoids instantiating a dense matrix if possible." For example with 4-bit quantization you prefer to use a specialized matmul kernel, but if you need to do something like calculate the norm of a matrix you can always full materialize it.

It's pretty easy to implement this pattern with tracers, but handling some of the stuff like dealing with pjit and remat primitives every time is annoying. Is this special case already implemented somewhere? For example I know there's a __jax_array__ API, so perhaps it's possible to implement it that way instead of needing to reinvent the wheel.

[jakevdp]

No, I don't believe there's any stock pattern for this kind of thing. One note on __jax_array__ is that it's not a fully-implemented API (and probably never will be, see #4725), so I would not recommend relying on it (cc/ @mattjj, who knows more about the plans for this).

[jakevdp]

I think rather than doing things manually, what you want is to call get_bind_params and then bind the original primitive, something like this:

subfuns, bind_params = primitive.get_bind_params(params)
result = primitive.bind(*subfuns, *vars, **bind_params)

[davisyoshida]

That's what I do when I want to execute a primitive with no modification, but in this case I need to at the very least replace the jaxpr which the pjit primitive is executing, and possibly do more things as well if the number of invars is changing (e.g. for a quantized matrix each original input might get transformed to 3 inputs, an int matrix + zero + scale).

Would the following be an okay strategy?

1. Get the bind params
2. Transform the original jaxpr using my transformation and trace its evaluation into a new jaxpr
3. Set the 'jaxpr' field in bind params equal to the new jaxpr
4. Take a guess at how to update 'in_shardings'/'out_shardings' such as duplicating entries in them based on the original input -> transformed inputs mapping.

I'm currently not even doing any sharding, it just happens that pjit primitives get emitted a lot and I want to make sure my code will handle them correctly.

[jakevdp]

Yeah, that sounds reasonable. I'm not sure how to correctly handle shardings in this case – shardings are a new enough concept (and custom final-style transformations are advanced enough) that, without exaggeration, I'd guess there's probably nobody who has experience doing this yet!

[davisyoshida]

Alright I'll give that a shot. Thanks a ton for the quick responses.

[davisyoshida]

"Easy to implement this pattern with tracers" may have been premature 😅 . Just in case anyone comes across this in the future, it seems like the "transform the jaxpr" approach works. The main thing to be careful of is accidentally transforming a jaxpr with no literals into one that has literals. This matters for remat2 since it never passes any literals to its jaxpr. The jax.interpreters.partial_eval.convert_constvars_jaxpr function can rectify this though.

[patrick-kidger]

I have two other cases where this would be a useful feature to have -- sparse arrays, and symbolic zeros. (The latter outside of AD, as a general jaxtype.)

I note that sparse arrays already exist and have some limited support. So would it be possible to build into JAX a single unified approach here?

[davisyoshida]

Is that something that would require modification of JAX internals? I'm not actually clear on what all the legal aval types are.

[patrick-kidger]

I think JAX has most of the notions for this already. We can already write custom primitives and it has avals, etc. The main things are (a) creating custom jaxtypes+their avals is still largely unexplored, and (b) ideally primitives would support multiple dispatch a bit more neatly than having us copy-paste it between their abstract/impl/lowering rules.

[davisyoshida]

Ah yeah it looks like all my tests pass if I make a custom aval inherits from ShapedArray. Cool!

I think to make it fully work properly I need to at least update jax.core.raise_to_shaped_mappings and probably some other places I'm not aware of.

As for the copy-pasting thing, aren't tracers intended for things like this? You can choose to change how a primitive executes, or defer to the default handler.

[patrick-kidger]

So JAX plays pretty fast-and-loose with stuff like inheritance. (And mostly does stuff like if type(x) is Foo rather than isinstance(x, Foo). It might be safer to create an entirely new aval.

Indeed I think there are probably multiple different registries that need to be told how to handle the new aval.

---

On tracers: hmm, maybe that is another option! Don't introduce another aval, but instead introduce a MultipleDispatchTrace with corresponding MultipleDispatchTracer. Aval-wise this is still just a a ShapedArray, but it carries additional metadata that the rule used to dispatch appropriately.

I think I like that! It's basically a generalisation of the existing jax.experimental.sparse.sparsify transform, or of your lora transform. (And using a tracer would make it initial-style rather than final-style, but I don't think that difference is important.)

[davisyoshida]

Have you taken a look at #16280? I'd definitely appreciate people's thoughts on it.