[Large Tensor] Fixed Embedding op (#17599)

* Switched from int to index_t for input_dim

* Implemented fix for output_dim

* Added nightly test for Embedding

* Set const value for output dim

* More standardization via const param

src/operator/tensor/indexing_op.h

@@ -66,8 +66,8 @@ enum QuantizedEmbeddingOpResource {kTempSpace};
 
 
 struct SparseEmbeddingParam: public dmlc::Parameter<SparseEmbeddingParam> {
- int input_dim;
- int output_dim;
+ index_t input_dim;
+ index_t output_dim;
  int dtype;
  bool deterministic;
  DMLC_DECLARE_PARAMETER(SparseEmbeddingParam) {
@@ -89,8 +89,8 @@ struct SparseEmbeddingParam: public dmlc::Parameter<SparseEmbeddingParam> {
 };
 
 struct EmbeddingParam: public dmlc::Parameter<EmbeddingParam> {
- int input_dim;
- int output_dim;
+ index_t input_dim;
+ index_t output_dim;
  int dtype;
  bool sparse_grad;
  DMLC_DECLARE_PARAMETER(EmbeddingParam) {

tests/nightly/test_large_array.py

@@ -38,6 +38,7 @@
 SMALL_X = 100
 SMALL_Y = 50
 LARGE_SIZE = LARGE_X * SMALL_Y
+LARGE_TENSOR_SHAPE = 2**32
 
 
 def test_nn():
@@ -467,6 +468,17 @@ def check_col2im():
  assert res.shape[2] == 2
  assert res.shape[3] == 2
  assert res.shape[4] == 1
+ def check_embedding():
+ data = nd.random_normal(shape=(LARGE_TENSOR_SHAPE, 1))
+ weight = nd.random_normal(shape=(LARGE_TENSOR_SHAPE, 1))
+ input_dim = LARGE_TENSOR_SHAPE
+ output_dim = 1
+
+ out = nd.Embedding(data=data, weight=weight, input_dim=input_dim, output_dim=output_dim)
+
+ assert out.shape[0] == LARGE_TENSOR_SHAPE
+ assert out.shape[1] == 1
+ assert out.shape[2] == 1
 
  check_gluon_embedding()
  check_fully_connected()
@@ -488,6 +500,7 @@ def check_col2im():
  check_l2_normalization()
  check_instance_norm()
  check_col2im()
+ check_embedding()
 
 
 def test_tensor():