embedding decomp

Author: samdow

functorch/_src/eager_transforms.py

@@ -1340,5 +1340,6 @@ def _register_python_decomposition_vmap(decomp):
 _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, use_python=True)
+_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

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