You can actually capture what you want implicitly.
Unfortunately  the checkpoint was built to work with much simpler function, usually the forward pass from a nn.Module.
You can go around it by doing something like this here:
import torch
from torch.utils import checkpoint

def fun(theta, t):
    return theta['alpha']**t + theta['beta']**t

def wrapper(x):
    x = torch.tensor(x, requires_grad=True)
    theta = {}
    theta['alpha'] = x[0]
    theta['beta'] = x[1]
    out = torch.empty(3)
    for t in range(3):
        out[t] = checkpoint.checkpoint(lambda x: fun(theta, t), x)
    y = out.sum()
    y.backward()
    grad = x.grad
    return (y.item(), grad)

print(wrapper([2., 3.]))