Fix Stack (#925)

* return the correct layer

* unskip the test

Author: wgb14

k2/csrc/ragged_test.cu

@@ -265,10 +265,7 @@ TEST(RaggedShapeOpsTest, UnstackRandom) {
       for (size_t i = 0; i < out.size(); ++i) {
         out_ptr.emplace_back(&(out[i]));
       }
-      // There is a bug in `Stack` for stacking a shape itself,
-      // not urgent, so skipping here.
-      // TODO: Remove this line when the bug fixed.
-      if (out.size() == 1) continue;
+
       auto dest = Stack(axis, out.size(), out_ptr.data());
       dest = RemoveEmptyLists(dest, axis);
 

k2/csrc/ragged_utils.cu

@@ -92,7 +92,9 @@ RaggedShape IntersperseRaggedLayer(int32_t layer,
     if (merge_map)
       *(reinterpret_cast<Array1<int32_t>*>(merge_map)) =
           Range(src[0]->Context(), src[0]->TotSize(layer + 1), 0);
-    return *src[0];
+    std::vector<RaggedShapeLayer> layers;
+    layers.emplace_back(src[0]->Layers()[layer]);
+    return RaggedShape(layers);
   }
 
   std::vector<int32_t*> row_splits_ptrs_vec(num_srcs);