Bug fixes for shared_embedding (#1941)

up

Author: metascroy

torchao/experimental/ops/embedding_xbit/op_embedding_xbit-impl.h

@@ -330,7 +330,7 @@ Tensor shared_embedding_out_cpu(
     } else {
       index = index64_ptr[idx];
     }
-    TORCHAO_CHECK(index >= 0 && index < k, "index out of bounds");
+    TORCHAO_CHECK(index >= 0 && index < n, "index out of bounds");
 #if defined(TORCHAO_BUILD_CPU_AARCH64)
     torchao::kernels::cpu::aarch64::embedding::
         shared_embedding<weight_nbit, nr, kr, sr>(

torchao/experimental/ops/embedding_xbit/op_embedding_xbit_executorch.cpp

@@ -56,7 +56,7 @@ DEFINE_OP(8);
   EXECUTORCH_LIBRARY(                                    \
       torchao,                                           \
       "_shared_embedding_" #weight_nbit "bit.out",       \
-      _op_out_##weight_nbit)
+      _shared_op_out_##weight_nbit)
 
 DEFINE_SHARED_OP(1);
 DEFINE_SHARED_OP(2);