Feedback on API for implicit arrays

[davisyoshida]

Edit (June 6): I've uploaded my transform as Qax in case anyone wants to try it out.

As I mentioned in #16259, I noticed that there were several cases where I wanted something which represented an array without fully instantiating it. This only requires a subset of the full power of custom interpreters, and requires a bunch of repeated tasks each time such as flattening args and passing the transformation to for higher level primitives such as remat.

I ended up with an abstract ImplicitArray type which may be subclassed by a class which implements:

1. A materialization function which can produce a full array if necessary
2. Any number of primitive handlers which compute the output of the target primitive without needing to instantiate the full array. Each handler should be decorated with implicit_op("<primitive_name>"), and can return either JAX types or further ImplicitArray values.

Here's the base class, I've omitted the actual Tracer implementation and some helper functions for brevity:
(You may want to just scroll past and look at the toy example since I think it gets the idea across.)

ImplicitArray interface

class ImplicitArray(ABC):
    op_registry = None
    commute_ops = True

    def __init__(self, shape, dtype):
        self._shape = shape
        self._dtype = dtype

    @property
    def shape(self):
        return self._shape

    @shape.setter
    def shape(self, new_shape):
        self._shape = new_shape

    @property
    def dtype(self):
        return self._dtype

    @dtype.setter
    def dtype(self, new_dtype):
        self._dtype = new_dtype

    @property
    def aval(self):
        return core.ShapedArray(self.shape, self.dtype)

    @classmethod
    def default_handler(cls, primitive, *args, params=None):
        if params is None:
            params = {}
        return materialize_handler(primitive, *args, params=params)

    def _materialize(self):
        #return type(self).materialize(self)
        wrapped = lu.wrap_init(type(self).materialize)
        flat, in_tree = _flatten_one_implicit_layer((self,))
        flat_fn, out_tree = flatten_fun_nokwargs(wrapped, in_tree)
        out_flat = use_implicit_args(flat_fn.call_wrapped)(*flat)
        result = jax.tree_util.tree_unflatten(out_tree(), out_flat)
        return result

    @abstractmethod
    def materialize(self):
        pass

    def flatten(self):
        """
        To control tree flattening, override this method rather than `tree_flatten`.
        The auxiliary data returned will have the shape and dtype added to it.
        """
        return [], ()

    def unflatten(self, aux_data, children):
        """
        To control unflattening override this method rather than `tree_unflatten`.
        This will avoid __init__ being called during tree unflattening.
        """
        pass

    def tree_flatten_with_keys(self):
        children, aux_data = self.flatten()
        aux_data = (self.shape, self.dtype, aux_data)
        return children, aux_data

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        obj = object.__new__(cls)
        shape, dtype, aux_data = aux_data
        obj.shape = shape
        obj.dtype = dtype
        obj.unflatten(aux_data, children)
        return obj

    def handle_primitive(self, primitive, *args, params):
        if primitive.name not in self.op_registry:
            return None

        handler = lu.wrap_init(partial(self.op_registry[primitive.name], primitive))
        use_params = params
        maybe_kwargs = {'params': params} if params else {}

        if len(args) == 2 and self.commute_ops:
            args, use_params = _maybe_swap_args(primitive.name, args, use_params)

        flat_args, in_tree = _flatten_one_implicit_layer((args, maybe_kwargs))
        flat_handler, out_tree = flatten_fun(handler, in_tree)

        result = use_implicit_args(flat_handler.call_wrapped)(*flat_args)
        return jax.tree_util.tree_unflatten(out_tree(), result)

    def __init_subclass__(cls, commute_ops=True, **kwargs):
        super().__init_subclass__(**kwargs)
        cls.op_registry = {}

        if 'tree_unflatten' in cls.__dict__:
            raise TypeError(
                'In order to define an ImplicitArray\'s unflattening behavior,'
                ' define the `unflatten` method instead of `tree_unflatten`.'
            )

        if 'tree_flatten_with_keys' in cls.__dict__:
            raise TypeError(
                'In order to define an ImplicitArray\'s flattening behavior,'
                ' define the `flatten` method instead of `tree_flatten_with_keys`.'
            )


        register_pytree_with_keys_class(cls)
        for imp_op in cls.__dict__.values():
            if not isinstance(imp_op, _ImplicitOp):
                continue
            for name in imp_op.names:
                cls.op_registry[name] = partial(imp_op, cls)

def _materialize_all(it):
    return [val._materialize() if isinstance(val, ImplicitArray) else val for val in it]

