Modify EnsembleAverageOp or create a new method that also returns Jacobian #145

@python_app
def ensemble_average_jacobian(
    system: fitlib.TensorSystem,
    frames_path: pathlib.Path,
    temperature: float,
    pressure: float | None,
    force_field: fitlib.TensorForceField,
) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
    import openmm.unit
    import torch

    import fitlib.mm
    from fitlib.mm._ops import _pack_force_field, _unpack_force_field

    # Use existing fitlib packing order, including attributes and optional v-sites.
    tensors, parameter_lookup, attribute_lookup, has_v_sites = _pack_force_field(
        force_field
    )
    parameters = torch.cat(
        [t.detach().reshape(-1) for t in tensors if t is not None]
    ).requires_grad_(True)
    pieces = iter(parameters.split([t.numel() for t in tensors if t is not None]))
    tensors = tuple(
        None if t is None else next(pieces).reshape(t.shape) for t in tensors
    )
    worker_ff = _unpack_force_field(
        tensors,
        parameter_lookup,
        attribute_lookup,
        has_v_sites,
        force_field,
    )
    means, _ = fitlib.mm.compute_ensemble_averages(
        system,
        worker_ff,
        frames_path,
        temperature * openmm.unit.kelvin,
        None if pressure is None else pressure * openmm.unit.atmosphere,
    )
    # Jacobian rows follow sorted observable names; columns follow packed parameters.
    jacobian = torch.stack(
        [
            torch.autograd.grad(means[name], parameters, retain_graph=True)[0]
            for name in sorted(means)
        ]
    )
    return {name: value.detach() for name, value in means.items()}, jacobian.detach()