Defining new JAX primitives: looking for help understanding vjp rule arg shapes

[oaniewiarowski]

Hello!
Similar to this post, I wish to wrap an external library function with jax primitives. The code provided in the linked solution to the post mostly makes sense, but it uses a scalar function as an example while my function is multidimensional. I am having trouble understanding the shapes that my arrays need to take when defining the vjp rule.

Here is a working pure jax minimal implementation of my function (for comparison purposes).

import jax
import jax.numpy as jnp

def jnp_lagrange(points):
  # calculate linear lagrange finite element basis on a unit triangle 
  # return array dimensions: 
  # (# of query points (1), # of degrees of freedom (3), value size of element (1))
  x = points[:,0]
  y = points[:,1]
  return jnp.array([[-x -y + 1, x, y]]).transpose(2,1,0)  # get shape to match the c++ func output

# a trivial scalar functional
def test_func(x): return jnp.take(jnp_lagrange(x), 0)

jax_pts = jnp.array([[0.1, 0.05]])
val, grad = jax.value_and_grad(test_func)(jax_pts)

print("basis functions: ", jnp_lagrange(jax_pts))
print("test_func value: ", val)
print("basis shape: ", jnp_lagrange(jax_pts).shape)
print("test_func grad: ", grad)  # this works as expected

basis functions:  [[[0.85]
  [0.1 ]
  [0.05]]]
test_func value:  0.85
basis shape:  (1, 3, 1)
test_func grad:  [[-1. -1.]]

Below, I import a wrapped C++ function that does the same thing from the basix library and follow the code from the previous post, but I run into trouble defining the vjp rule. I find the shapes of the cotangent and my jacobian don't match. I would very much appreciate it if someone could explain the shape of the cotangent (e.g where do the outer dimensions come from in jax) and how to write the transpose correctly!)

Thank you,
Olek

! pip install fenics_basix
import basix
from basix import CellType, ElementFamily, LagrangeVariant
import numpy as onp
import jax
import jax.numpy as jnp
from jax.interpreters import ad

######### code as from previous post, with new foo_p and foo_vjp_impl functions
foo_p = jax.core.Primitive('foo')
foo_jvp_p = jax.core.Primitive('foo_jvp')  # jvp rule primititive
foo_vjp_p = jax.core.Primitive('foo_vjp')  # vjp ''

def foo(x):
  return foo_p.bind(x)

@foo_p.def_impl
def foo_impl(x_arr):
  lagrange = basix.create_element(
    ElementFamily.P, CellType.triangle, 1, LagrangeVariant.equispaced)
    # the first arg is the number of derivatives requested (0 = vals only)
  return lagrange.tabulate(0, x_arr)

def foo_jvp(primals, tangents):
  (x,), (xdot,) = primals, tangents
  y = foo(x)
  y_dot = foo_jvp_p.bind(x, xdot)
  return y, y_dot
  
ad.primitive_jvps[foo_p] = foo_jvp

@foo_jvp_p.def_abstract_eval
def foo_jvp_abstract_eval(x_aval, x_dot_aval):
  y_dot_aval = jax.core.ShapedArray(x_dot_aval.shape, x_dot_aval.dtype)
  return y_dot_aval

def foo_jvp_transpose(y_bar, x, x_dot_dummy):
  assert ad.is_undefined_primal(x_dot_dummy)  # just a dummy input
  x_bar = foo_vjp_p.bind(x, y_bar)  # y_bar aka y_grad
  return None, x_bar  # None for nonlinear primal input x

ad.primitive_transposes[foo_jvp_p] = foo_jvp_transpose

@foo_vjp_p.def_impl
def foo_vjp_impl(x, y_bar):
  lagrange = basix.create_element(
    ElementFamily.P, CellType.triangle, 1, LagrangeVariant.equispaced)
  vals_and_grads = lagrange.tabulate(1, x)  # get vals and first deriv
  print("y_bar shape: ", y_bar.shape)
  print("y_bar: ", y_bar)
  # the first element corresponds to the basis func vals, followed by the derivs
  basis_grads = vals_and_grads[1:]
  print("basis_grads shape: ", basis_grads.shape)
  print("basis_grads: ", basis_grads)
  return y_bar * basis_grads

##########

onp_pts = onp.array([[0.1, 0.05]])
def test_func_onp(x): return jnp.take(foo(x), 0)

grad = jax.grad(test_func_onp)(onp_pts)
print("incorrect grad shape: ", grad.shape)
print("incorrect grad: ", grad)

y_bar shape:  (1, 1, 3, 1)
y_bar:  [[[[1.]
   [0.]
   [0.]]]]
basis_grads shape:  (2, 1, 3, 1)
basis_grads:  [[[[-1.00000000e+00]
   [ 1.00000000e+00]
   [-4.30845614e-17]]]


 [[[-1.00000000e+00]
   [ 0.00000000e+00]
   [ 1.00000000e+00]]]]
incorrect grad shape:  (2, 1, 3, 1)
incorrect grad:  [[[[-1.]
   [ 0.]
   [-0.]]]


 [[[-1.]
   [ 0.]
   [ 0.]]]]

[jakevdp]

Hi - it looks like the array shapes here are of the expected shape: you define f_jvp: (1, 2)->(1, 1, 3, 1), and so the shape of x is (1, 2) and the shape of y_bar is (1, 1, 3, 1). I must admit I'm not clear on what your lagrange tabulation is meant to do and why you've transposed and reshaped the results the way you have, so I don't think I can help with the details of why the outputs of those functions are of an unexpected shape, but I can confirm that the inputs are of the expected shape.

Backing up, though, I think the solution in the post you link to is probably not the right template for you to follow here, because it uses non-public APIs, and because Primitives are a pretty heavy solution for the problem you have. You might have more luck instead using an approach based on pure_callback; you can find an example at https://jax.readthedocs.io/en/latest/notebooks/external_callbacks.html#example-pure-callback-with-custom-jvp

[oaniewiarowski]

(I edited my original post because I left those transpose and reshape operations in foo_vjp_impl by mistake trying to wrangle the jacobian shape into something that 'works.')

@jakevdp Thank you for pointing out the pure callback option, but perhaps my question is more fundamental: how to write the vjp itself. In my understanding, the vjp is the transpose of a jacobian times a cotangent vector, $J^T \bar{v}$. However, how does this transposition work for arbitrarily shaped arrays? In this particular case, the jacobian basis_grad is (2, n, 3, 1) and y_bar is (1, n, 3, 1), where n is the number of points . It is not obvious to me how to transpose the jacobian such that the product (the gradient of test_func_onp) will have shape (n, 2). Maybe I misunderstood the purpose of foo_jvp_transpose and there is something missing there?

Even if I take the pure_callback approach, I will still need to implement the vjp correctly, so I really would appreciate any insight on this fundamental matter. Thank you again