embedding decomp

samdow · samdow · commit e072fd099192 · 2022-06-15T13:41:58.000-04:00
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);
 }
 
 

Original file line number	Diff line number	Diff line change
`@@ -503,6 +503,7 @@ TORCH_LIBRARY_IMPL(aten, FT_DYNAMIC_LAYER_FRONT_MODE_KEY, m) {`
`503`	`503`	`JVP_DECOMP(log_sigmoid_forward);`
`504`	`504`	`JVP_DECOMP(native_layer_norm_backward);`
`505`	`505`	`JVP_DECOMP(native_batch_norm_backward);`
	`506`	`+ JVP_DECOMP(embedding_dense_backward);`
`506`	`507`	`}`
`507`	`508`
`508`	`509`