Fix LookupTable GPU for scalar inputs

dali/operators/generic/lookup_table.cu

@@ -1,4 +1,4 @@
-// Copyright (c) 2019-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2019-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -48,14 +48,14 @@ void LookupTable<GPUBackend>::RunImpl(Workspace &ws) {
 
  auto num_samples = shape.num_samples();
  samples_.resize(num_samples);
+ TensorListShape<1> collapsed_shape(num_samples);
  for (int sample_id = 0; sample_id < num_samples; sample_id++) {
  samples_[sample_id].output = output.raw_mutable_tensor(sample_id);
  samples_[sample_id].input = input.raw_tensor(sample_id);
+ collapsed_shape.tensor_shape_span(sample_id)[0] = shape.tensor_size(sample_id);
  }
  samples_dev_.from_host(samples_, stream);
 
- auto collapsed_shape = collapse_dims<1>(shape, {std::make_pair(0, shape.sample_dim())});
-
  block_setup_.SetupBlocks(collapsed_shape, true);
  blocks_dev_.from_host(block_setup_.Blocks(), stream);
 

dali/test/python/operator_1/test_lookup_table.py

@@ -12,15 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import numpy as np
-import nvidia.dali.ops as ops
-import nvidia.dali.types as types
+import numpy as np single batch of 4 scalars
 import random
+from nvidia.dali import fn, types, ops, pipeline_def
 from nvidia.dali.pipeline import Pipeline
-
 from test_utils import RandomlyShapedDataIterator
 from test_utils import compare_pipelines
-
+from nose2.tools import params
 
 class LookupTablePipeline(Pipeline):
  def __init__(
@@ -175,3 +173,27 @@
  dictionary_type,
  default_value,
  )
+
+@params("cpu", "gpu")
+def test_scalar(device):
+ @pipeline_def(batch_size=64, num_threads=2, device_id=0)
+ def pipe():
+ raw = np.array([[0, 1, 2, 3]]) # single batch of 4 scalars
+ ids = fn.external_source(source=raw, device=device)
+ scale_keys = [0, 1]
+ scale_values = [100, 200]
+ scale_mat = fn.lookup_table(
+ ids,
+ keys=scale_keys,
+ values=scale_values,
+ device=device,
+ dtype=types.INT64,
+ )
+ return scale_mat, ids
+
+ p = pipe()
+ p.build()
+ scaled, _ = p.run()
+ if device == "gpu":
+ scaled = scaled.as_cpu()
+ assert (scaled.as_array() == [100, 200, 0, 0]).all()