Hi Chris!
Luckily you don’t need to fork nor wait for the team to change that for you. You will however have the responsibility of writing a function to compute grad_sample for your layer. Once you write it, you simply register it for your layer using @register_grad_sampler 8 and you should be good to go!