TensorList shape

dali/python/backend_impl.cc

@@ -88,6 +88,15 @@ py::list py_shape(const Tensor<Backend> &t) {
   return as_py_list(t.shape());
 }
 
+template <typename Backend>
+py::list py_shape_list(const TensorList<Backend> &tl) {
+  py::list ret(tl.shape().size());
+  for (int i = 0; i < tl.shape().size(); ++i) {
+    ret[i] = py::tuple(as_py_list(tl.tensor_shape(i)));
+  }
+  return ret;
+}
+
 static string TensorLayoutRepr(const TensorLayout &tl) {
   std::stringstream ss;
   ss << "nvidia.dali.types.TensorLayout('";
@@ -693,6 +702,10 @@ void ExposeTensorList(py::module &m) {
     .def("layout", [](TensorList<CPUBackend> &t) {
       return t.GetLayout().str();
     })
+    .def("shape", &py_shape_list<CPUBackend>,
+      R"code(
+      Shape of the tensor list.
+      )code")
     .def("at", [](TensorList<CPUBackend> &tl, Index id) -> py::array {
           DALI_ENFORCE(IsValidType(tl.type()), "Cannot produce "
               "buffer info for tensor w/ invalid type.");
@@ -909,6 +922,10 @@ void ExposeTensorList(py::module &m) {
       Returns a `TensorListCPU` object being a copy of this `TensorListGPU`.
       )code",
       py::return_value_policy::take_ownership)
+    .def("shape", &py_shape_list<GPUBackend>,
+      R"code(
+      Shape of the tensor list.
+      )code")
     .def("__len__", [](TensorList<GPUBackend> &t) {
           return t.num_samples();
         })

dali/python/nvidia/dali/_debug_mode.py

@@ -35,6 +35,9 @@ def gpu(self):
     def get(self):
         return self._data
 
+    def shape(self):
+        return self._data.shape()
+
     def __add__(self, other):
         return _PipelineDebug.current()._wrap_op_call(_arithm_op, ["add", self, other], {})
 

dali/test/python/test_pipeline_debug.py

@@ -315,3 +315,26 @@ def test_init_config_pipeline():
     pipe_standard = init_config_pipeline(batch_size=8, num_threads=3, device_id=0)
     pipe_debug = init_config_pipeline(batch_size=8, num_threads=3, device_id=0, debug=True)
     compare_pipelines(pipe_standard, pipe_debug, 8, 10)
+
+
+@pipeline_def(batch_size=8, num_threads=3, device_id=0, seed=47, debug=True)
+def shape_pipeline(output_device):
+    jpegs, _ = fn.readers.file(
+        file_root=file_root, shard_id=0, num_shards=2)
+    images = fn.decoders.image(jpegs, device=output_device, output_type=types.RGB)
+    assert images.shape() == [tuple(im.shape()) for im in images.get()]
+    return images
+
+
+def _test_shape_pipeline(device):
+    pipe = shape_pipeline(device)
+    pipe.build()
+    res, = pipe.run()
+
+    # Test TensorList.shape() directly.
+    assert res.shape() == [tuple(im.shape()) for im in res]
+
+
+def test_shape_pipeline():
+    for device in ['cpu', 'mixed']:
+        yield _test_shape_pipeline, device