def METHOD_NAME(self) -> Optional[PyTree]:
    r"""This methods should return a PyTree that will be passed as the `data` argument
    to the `_expect_kernel`. The PyTree should be composed of jax arrays or hashable
    objects.
    For example for the kinetic energy this method would return the masses of the
    individual particles."""