embedding decomp

pytorch · Jun 15, 2022 · e072fd0 · e072fd0
1 parent 915aecb
commit e072fd0
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/functorch/_src/eager_transforms.py b/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)
diff --git a/functorch/csrc/DynamicLayer.cpp b/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);
 }