class _ImplicitOp:
    def __init__(self, names, op_fn):
        self.names = names
        self.op_fn = op_fn

    def __call__(self, cls, *args, **kwargs):
        if isinstance(self.op_fn, classmethod):
            return self.op_fn.__func__(cls, *args, **kwargs)
        print(f'Calling with {args} and {kwargs}')
        return self.op_fn(*args, **kwargs)

def implicit_op(op_names):
    if isinstance(op_names, str):
        op_names = [op_names]
    def decorator(op_fn):
        return _ImplicitOp(op_names, op_fn)
    return decorator

Here's a toy example of using it to represent a symbolic zero (I think you can probably get XLA to accomplish the same thing with constant folding, so it's just for demonstration purposes):

Symbolic zero example

import jax
import jax.numpy as jnp

from implicit_arrays.implicit_array import ImplicitArray, implicit_op, use_implicit_args

class ImplicitZeros(ImplicitArray):
    def materialize(self):
        return jnp.zeros(self.shape, dtype=self.dtype)

    @implicit_op('mul')
    def mul(primitive, *args):
        result = jax.eval_shape(
            (lambda x, y: x * y),
            *(x.aval for x in args)
        )
        return ImplicitZeros(result.shape, result.dtype)

def main():
    @use_implicit_args
    def f(x, y):
        return jnp.sum(x * y)

    zeros = ImplicitZeros((2, 3), dtype=jnp.float32)

    result = jax.jit(f)(zeros, jnp.ones(zeros.shape))
    # UserWarning: Primitive reduce_sum was not handled by class ImplicitZeros, so implicit args will be materialized

    assert result == 0

if __name__ == '__main__':
    main()

Since I've implemented mul only, the output of x * y will be another ImplicitZeros, but then it gets materialized into a full matrix when the unknown reduce_sum op is hit.

The things I wasn't sure about were:

1. Is it a bad idea to modify jax.core.pytyp_aval_mappings?
2. Is it undesirable to register subclasses as pytrees myself rather than letting the subclasser do it?
3. I went with passing all the primitive arguments to the handlers in order, rather than separately having handlers for left and right application of binary ops (as with __add__ and __radd__ in python). Maybe the separate handlers setup would be preferable?
4. One thing that isn't clear from the above example is that an ImplicitArray can look at the params for a primitive and refuse to handle it by returning None. For example in my LoRA implementation I only allow gathers which look like embedding lookups, otherwise I materialize the array. Should the "decide whether to handle this op" be separate from the handling itself? Also, is None ever a valid return value for an op? If so I'll need to switch to some other null result, or maybe exceptional control flow.

I'm interested in any other feedback as well.

P.S. This should be compatible with Equinox, just in case anyone pops through to ask about that ;)

Edit: Here's another example, based on the request from #16289. It's not feature complete, but it demonstrates the idea:

Masked array example

import jax
import jax.numpy as jnp

from implicit_arrays import ImplicitArray, implicit_op, use_implicit_args
from implicit_arrays.constants import ELEMENTWISE_BINOPS, ELEMENTWISE_UNOPS, REDUCTION_OPS

class NullableArray(ImplicitArray):
    def __init__(self, val, mask):
        super().__init__(shape=val.shape, dtype=val.dtype)
        self.val = val
        self.mask = mask

    def flatten(self):
        return [('val', self.val), ('mask', self.mask)], ()

    def unflatten(self, aux_data, children):
        self.val, self.mask = children

    def materialize(self):
        return self.val

    @implicit_op(ELEMENTWISE_UNOPS)
    @classmethod
    def handle_unop(cls, primitive, nullable_val, params=None):
        val = cls.default_handler(primitive, nullable_val.val, params=params)
        return NullableArray(val, nullable_val.mask)

    @implicit_op(ELEMENTWISE_BINOPS)
    @classmethod
    def handle_binop(cls, primitive, lhs, rhs, params=None):
        lhs_is_nullable = isinstance(lhs, NullableArray)
        rhs_is_nullable = isinstance(rhs, NullableArray)
        mask = lhs.mask if lhs_is_nullable else None

        if lhs_is_nullable:
            lhs = lhs.val

        if rhs_is_nullable:
            mask = rhs.mask if mask is None else mask & rhs.mask
            rhs = rhs.val

        out_val = cls.default_handler(primitive, lhs, rhs, params=params)
        return NullableArray(out_val, mask)

    @implicit_op(REDUCTION_OPS)
    @classmethod
    def handle_reduction(cls, primitive, val, params=None):
        new_val = cls.default_handler(primitive, val, params=params)
        new_mask = cls.default_handler(jax.lax.reduce_and_p, val.mask, params=params)
        return NullableArray(new_val, new_mask)

@jax.jit
@use_implicit_args
def f(x, y):
    return jnp.sum(-x * y, axis=0)

