Immediate solution
Laplace operator is important for computing the kinetic energy, currently we have to use two nabla operators together with the trace. For example
-
in the transpose rule of nabla,
https://github.com/sail-sg/autofd/blob/main/autofd/operators/operators.py#L773
-
in mgga functional of jax_xc,
https://github.com/sail-sg/jax_xc/blob/main/gen_repo/experimental.jinja#L66
While we can do better with the jet rules.
In the long term
In the long term, I believe we need to have a jaxpr level graph rewriting tool, we can then introduce algebraic relationships that works more generally for any nested nabla operations.
Immediate solution
Laplace operator is important for computing the kinetic energy, currently we have to use two nabla operators together with the trace. For example
in the transpose rule of nabla,
https://github.com/sail-sg/autofd/blob/main/autofd/operators/operators.py#L773
in mgga functional of
jax_xc,https://github.com/sail-sg/jax_xc/blob/main/gen_repo/experimental.jinja#L66
While we can do better with the jet rules.
In the long term
In the long term, I believe we need to have a jaxpr level graph rewriting tool, we can then introduce algebraic relationships that works more generally for any nested nabla operations.