Primitive that makes XLA custom call with parameters

[ebilionis]

Hi Jax community,

I am writing a primitive that makes an XLA custom call (to a C++ function).
I know how to do it when my primitive does not have any keywords (parameters).
My problem is how to make it work when the primitive has parameters.

Here are some details:

• My primitive has the form y = foo(x, *, p).
• The input x is a regular array. I want to be able to differentiate with respect to it. etc.
• The input p is a parameter. In my case, it is a an array of integers.
• The shape of y is the same for all x and p.
• My C++ function needs to know the shape of x, say it is 1D array of size N, a pointer to the data in x and the same for p (say also 1D, integers, of size D).
• My xla translation function looks like this:

from jaxlib import xla_client

def foo_xla_translation(c, x, *, p):
   # sanity checks
   # definition of shapes, N and D
   return xla_client.ops.CustomCallWithLayout(
              c,
              b"my_cpp_func",
              operands=(xla_client.ops.ConstantLiteral(N), x, xla_client.ops.ConstantLiteral(D), p),
              shape_with_layout=out_shape,
              operand_shapes_with_layout=xla_shape_of_N, xla_shape_of_x, xla_shape_of_D, xla_shape_of_p)
    )

This does not work because p is not an XlaOp. It is a Traced<ShapedArray(int32[1000])>with<DynamicJaxprTrace(level=0/1)>.

Note that everything works fine if I treat p just like x (as a regular input, not as a parameter).

I'm okay with a solution that treats p just like x for the XLA call if I can exclude p from JVP/VJP.

Regards,

ib

[froystig]

I'm okay with a solution that treats p just like x for the XLA call if I can exclude p from JVP/VJP.

I think this will be the way to go. In jax internals, primitive "parameters" supply static information, whereas "operands" are possibly-dynamic. The fact that p is a Traced<...>with<...> suggests that it's not actually static info, and that what you're passing for p at primitive binding time is a value derived from dynamic inputs/intermediates of the computation being staged out for compilation.

To exclude p from autodiff, have you tried accepting it as an operand instead, just like x, but without defining a JVP rule for it?

By the way, we have a proposal for adding a proper FFI API in #12632, so that hopefully you won't need to dig into jax internals/primitives in the future in order to set up custom calls. Feedback there is welcome. (cc @sharadmv)

[ebilionis]

Thanks for the response @froystig.

It would really be great to have the FFI API! I have spent many days writing primitives with a custom call... and every little variation I have to dig deeper into jaxlib, xla, etc.

As I wait for the FFI API, here is a minimal example of what my problem is.
Maybe it is an obvious mistake of how to make primitives with parameters.

First, I know that this is possible to do what I want, because I see done extensively here:
https://github.com/google/jax/blob/main/jax/_src/lax/linalg.py

It is just that at the end there is always a call to an xla op that is opaque to me.
For example, see the QR

def _qr_translation_rule(ctx, avals_in, avals_out, operand, *, full_matrices):
  operand_aval, = avals_in
  shape = operand_aval.shape
  m, n = shape[-2:]
  if m == 0 or n == 0:
    return [_eye_like_xla(ctx.builder, avals_out[0]),
            _zeros_like_xla(ctx.builder, avals_out[1])]
  return xops.QR(operand, full_matrices)

The full_matrices is a bool here. My case is like that.
But I have an int instead of a bool.
(Actually an array of ints, but I would be happy if I first get an int).

Here is a simple code demonstrating what happens.

import jax
import jax.numpy as jnp
from jaxlib import xla_client
from jax.interpreters import xla

def foo(x, p=2):
    return foo_p.bind(x, p=p)

foo_p = jax.core.Primitive("foo")

@foo_p.def_impl
def foo_p_impl(x, *, p):
    return x ** p

@foo_p.def_abstract_eval
def foo_p_eval(x, **unused_kwg):
    return jax.abstract_arrays.ShapedArray(x.shape, x.dtype)

def foo_p_xla(ct, avals_in, avals_out, xc, *, p):
    print(p)
    # The following does not work because p is a
    # Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)
    return [xla_client.ops.Pow(xc, p)]

xla.register_translation(foo_p, foo_p_xla)

if __name__ == "__main__":
    x = 2.0 * jnp.ones((10,))
    y = foo(x, p=5)

    # This does not work
    jfoo = jax.jit(foo)

    y = jfoo(x, p=5)

[mattjj]

Also cc-ing @zhangqiaorjc @cky9301

[froystig]

So could you pass and accept p as an operand to the primitive instead of as a parameter? For autodiff, you can use ad.defjvp but give it only one rule. That would correspond to dropping the second input tangent's contribution to the output tangent. In other words, this would ignore any perturbations in p, treating them as zero.

(The usual disclaimer applies: these are internals and may change without notice. But this is how they behave today.)