if __name__ == '__main__':
    x = NullableArray(
        val=jnp.ones((2, 3)),
        mask=jnp.asarray(
            [[True, False, True],
             [False, True, True]]
        )
    )

    y = NullableArray(
        val=jnp.full((2, 3), 3),
        mask=jnp.asarray(
            [[False, True, True],
             [True, True, True]]
        )
    )

    output = f(x, y)
    print(f'Result: {output.val}')
    print(f'Mask: {output.mask}')

    # Result: [-6. -6. -6.]
    # Mask: [False False  True]

You can also nest ImplictArrays inside each other:

Masked symbolic zero

@jax.jit
@use_implicit_args
def f(x, y):
    return (x * y) + 3

shape = (2, 3)
x = NullableArray(
    val=ImplicitZeros(shape, jnp.float32),
    mask=jnp.asarray(
        [[True, False, True],
         [False, True, True]],
    )
)

y = NullableArray(
    val=jnp.ones(shape),
    mask=jnp.asarray(
        [[True, True, True],
         [False, False, True]]
    )
)

result = f(x, y)
print(f'Result:\n{result.val}')
print(f'Mask:\n{result.mask}')

# Result:
# [[3. 3. 3.]
#  [3. 3. 3.]]
# Mask:
# [[ True False  True]
#  [False False  True]]

Edit June 13: I used Qax to implement 4-bit absmax quantization here. It only has a materialize rule since my adventures with triton kernels have mostly been unsuccessful so far. This isn't really new since my GPTQ code works the same way, but it was really easy to rewrite it this way.

[davisyoshida]

@chaokunyang re: #16289, would this be applicable to your use case? I don't know if there's enough demand for nullability for it to become a first class JAX feature but I think you could implement it in "user space" this way.

Edit: I took a crack at prototyping it and it seems to work alright!

[davisyoshida]

PRNGKeys are array-likes with special properties

Can you elaborate on how this would help with PRNGKey usage? Mostly it just seems like you could get some better error messages if you passed one in the wrong argument.

BTW, can you explain what you see as the main advantages of the multiple dispatch approach? I do like the ergonomics of separately defining f(A, B) and f(A, B), but it seems like there are some benefits that you could get from inheritance as well.

[davisyoshida]

The only downside seems to be that we'd need to insert another transformation between our array-like objects and the code they're being run on: adding another decorator means the end-user ergonomics aren't as clean.

I think there are some good reasons to want this to be a transformation instead of default behavior. One that came up in testing is that you can have a cond where both branches emit the same shape, but if one branch leads to an ImplicitArray and the other doesn't it's no longer valid. Also I think the ordering of "implicit-ing" and "grad/vmap/etc-ing" is important. (Although currently I can't come up with a case where you wouldn't want implicit-ify to be the inner transformation).

[patrick-kidger]

Can you elaborate on how this would help with PRNGKey usage? Mostly it just seems like you could get some better error messages if you passed one in the wrong argument.

Right now, JAX PRNG keys are arrays of shape (2,) and dtype uint32. This is kind of weird, in that it allows you to write stuff like key + 1, despite the fact that addition isn't really defined for keys -- that is to say, we'd like to type them more strongly. As such, things are about to change: PRNG keys will become magical sort-of-arrays with shape () and dtype key. (Which is an "opaque dtype".) They are allowed to bind to some special ops (jax.random.normal) but have default behaviour when used with other ops (e.g. raise an error when used with jax.lax.add_p).

And this is exactly the same scenario we find ourselves in here! We want our sort-of-arrays to have special behaviour when interacting with certain primitives (implicit_op("mul")) and default behaviour when used with other ops (materialize -- which could raise an error if we wanted).

So far the JAX team have been implementing PRNGKeys by introducing a whole slew of new systems, like the aforementioned opaque dtypes and such. But I think that it could actually be implemented as a special case of what we're discussing here! (Although possibly it would require this extra decorator, so it wouldn't be backward compatible -- good reasons that the JAX team didn't do it this way in their case.)

---

Multiple dispatch: this is important even in trivial examples. I want both mul_p.bind(Zero(...), x) and mul_p.bind(x, Zero(...),) to return a Zero to me. I can't do that with single dispatch -- the latter example will fail as the dispatch hits the x argument. (Which might be one of our sort-of-arrays itself, of a different type!)

---

Being a transformation: I think having a cond which emits two arrays with the shape, but different dtypes, should indeed be an error! I wouldn't want that to silently coerce to one type or the other. So I wouldn't mind if multiple dispatch was default behaviour and it just raised an error here.

Ordering of transformations: I think I would usually want the order to be jit-implicit-{everything else}, for which we would have implicit as an outer transformation. Indeed I don't think we necessarily want our sort-of-array to be a pytree except to have it pass the JIT boundary cleanly.

