From ffcae3f7feff3933d45729a7c37b99834ed4cfe7 Mon Sep 17 00:00:00 2001 From: Yunjiang Jiang Date: Mon, 8 Dec 2025 09:24:58 -0800 Subject: [PATCH] Fix jagged_to_padded_dense autograd (#5191) Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2189 This change will remove the warning that those ops don't have registered autograd. ``` torch/autograd/graph.py:857: 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 fbcode/caffe2/torch/csrc/autograd/autograd_not_implemented_fallback.cpp:71.) ``` Reviewed By: r-barnes Differential Revision: D85714299 --- .../src/jagged_tensor_ops/jagged_tensor_ops_autograd.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/fbgemm_gpu/src/jagged_tensor_ops/jagged_tensor_ops_autograd.cpp b/fbgemm_gpu/src/jagged_tensor_ops/jagged_tensor_ops_autograd.cpp index ac14fdd975..740be8956c 100644 --- a/fbgemm_gpu/src/jagged_tensor_ops/jagged_tensor_ops_autograd.cpp +++ b/fbgemm_gpu/src/jagged_tensor_ops/jagged_tensor_ops_autograd.cpp @@ -989,6 +989,8 @@ std::tuple jagged_slice( TORCH_LIBRARY_IMPL(fbgemm, Autograd, m) { m.impl("jagged_2d_to_dense", TORCH_FN(fbgemm_gpu::jagged_2d_to_dense)); m.impl("jagged_1d_to_dense", TORCH_FN(fbgemm_gpu::jagged_1d_to_dense)); + m.impl( + "jagged_to_padded_dense", TORCH_FN(fbgemm_gpu::jagged_to_padded_dense)); m.impl( "jagged_dense_dense_elementwise_add_jagged_output", TORCH_FN(fbgemm_gpu::jagged_dense_dense_elementwise_add_jagged_output)); @@ -1016,6 +1018,4 @@ TORCH_LIBRARY_IMPL(fbgemm, Autograd, m) { TORCH_LIBRARY_IMPL(fbgemm, CompositeImplicitAutograd, m) { m.impl("jagged_index_select", TORCH_FN(fbgemm_gpu::jagged_index_select_2d)); m.impl("dense_to_jagged", TORCH_FN(fbgemm_gpu::dense_to_jagged)); - m.impl( - "jagged_to_padded_dense", TORCH_FN(fbgemm_gpu::jagged_to_padded_dense)); }