jagged_to_padded_dense: an autograd kernel was not registered to the Autograd key(s) but we are trying to backprop through it

I'm using TorchRec's JaggedTensor in an unconventional way. I use it to reshape and pad my labels and logits for a custom loss function, but this leads to the following warning:

UserWarning: fbgemm::jagged_to_padded_dense: an autograd kernel was not registered to the Autograd key(s) but we are trying to backprop through it. This may lead to silently incorrect behavior. This behavior is deprecated and will be removed in a future version of PyTorch. If your operator is differentiable, please ensure you have registered an autograd kernel to the correct Autograd key (e.g. DispatchKey::Autograd, DispatchKey::CompositeImplicitAutograd). If your operator is not differentiable, or to squash this warning and use the previous behavior, please register torch::CppFunction::makeFallthrough() to DispatchKey::Autograd. (Triggered internally at /pytorch/torch/csrc/autograd/autograd_not_implemented_fallback.cpp:62.)

The loss function returns a scalar value for a batch, just like typical loss functions. However, I don't know C++ and am not sure how to "register an autograd kernel" -- that sounds difficult and risky. Is there a way I could back-propagate the scalar loss from the logits instead? Is what I'm proposing sensible?

Additional context

I'm training a recommender model that ranks candidate items for users. I'm using TorchRec's DLRM model that is designed for point-wise loss, which means each user-item pair's loss can be calculated independently. However, I want to optimize ranking, and so I'm using ListNet loss. This means I need to calculate loss for all of a user's items jointly. The DLRM forward pass produces logits of shape batch_size, but my loss function needs a shape of users_in_batch, max_candidate_items. Not all users have the same number of candidate items, so I pad up to the largest number of candidate items in the batch. Anyways, padding and reshaping are not differentiable, which makes sense, but in the end I get a scalar loss, so I'm hoping that's sufficient.

Here is what my reshaping code looks like in the forward pass of my version of Meta's DLRMTrain class:

from torchrec.models.dlrm import DLRM
from torchrec.sparse.jagged_tensor import JaggedTensor

...

class DLRMTrain(nn.Module):

...

    def forward(
            self, batch: Batch
        ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
    
        logits = model(batch.dense_features, batch.sparse_features)
        logits = logits.squeeze(-1)
        
        # The data loader / dataset will create JaggedTensors
        # of the labels. Since labels and logits should be in
        # the same order, I can apply the lengths from labels
        # to the logits too. Once both are JaggedTensors, I can
        # use JaggedTensor.to_padded_dense to convert them to
        # the batch_length x slate_length shape that the
        # ListNet loss function requires. Note that my collate
        # function keeps all of a user's items next to each
        # other, so the lengths properly define user / item
        # boundaries
        lengths = batch.labels.lengths()
        
        # Labels are already in JaggedTensor format, so I only
        # need to reshape them
        listnet_labels = batch.labels.to_padded_dense(
            desired_length=torch.max(lengths),
            padding_value=MAX_NEGATIVE_64INT,
        )
        
        # Convert logits to JaggedTensor and then into the
        # shape that ListNet requires
        listnet_logits = JaggedTensor(
            values=logits,
            lengths=lengths,
        ).to_padded_dense(
            desired_length=torch.max(lengths),
            padding_value=MAX_NEGATIVE_64INT,
        )
        
        # Pass the reshaped logits and labels to the ListNet
        # loss function. The original DLRM configuration used
        # pointwise loss, but we want to optimize rank order,
        # so we use ListNet. However, ListNet requires looking
        # at a request's entire slate, so we need to reshape the
        # logits and labels to be batch_length x slate_length
        # instead of batch_length x 1. It's easier to reshape
        # only the logits and labels than to change the shape of
        # intermediate tensors all throughout the DLRM model.
        loss = self.loss_fn(listnet_logits, listnet_labels.float())

Start with the fbgemm::jagged_to_padded_dense entry point and the provided DLRMTrain forward snippet; inspect how the operator participates in autograd. Reproduce the warning through the shown JaggedTensor.to_padded_dense path and verify that the resolved behavior preserves gradients for logits without the warning.