Can you confirm that I understood what you want properly:

output is computed as candidate_ops[active_idx](x).

Where does active_idx come from?
The gradient should be computed as the gradients of the selected op (which is differentiable so autograd is good here.


active_idx is just used for indexing but you want to get gradients for it

Since this is not differentiable, you want to use a custom Function for that.
The gradient formula is just sum(output_of_that_op * grad_output) for each score associated with an op (whether or not it was the selected one).