autograd.function with setup_context has a number of issues with torch.compile