embedding decomp

Author: samdow

functorch/_src/eager_transforms.py

@@ -1339,7 +1339,6 @@ def _register_python_decomposition_vmap(decomp):
 _register_jit_decomposition(torch.ops.aten._softmax_backward_data.default)
 _register_jit_decomposition(torch.ops.aten.log_sigmoid_forward.default)
 _register_jit_decomposition(torch.ops.aten.native_layer_norm_backward.default)
-_register_jit_decomposition(torch.ops.aten.native_batch_norm_backward.default)
-_register_jit_decomposition(torch.ops.aten.cudnn_batch_norm_backward.default)
+_register_jit_decomposition(torch.ops.aten.embedding_dense_backward.default)
 _register_python_decomposition_vmap(torch.ops.aten.mse_loss_backward.default)
 _register_python_decomposition_vmap(torch.ops.aten.addr.default)

functorch/csrc/DynamicLayer.cpp

@@ -504,6 +504,7 @@ TORCH_LIBRARY_IMPL(aten, FT_DYNAMIC_LAYER_FRONT_MODE_KEY, m) {
   JVP_DECOMP(native_layer_norm_backward);
   JVP_DECOMP(native_batch_norm_backward);
   JVP_DECOMP(cudnn_batch_norm_backward);
+  JVP_DECOMP(embedding_dense_backward);
 }