Product of grad operator

[exenGT]

Hi,

Recently I have encountered a scenario where I need to process expressions like this:

[(a0 + a1*(d/dx) + a2*(d^2/dx^2)) * (b1*(d/dx) + b3*(d^3/dx^3))] f(x)

which, when expanded, involves products of the grad operator (for example, one of the terms is a2*(d^2/dx^2) * b3*(d^3/dx^3) = a2*b3*(d^5/dx^5).). In this situation, I imagine that an implementation of grad with differentiation order as an argument would be very helpful. Specifically, it should have the property

[grad(m1, f) * grad(m2, f)](x) = grad(m1+m2, f)(x)

where m1 and m2 are orders of differentiation.

I would like to know if something like this can be implemented? If so, how?

[jakevdp]

You could write it this way if you like:

def grad_n(n: int, f, **kwargs):
  for i in range(n):
    f = grad(f, **kwargs)
  return f

I think we'd be unlikely to add such a helper function to the JAX API, since it's straightforward to define a function like this if you need it.