Regression test on matrix element and gradient

[matthewfeickert]

A useful test to help establish use case and template out APIs would be to take some notes from the Matrix Element Regression with Deep Neural Networks paper and try to do a regression on the matrix element and the gradient using JAX. This could serve as a good first test and way to help develop an understanding of useful APIs.