Start adding output metadata inference to operator schemas

rostan-t · rostan-t · commit 5e10607e6af6 · 2026-03-16T12:15:49.000Z
Signed-off-by: Rostan Tabet &lt;rtabet@nvidia.com&gt;
diff --git a/dali/operators/bbox/bb_flip.cc b/dali/operators/bbox/bb_flip.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2017-2026, 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.
@@ -37,7 +37,10 @@ system, that is 0.0-1.0)code")
                     1, true)
     .AddOptionalArg("vertical",
                     R"code(Flip vertical dimension.)code",
-                    0, true);
+                    0, true)
+    .OutputDType(0, [](const OpSpec &, span<const DALIDataType> in) { return in[0]; })
+    .OutputNdim(0, [](const OpSpec &, span<const int> in) { return in[0]; })
+    .OutputLayout(0, [](const OpSpec &, span<const TensorLayout> in) { return in[0]; });
 
 void BbFlipCPU::RunImpl(Workspace &ws) {
   const auto &input = ws.Input<CPUBackend>(0);
diff --git a/dali/operators/decoder/image_decoder.cc b/dali/operators/decoder/image_decoder.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2019-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2019-2026, 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.
@@ -191,7 +191,10 @@ Please note that GPU acceleration for JPEG 2000 decoding is only available for C
   .NumInput(1)
   .NumOutput(1)
   .AddParent("ImageDecoderAttr")
-  .AddParent("CachedDecoderAttr");
+  .AddParent("CachedDecoderAttr")
+  .OutputDType(0, [](const OpSpec &, span<const DALIDataType>) { return DALI_UINT8; })
+  .OutputNdim(0, [](const OpSpec &, span<const int>) { return 3; })
+  .OutputLayout(0, [](const OpSpec &, span<const TensorLayout>) { return "HWC"; });
 
 // Fused
 
@@ -309,7 +312,10 @@ of the slice (s0, s1, s2, …).
 
 Integer coordinates are interpreted as absolute coordinates, while float coordinates can be
 interpreted as absolute or relative coordinates, depending on the value of
-`normalized_shape`.)code");
+`normalized_shape`.)code")
+  .OutputDType(0, [](const OpSpec &, span<const DALIDataType>) { return DALI_UINT8; })
+  .OutputNdim(0, [](const OpSpec &, span<const int>) { return 3; })
+  .OutputLayout(0, [](const OpSpec &, span<const TensorLayout>) { return "HWC"; });
 
 
 // Deprecated aliases
diff --git a/dali/operators/generic/cast.cc b/dali/operators/generic/cast.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2017-2026, 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.
@@ -71,7 +71,12 @@ DALI_SCHEMA(Cast)
     .NumOutput(1)
     .AllowSequences()
     .SupportVolumetric()
-    .AddTypeArg("dtype", R"code(Output data type.)code");
+    .AddTypeArg("dtype", R"code(Output data type.)code")
+    .OutputDType(0, [](const OpSpec &spec, span<const DALIDataType>) {
+      return spec.GetArgument<DALIDataType>("dtype");
+    })
+    .OutputNdim(0, [](const OpSpec &, span<const int> in) { return in[0]; })
+    .OutputLayout(0, [](const OpSpec &, span<const TensorLayout> in) { return in[0]; });
 
 DALI_SCHEMA(CastLike)
     .DocStr("Cast the first tensor to the type of the second tensor.")
diff --git a/dali/operators/generic/flip.cc b/dali/operators/generic/flip.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2019-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2019-2026, 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.
@@ -34,7 +34,10 @@ and depthwise).)code")
     .AddOptionalArg("depthwise", R"code(Flip the depthwise dimension.)code", 0, true)
     .InputLayout({"FDHWC", "FHWC", "DHWC", "HWC", "FCDHW", "FCHW", "CDHW", "CHW"})
     .AllowSequences()
-    .SupportVolumetric();
+    .SupportVolumetric()
+    .OutputDType(0, [](const OpSpec &, span<const DALIDataType> in) { return in[0]; })
+    .OutputNdim(0, [](const OpSpec &, span<const int> in) { return in[0]; })
+    .OutputLayout(0, [](const OpSpec &, span<const TensorLayout> in) { return in[0]; });
 
 
 template <>