[davisyoshida]

RNG keys will become magical sort-of-arrays

I hadn't realized this, makes sense

I want both mul_p.bind(Zero(...), x) and mul_p.bind(x, Zero(...),) to return a Zero to me

So currently the way I'm handling this is that the handler of whichever argument is an ImplicitArray is invoked. This is no problem for Zero * Array, Array * Zero, or Zero * Zero. The only question is what should happen if there's another ImplicitArray, perhaps Infinity. Should Infinity * Zero be handled by the Zero handler or the Infinity handler?

Actually now I'm sorta talking myself into the multiple dispatch thing since I assume the answer is that a user who wants to use these together should register a mul(Infinity, Zero) handler...

a cond which emits two arrays with the shape, but different dtypes, should indeed be an error!

Well it's pretty annoying if the goal is to support running these objects through functions which had no idea they'd be used with anything other than vanilla arrays. The problem isn't different dtypes per-se, but more that one branch might return an explicit MxN float32 array, and the other might end up giving an implicit one. Thinking about it more though, I think this should be resolved by materializing stuff until the branches are in agreement again (as opposed to being an error or materializing all the cond arguments instead).

I don't think we necessarily want our sort-of-array to be a pytree except to have it pass the JIT boundary cleanly

I think that them being pytrees and using implicit as the inner transformation is really nice for grad in particular. This way anything you've implemented in pure JAX in your handlers gets differentiated through and the actual JAX arrays that make up your ImplicitArray receive gradient information. You can also use all the available optimizers and so on because these subarrays are part of a pytree.

The alternative of implicit-grad requires (I think) quite a bit more manipulation to get the behavior you want. In general, the output of implicit(grad(f))(MyImplicitType) will not be a MyImplicitType instance. So then you have to implement some extra machinery to compute the grad of the constituent subarrays. Admittedly I'm not well-versed enough in JAX's internals to know how easy or hard this is to get working. I think everything being pytrees and user "userland" JAX to accomplish things makes it more likely that it will be compatible with other user-made transformations as well, not just JIT and grad.

[patrick-kidger]

This way anything you've implemented in pure JAX in your handlers gets differentiated through

In general, the output of implicit(grad(f))(MyImplicitType) will not be a MyImplicitType instance.

I think everything being pytrees and user "userland" JAX to accomplish things makes it more likely that it will be compatible with other user-made transformations as well, not just JIT and grad.

All good points, agreed. PyTrees it is!

[patrick-kidger]

Is it a bad idea to modify jax.core.pytyp_aval_mappings?

I can see this must be part of the code you've omitted. Why is this necessary? As per #16259 we can probably do this without new avals.

Is it undesirable to register subclasses as pytrees myself rather than letting the subclasser do it?

I think this is definitely desirable! If I were doing this myself I'd using eqx.Module as a base class for ImplicitArray, which does exactly this. ;)

I went with passing all the primitive arguments to the handlers in order, rather than separately having handlers for left and right application of binary ops (as with add and radd in python). Maybe the separate handlers setup would be preferable?

Actually, my main criticism of the above approach is that I think we should go full multiple dispatch. For example this means not defining ImplicitZeros.mul, it means something like

@register_multiple_dispatch(jax.lax.mul_p)
def _(x: ImplicitZeros, y):
    return x

@register_multiple_dispatch(jax.lax.mul_p)
def _(x, y: ImplicitZeros):
    return y

(side note, I'm also using primitives themselves, not strings, to refer to the operaiton to be overloaded).
(side side note, maybe use plum for implementing multiple dispatch?)

Should the "decide whether to handle this op" be separate from the handling itself?

I'd say probably not. I don't think this is needed.

Also, is None ever a valid return value for an op? If so I'll need to switch to some other null result, or maybe exceptional control flow.

return NotImplemented perhaps?

P.S. This should be compatible with Equinox, just in case anyone pops through to ask about that ;)

Hehehe.

[davisyoshida]

Why is this necessary?

I've just realized that it isn't since Tracers are supposed to have a custom aval method which I had failed to implement. Whoops!

I think we should go full multiple dispatch

I don't have a super strong opinion here. I don't think it affects the implementation much to be honest, since either the trace code checks with the instance or it can go look it up in whatever registry. One nice thing about having the implementation live on the class is that it makes it easy to partially defer to the superclass's implementation. I'm about to update the main post with another example which takes advantage of this. I've also made a bunch of assorted changes which I think improved usability/conciseness in the examples I've tested with so far.

I'm also using primitives themselves, not strings, to refer to the operation to be overloaded

This is probably better since you could technically run into ambiguity, but I do like the ease of use you get from strings.