Define function of grad operator

[exenGT]

Hi all,

Is it possible to define your own grad operator that satisfies certain user-defined rules (e.g. chain rule, product rule, etc.)?

Thanks,
JW

[jakevdp]

Could you say more about what you have in mind? The build-in jax.grad does satisfy the chain rule and product rule.

Or perhaps you're asking whether you can define your own derivative rule for a function? If so, this might be helpful: https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html

[exenGT]

Hi Jake,

Thank you for your reply! I think the @custom_jvp decorator should be able to do what I want, except for the fact that the modification should be imposed only on the grad operator but not on the original function. An example I have in mind is the "spherical tensor gradient operator" (https://arxiv.org/pdf/math-ph/0505018.pdf), which defines a particular type of derivative for any ordinarily-differentiable function. Is there a way to do this without redefining every possible function which would undergo this new derivative operation?

Thanks,
JW

[exenGT]

An alternative way to look at this is that the new derivative can be written as a function (in this case the regular solid harmonics, which is a polynomial) of the grad operator. Is there a way to express this in JAX?

[exenGT]

To make it clearer, suppose I would like to define an operator new_grad whose action on any function f(x, y) is defined by new_grad(f)(x0, y0) = (dx)^2 f(x0, y0) - (dy)^2 f(x0, y0), where (dx)^2 f(x0, y0) denotes second-order partial derivative of f with respect to x evaluated at x = x0, y = y0 (similarly for (dy)^2). I am trying to think of a way to implement this function for generic f, but without success. Any suggestion?

[jakevdp]

Is something like this what you have in mind?

def new_grad(f):
  def f_grad(x, y):
    return jax.hessian(f, argnums=0)(x, y) + jax.hessian(f, argnums=1)(x, y)
  return f_grad

f = lambda x, y: x * y
result = new_grad(f)(1.0, 2.0)

new_grad(lambda x, y: x ** y)(1.0, 2.0)

[exenGT]

I think this works! Thank you very much :)