diff --git a/dali/operators/image/color/brightness_contrast.cc b/dali/operators/image/color/brightness_contrast.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2019-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2019-2026, 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.
@@ -89,7 +89,15 @@ This operator can also change the type of data.)code")
     .NumOutput(1)
     .AllowSequences()
     .SupportVolumetric()
-    .InputLayout({"FHWC", "DHWC", "HWC"});
+    .InputLayout({"FHWC", "DHWC", "HWC"})
+    .OutputDType(0, [](const OpSpec &spec, span<const DALIDataType> in) {
+      DALIDataType dtype;
+      if (spec.TryGetArgument(dtype, "dtype"))
+        return dtype;
+      return in[0];
+    })
+    .OutputNdim(0, [](const OpSpec &, span<const int> in) { return in[0]; })
+    .OutputLayout(0, [](const OpSpec &, span<const TensorLayout> in) { return in[0]; });
 
 DALI_REGISTER_OPERATOR(BrightnessContrast, BrightnessContrastCpu, CPU)
 DALI_REGISTER_OPERATOR(Brightness, BrightnessContrastCpu, CPU);
diff --git a/dali/operators/image/color/color_twist.cc b/dali/operators/image/color/color_twist.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2020-2026, 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.
@@ -50,7 +50,15 @@ they would in case of rotation.)code",
 If a value is not set, the input type is used.)code",
                     DALI_UINT8)
     .InputLayout(0, {"HWC", "FHWC", "DHWC"})
-    .AllowSequences();
+    .AllowSequences()
+    .OutputDType(0, [](const OpSpec &spec, span<const DALIDataType> in) {
+      DALIDataType dtype;
+      if (spec.TryGetArgument(dtype, "dtype"))
+        return dtype;
+      return in[0];
+    })
+    .OutputNdim(0, [](const OpSpec &, span<const int> in) { return in[0]; })
+    .OutputLayout(0, [](const OpSpec &, span<const TensorLayout> in) { return in[0]; });
 
 DALI_SCHEMA(ColorTransformBase)
     .DocStr(R"code(Base Schema for color transformations operators.)code")
diff --git a/dali/operators/image/crop/bbox_crop.cc b/dali/operators/image/crop/bbox_crop.cc
@@ -364,7 +364,16 @@ if the fraction of their area within the ROI is greater than or equal to the thr
 For example, when `bbox_prune_threshold=0.2` bboxes that have at least 20% of their original area within
 the ROI are kept, bboxes less than or equal to are pruned. If `bbox_prune_threshold=0.0`, all boxes that
 have some presence in the ROI are kept.)code",
-        nullptr);
+        nullptr)
+    .OutputDType(0, [](const OpSpec &, span<const DALIDataType>) { return DALI_FLOAT; })
+    .OutputDType(1, [](const OpSpec &, span<const DALIDataType>) { return DALI_FLOAT; })
+    .OutputDType(2, [](const OpSpec &, span<const DALIDataType>) { return DALI_FLOAT; })
+    .OutputNdim(0, [](const OpSpec &, span<const int>) { return 1; })
+    .OutputNdim(1, [](const OpSpec &, span<const int>) { return 1; })
+    .OutputNdim(2, [](const OpSpec &, span<const int>) { return 2; })
+    .OutputLayout(0, [](const OpSpec &, span<const TensorLayout>) { return TensorLayout{}; })
+    .OutputLayout(1, [](const OpSpec &, span<const TensorLayout>) { return TensorLayout{}; })
+    .OutputLayout(2, [](const OpSpec &, span<const TensorLayout>) { return TensorLayout{}; });
 
 template <int ndim>
 class RandomBBoxCropImpl : public OpImplBase<CPUBackend> {
diff --git a/dali/operators/image/crop/crop.cc b/dali/operators/image/crop/crop.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2019-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2019-2026, 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.
@@ -36,7 +36,15 @@ DALI_SCHEMA(Crop)
     .DeprecateArg("image_type", "0.24")
     .AddParent("CropAttr")
     .AddParent("OutOfBoundsAttr")
-    .AddParent("SliceBase");
+    .AddParent("SliceBase")
+    .OutputDType(0, [](const OpSpec &spec, span<const DALIDataType> in) {
+      DALIDataType dtype;
+      if (spec.TryGetArgument(dtype, "dtype"))
+        return dtype;
+      return in[0];
+    })
+    .OutputNdim(0, [](const OpSpec &, span<const int> in) { return in[0]; })
+    .OutputLayout(0, [](const OpSpec &, span<const TensorLayout> in) { return in[0]; });
 
 // Register operator
 DALI_REGISTER_OPERATOR(Crop, Crop<CPUBackend>, CPU);
diff --git a/dali/operators/image/crop/crop_mirror_normalize.cc b/dali/operators/image/crop/crop_mirror_normalize.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2019-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2019-2026, 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.
@@ -71,7 +71,14 @@ This argument is useful when using integer outputs to improve dynamic range util
 This argument is useful when using unsigned integer outputs to improve dynamic range utilization.)",
     0.0f)
   .AddParent("CropAttr")
