Higher order Jacobians and interaction with grad

[Lun4m]

Hi, I have two questions and I'm a bit confused about autodiff.
Is there an easy way to calculate Jacobian tensors of order 3 or higher?
Take for example this function that has as input an array of positions (Nx3 matrix) and outputs a 3D mesh (MxMxM tensor).

import jax.numpy as jnp
import jax.lax as lax
from jax import jacfwd, random, value_and_grad

def paint(positions, mesh_size, box, mass=1):
    # Paint particle positions to 3D mesh
    # Taken from https://github.com/DifferentiableUniverseInitiative/JaxPM
    scale = mesh_size / box
    positions = scale * positions

    positions = jnp.expand_dims(positions, 1)
    floor = jnp.floor(positions)

    connections = jnp.array(
        [[[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1],
          [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1]]]
    )

    neighboor_coords = floor + connections
    kernel = 1.0 - jnp.abs(positions - neighboor_coords)
    kernel = mass * kernel[..., 0] * kernel[..., 1] * kernel[..., 2]
    neighboor_coords = jnp.mod(neighboor_coords, mesh_size)

    dnums = lax.ScatterDimensionNumbers(
        update_window_dims=(),
        inserted_window_dims=(0, 1, 2),
        scatter_dims_to_operand_dims=(0, 1, 2),
    )

    mesh = lax.scatter_add(
        jnp.zeros(mesh_size),
        neighboor_coords.reshape([-1, 8, 3]).astype("int32"),
        kernel.reshape([-1, 8]),
        dnums,
    )
    return mesh

If I write the following, the process gets killed because it eats up all the memory.

seed = 874
key = random.PRNGKey(seed)
n_particles = 5000
particles = random.uniform(key, minval=0, maxval=9.9, shape=(5000, 3))

box = jnp.array([10, 10, 10])
mesh_size = jnp.array([24, 24, 24])
fun = lambda x: paint(x, mesh_size, box)

J = jacfwd(fun)(particles)

But if I wrap the paint function around another function that returns a scalar I can calculate the gradient without problems.
So we get to the second question, what's really happening here? Why does this "work" while the previous code block hangs? Isn't jax chain ruling through the paint function?

def sum_mesh_points(positions, mesh_size, box):
    mesh = paint(positions, mesh_size, box)
    return jnp.sum(mesh)

value, grad = value_and_grad(sum_mesh_points)(particles, mesh_size, box)

[Lun4m]

Okay, I was playing around with make_jaxpr and it seems that jax.lax doesn't like the jnp.zeros(mesh_size) inside scatter_add:

Traceback (most recent call last):
  File "paint.py", line 46, in <module>
    print(make_jaxpr(jacfwd(fun))(particles))
  File "paint.py", line 31, in paint
    jnp.zeros(mesh_size),
  File "~/.local/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 3665, in zeros
    return lax.full(shape, 0, dtype)
TypeError: Shapes must be 1D sequences of concrete values of integer type, got Traced<ShapedArray(int32[3])>with<DynamicJaxprTrace(level=1/0)>.

I added the mesh array as an argument, paint(positions, mesh_size, mesh, box, mass=1), and this seems to solve the issue.
The program still consumes all the memory, but at least it doesn't crash.
Anyway it kind of makes sense that it requires a lot of memory because with the numbers I provided it has to evaluate ~200 million derivatives.

By checking the jaxpr it seems indeed that the higher order jacobians can be computed directly with jacfwd.