Merge branch 'main' into chhwang/bin-caching

microsoft · Oct 16, 2023 · fee2c96 · fee2c96
2 parents 290eee8 + fdc76bd
commit fee2c96
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/ark/include/kernels/embedding.h b/ark/include/kernels/embedding.h
@@ -171,6 +171,10 @@ DEVICE void embedding(DataType *output, int *input, DataType *weight,
 
     // pWeight: Vec<1, 1, 1, EmbeddingDim>
     int emb_idx = input[un * InDims::CH + uc * InDims::H + uh];
+    if (emb_idx < 0) {
+        emb_idx += WeightShape::H;
+    }
+    // TODO: assert if emb_idx is still negative
     DataType *pWeight = &weight[emb_idx * WeightDims::W];
 
     Broadcast1<Vec<1, 1, 1, WeightDims::W>, Vec<1, 1, 1, EmbeddingDim>, OutDims,

diff --git a/ark/ops/ops_embedding_test.cc b/ark/ops/ops_embedding_test.cc
@@ -28,6 +28,9 @@ void baseline_embedding(std::vector<void *> &outputs,
         for (ark::DimType c = 0; c < osh[1]; ++c) {
             for (ark::DimType h = 0; h < osh[2]; ++h) {
                 int weight_idx = in[in_idx++];
+                if (weight_idx < 0) {
+                    weight_idx += wsh[2];
+                }
                 T *ptr = &weight[weight_idx * wsh[3]];
                 for (ark::DimType w = 0; w < osh[3]; ++w) {
                     out[n * osh[1] * osh[2] * osh[3] + c * osh[2] * osh[3] +
@@ -60,7 +63,12 @@ ark::unittest::State test_embedding() {
     std::vector<int> ti_data;
     for (auto i = 0; i < ti->shape.size(); ++i) {
         // Random indices in [0, num_emb)
-        ti_data.push_back(ark::rand() % num_emb);
+        int rand_idx = ark::rand() % num_emb;
+        if (i % 9 == 0) {
+            // test negative tokens (padding)
+            rand_idx = -rand_idx;
+        }
+        ti_data.push_back(rand_idx);
     }
     auto tw_data = ark::utils::rand_array<T>(tw->shape.size(), 1.0);
     auto result =