-  .AddParent("OutOfBoundsAttr");
+  .AddParent("OutOfBoundsAttr")
+  .OutputDType(0, [](const OpSpec &spec, span<const DALIDataType>) {
+    return spec.GetArgument<DALIDataType>("dtype");
+  })
+  .OutputNdim(0, [](const OpSpec &, span<const int> in) { return in[0]; })
+  .OutputLayout(0, [](const OpSpec &spec, span<const TensorLayout>) {
+    return spec.GetArgument<TensorLayout>("output_layout");
+  });
 
 DALI_REGISTER_OPERATOR(CropMirrorNormalize, CropMirrorNormalize<CPUBackend>, CPU);
 
diff --git a/dali/operators/image/resize/resize.cc b/dali/operators/image/resize/resize.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2017-2026, 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.
@@ -34,7 +34,15 @@ DALI_SCHEMA(Resize)
   .SupportVolumetric()
   .AllowSequences()
   .AddParent("ResizeAttr")
-  .AddParent("ResamplingFilterAttr");
+  .AddParent("ResamplingFilterAttr")
+  .OutputDType(0, [](const OpSpec &spec, span<const DALIDataType> in) {
+    DALIDataType dtype;
+    if (spec.TryGetArgument(dtype, "dtype"))
+      return dtype;
+    return in[0];
+  })
+  .OutputNdim(0, [](const OpSpec &, span<const int> in) { return in[0]; })
+  .OutputLayout(0, [](const OpSpec &, span<const TensorLayout> in) { return in[0]; });
 
 template<typename Backend>
 Resize<Backend>::Resize(const OpSpec &spec)
diff --git a/dali/operators/random/coin_flip_cpu.cc b/dali/operators/random/coin_flip_cpu.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2020-2026, 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.
@@ -35,7 +35,24 @@ a single value per sample is generated.
     .AddOptionalArg<float>("probability",
       R"code(Probability of value 1.)code",
       0.5f, true)
-    .AddParent("RNGAttr");
+    .AddParent("RNGAttr")
+    .OutputDType(0, [](const OpSpec &spec, span<const DALIDataType>) {
+      DALIDataType dtype;
+      if (spec.TryGetArgument(dtype, "dtype"))
+        return dtype;
+      return DALI_INT32;
+    })
+    .OutputNdim(0, [](const OpSpec &spec, span<const int> in) -> std::optional<int> {
+      if (!in.empty())
+        return in[0];
+      std::vector<int> shape;
+      if (spec.TryGetArgument(shape, "shape"))
+        return static_cast<int>(shape.size());
+      if (spec.ArgumentDefined("shape"))  // shape specified but no workspace
+        return std::nullopt;
+      return 0;
+    })
+    .OutputLayout(0, [](const OpSpec &, span<const TensorLayout>) { return TensorLayout{}; });
 
 DALI_REGISTER_OPERATOR(random__CoinFlip, CoinFlip<CPUBackend>, CPU);
 
diff --git a/dali/operators/random/uniform_distribution_cpu.cc b/dali/operators/random/uniform_distribution_cpu.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2021-2026, 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.
@@ -49,7 +49,24 @@ This argument is mutually exclusive with `values`.
 
 This argument is mutually exclusive with `range`.)code",
       nullptr, true)
