Skip to content

Allow defining output metadata in the OpSchema#6244

Draft
rostan-t wants to merge 3 commits intoNVIDIA:mainfrom
rostan-t:opschema-metadata
Draft

Allow defining output metadata in the OpSchema#6244
rostan-t wants to merge 3 commits intoNVIDIA:mainfrom
rostan-t:opschema-metadata

Conversation

@rostan-t
Copy link
Collaborator

@rostan-t rostan-t commented Mar 2, 2026

Category:

New feature (non-breaking change which adds functionality)

Description:

It is currently necessary to run operators in order to know the output ndim, dtype and layouts. This PR adds the possibility to specify those metadata when defining an OpSchema.

This is especially relevant for dynamic mode because accessing those metadata currently causes immediate evaluation of tensors, even when EvalMode.eager or EvalMode.defer is used.

Additional information:

Affected modules and functionalities:

Key points relevant for the review:

Tests:

  • Existing tests apply
  • New tests added
    • Python tests
    • GTests
    • Benchmark
    • Other
  • N/A

Checklist

Documentation

  • Existing documentation applies
  • Documentation updated
    • Docstring
    • Doxygen
    • RST
    • Jupyter
    • Other
  • N/A

DALI team only

Requirements

  • Implements new requirements
  • Affects existing requirements
  • N/A

REQ IDs: N/A

JIRA TASK: DALI-2778

rostan-t added 3 commits March 2, 2026 17:14
Signed-off-by: Rostan Tabet <rtabet@nvidia.com>
Signed-off-by: Rostan Tabet <rtabet@nvidia.com>
Signed-off-by: Rostan Tabet <rtabet@nvidia.com>
@rostan-t rostan-t force-pushed the opschema-metadata branch from 6f6cb95 to 98706e8 Compare March 2, 2026 17:16
* @param index Index of the output to set the function for.
* @param fn Function that returns the data type for the given output.
*/
OpSchema &OutputDType(int index, OutputDTypeFunc fn);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think that some kind of policy would be nice - and a sane default, too. Most operators will follow the pattern:

  • use the value of dtype argument, if not None (and it's already a special argument)
  • use the data type of the 1st input if dtype is None.
    I say "is not None" as opposed to "is present", because some operators have a default dtype and adhere to it, even if the user doesn't provide one.

The operators that do NOT follow this pattern are:

  • readers (typically uint8, unpredictable type for NumPy reader, something else for video)
  • arithmetic/math (use type promotion rules)
  • sum (promotes all integers to int64 of the same signeddness; keeps floating point type)
  • normalize, mean, rms, etc (promotes integers to float, keeps floating point type).


/** Try calculating the data type of a given output */
std::optional<DALIDataType> CalculateOutputDType(int index, const OpSpec &spec,
span<const DALIDataType> input_dtypes) const;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. I think that the metadata should be a part of what is now called OpSpec::InOutDeviceDesc
  2. Even if it's not, input_dtypes should all be optional - in many cases we only need the data type of the 1st input. Inability to compute all of them will ruin the static evaluation for pipeline mode.


/** Try calculating the ndim of a given output */
std::optional<int> CalculateOutputNdim(int index, const OpSpec &spec,
span<const int> input_ndims) const;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Similar.


/** Try calculating the layout of a given output */
std::optional<TensorLayout> CalculateOutputLayout(int index, const OpSpec &spec,
span<const TensorLayout> input_layouts) const;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Similar.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants