As the creation of the learned index tensor is a multi-stage process (involving the learning of a float tensor which is then transformed into a tensor of intigers), do you think there is any way to bypass the intiger-index tensor for calculating the gradient from inside the gradient tree? Perhaps a subunit of somekind?
Something like the following, where the green arrows show the gradient flow and the black arrows show the forward pass…

image1351×540 52.5 KB