-    .AddParent("RNGAttr");
+    .AddParent("RNGAttr")
+    .OutputDType(0, [](const OpSpec &spec, span<const DALIDataType>) {
+      DALIDataType dtype;
+      if (spec.TryGetArgument(dtype, "dtype"))
+        return dtype;
+      return DALI_FLOAT;
+    })
+    .OutputNdim(0, [](const OpSpec &spec, span<const int> in) -> std::optional<int> {
+      if (!in.empty())
+        return in[0];
+      std::vector<int> shape;
+      if (spec.TryGetArgument(shape, "shape"))
+        return static_cast<int>(shape.size());
+      if (spec.ArgumentDefined("shape"))  // shape specified but no workspace
+        return std::nullopt;
+      return 0;
+    })
+    .OutputLayout(0, [](const OpSpec &, span<const TensorLayout>) { return TensorLayout{}; });
 
 DALI_REGISTER_OPERATOR(random__Uniform, UniformDistribution<CPUBackend>, CPU);
 
diff --git a/dali/test/python/experimental_mode/test_output_metadata.py b/dali/test/python/experimental_mode/test_output_metadata.py
@@ -0,0 +1,100 @@
+# Copyright (c) 2026, 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import nvidia.dali.experimental.dynamic as ndd
+from ndd_utils import eval_modes
+from collections.abc import Iterable
+from nose2.tools import params
+
+
+def assert_correct_metadata(
+    *tensors: ndd.Tensor,
+    attributes: Iterable[str] = ("dtype", "ndim", "layout"),
+):
+    # First collect all metadata to avoid triggering transitive evaluations
+    actual_metadata = []
+    for tensor in tensors:
+        actual = {attr: getattr(tensor, attr) for attr in attributes}
+        assert tensor._invocation_result._invocation._results is None
+        actual_metadata.append(actual)
+
+    for tensor, actual in zip(tensors, actual_metadata):
+        tensor.evaluate()
+        assert tensor._invocation_result._invocation._results is not None
+        expected = {attr: getattr(tensor, attr) for attr in attributes}
+
+        for attr in attributes:
+            assert (
+                actual[attr] == expected[attr]
+            ), f"Expected {attr} to be {expected[attr]}, but got {actual[attr]}"
+
+
+# --- Image operators ---
+
+
+@eval_modes(ndd.EvalMode.deferred)
+@params(
+    (ndd.flip, dict(horizontal=1)),
+    (ndd.crop, dict(crop=(2, 2))),
+    (ndd.crop, dict(crop=(2, 2), dtype=ndd.float32)),
+    (ndd.cast, dict(dtype=ndd.float32)),
+    (ndd.resize, dict(resize_x=2, resize_y=2)),
+    (ndd.hsv, dict(dtype=ndd.float32, hue=10.0)),
+    (ndd.hsv, dict(hue=10.0)),
+    (ndd.brightness_contrast, dict(brightness=1.2)),
+    (ndd.crop_mirror_normalize, dict(crop=(2, 2))),
+    (ndd.crop_mirror_normalize, dict(crop=(2, 2), dtype=ndd.float32)),
+)
+def test_image_ops(func, kwargs):
+    tensor = ndd.zeros(shape=(4, 4, 3), dtype=ndd.uint8, layout="HWC")
+    assert_correct_metadata(func(tensor, **kwargs))
+
+
+# --- Random generators ---
+
+
+@eval_modes(ndd.EvalMode.deferred)
+@params(
+    ndd.random.coin_flip,
+    ndd.random.uniform,
+)
+def test_random_generators(func):
+    assert_correct_metadata(func())
+
+
+@eval_modes(ndd.EvalMode.deferred)
+@params(
+    (ndd.random.coin_flip, dict(shape=[4])),
+    (ndd.random.uniform, dict(shape=[4])),
+)
+def test_random_generators_with_shape(func, kwargs):
+    # ndim is excluded because the `shape` kwarg determines ndim at run time,
+    # and the schema cannot infer it without executing the operator.
+    assert_correct_metadata(func(**kwargs), attributes=("dtype", "layout"))
+
+
+# --- Other operators ---
+
+
+@eval_modes(ndd.EvalMode.deferred)
+def test_bb_flip():
+    bbox = ndd.zeros(shape=(1, 4), dtype=ndd.float32)
+    assert_correct_metadata(ndd.bb_flip(bbox))
+
+
+@eval_modes(ndd.EvalMode.deferred)
+def test_random_bbox_crop():
+    bboxes = ndd.zeros(shape=(1, 4), dtype=ndd.float32)
+    anchor, shape, out_bboxes = ndd.random_bbox_crop(bboxes)
+    assert_correct_metadata(anchor, shape, out_bboxes)
