Similar problem, https://arxiv.org/pdf/2311.00212 see section 6. Note: the paper is in the context of machine learning, but they develop the results generically.