def METHOD_NAME(equation, *operands, **kwargs):
    """
    Runs einsum to compute the partition function.
    This is cheap but less numerically stable than using the torch_log backend.
    """
    key = "prob", equation, kwargs["plates"]
    if key not in _CACHE:
        # This simply wraps einsum for jit compilation.
        def _einsum(*operands):
            return einsum(equation, *operands, **kwargs)
        _CACHE[key] = torch.jit.trace(_einsum, operands, check_trace=False)
    return _CACHE[key](*operands)