jax_md_mod.custom_quantity.total_energy_wrapper# total_energy_wrapper(energy_fn_template)[source]# Wrapper around energy_fn to allow total energy computation via traj_util.quantity_traj.