Would this approach of looping through the inputs to calculate the gradients (or using jacobian) work for your use case?
jacobian