Add unique op

[a-gardner1]

Add support for exporting torch.unique following the conclusion of pytorch/pytorch#113118.

[codecov]

Codecov Report

Attention: Patch coverage is 57.14286% with 18 lines in your changes missing coverage. Please review.

Project coverage is 77.50%. Comparing base (69ae7f4) to head (f9885f1).
Report is 171 commits behind head on main.

Files with missing lines	Patch %	Lines
onnxscript/function_libs/torch_lib/ops/core.py	57.14%	16 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1547      +/-   ##
==========================================
- Coverage   77.56%   77.50%   -0.07%     
==========================================
  Files         214      216       +2     
  Lines       23186    23381     +195     
  Branches     3975     4033      +58     
==========================================
+ Hits        17984    18121     +137     
- Misses       4433     4477      +44     
- Partials      769      783      +14     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[justinchuby]

Thanks for your contribution! Could you follow the CLA bot's instruction to get that cleared?

[justinchuby]

I would remove this try-catch as the function here is symbolic; we don't expect them to raise any errors

[a-gardner1]

Addressed in b528a6a

[a-gardner1]

Thanks for your contribution! Could you follow the CLA bot's instruction to get that cleared?

Yea, I may have jumped the gun a bit. Working on officially getting permission from my employer.

[a-gardner1]

@a-gardner1 please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.

@microsoft-github-policy-service agree [company="{your company}"]

Options:

• (default - no company specified) I have sole ownership of intellectual property rights to my Submissions and I am not making Submissions in the course of work for my employer.

@microsoft-github-policy-service agree

• (when company given) I am making Submissions in the course of work for my employer (or my employer has intellectual property rights in my Submissions by contract or applicable law). I have permission from my employer to make Submissions and enter into this Agreement on behalf of my employer. By signing below, the defined term “You” includes me and my employer.

@microsoft-github-policy-service agree company="Microsoft"

Contributor License Agreement

@microsoft-github-policy-service agree [company="Radiance Technologies"]

@microsoft-github-policy-service agree company="Radiance Technologies"

[a-gardner1]

@microsoft-github-policy-service agree company="Radiance Technologies"

[a-gardner1]

@justinchuby I'm not sure what the preferred behavior is here. Should we match torch.unique and ignore the sorted argument (i.e., always sort in aten_unique) or respect the argument and deviate in accordance with this matcher?

[justinchuby]

I wonder if the argument leads to different behavior in cuda/cpu etc? I assume sorted=False means it can be sorted, but it doesn't need to be; and there are some potential performance gain by turning it off. If that's the interpretation I would keep the argument. Otherwise ignoring the argument and matching behavior would also be nice.

[a-gardner1]

I am investigating differences in behavior between cuda/cpu and have found at least one already (unique_dim on CPU ignores the return_inverse and return_counts arguments whereas the CUDA impl does not). How should these differences be handled? Can the op registration be conditioned by the device somehow, or should I favor CUDA over CPU?

[justinchuby]

Matching CUDA for now is preferable. Thanks!

[a-gardner1]

I want to note that the way that this function was written in 1d74d59 is functionally equivalent but yields an error in onnxscript.Scope.lookup_or_create because it causes modified to be True in onnxscript.optimizer.optimize, thus causing a second loop of optimization that crashes in the first call to inline_simple_functions.

This seems indicative of a potential bug to me, but I am not knowledgeable enough about the codebase to suggest a cause or fix.

[justinchuby]

cc @gramalingam

[justinchuby]

Thanks for completing the CLA. I will take a look next week

[justinchuby]

I think we need to always return the same number of values. Consider returning None when they are not available?

[a-gardner1]

Doing so deviates from the behavior of torch.unique and causes this assertion in the unit tests to fail:

onnxscript/tests/function_libs/torch_lib/ops_test.py

Line 251 in 69ae7f4

assert len(flattened_torch_outputs) == len(flattened_function_outputs)

Please advise on how to address this.

[justinchuby]

Does torch.ops.aten.unique exhibit the same behavior? If it always returns three variables, consider creating a new OpInfo for torch.ops.aten.unique similar to

onnxscript/tests/function_libs/torch_lib/extra_opinfo.py

Lines 2105 to 2114 in a5ed079

	opinfo_core.OpInfo(
	"ops.aten._native_batch_norm_legit.no_stats",
	aten_name="_native_batch_norm_legit.no_stats",
	dtypes=common_dtype.floating_types_and(torch.bfloat16),
	dtypesIfCUDA=common_dtype.floating_types_and(torch.float16, torch.bfloat16),
	supports_forward_ad=True,
	supports_fwgrad_bwgrad=True,
	assert_jit_shape_analysis=True,
	sample_inputs_func=sample_inputs__native_batch_norm_legit_no_stats,
	),

. You may remove the xfail with the custom OpInfo too because you may simply remove the xfail cases.

You may adapt the sample function from https://github.com/pytorch/pytorch/blob/b948b1ad7a9cf61c9692506c60c295fd40e00f43/torch/testing/_internal/common_methods_invocations.py#L3346-L3372

[a-gardner1]

Thanks for the pointer to extra_opinfo. It turns out torch.ops.aten.unique does not exist, but torch.ops.aten._unique does. Added OpInfo for it, _unique2, and unique_dim in 14d03b5

[justinchuby]

Is this a bug we should fix elsewhere? saw comment below

[a-gardner1]

I think this could possibly be considered a different bug. The other one is a side-effect of onnxscript.optimizer.constant_folding.fold_constants, whereas this one is a side-effect of the function linked below, which converts the names of unused outputs to empty strings but only removes them if they are trailing. Since inverse_indices and counts are used, it leads to an error being raised in onnxscript.Scope.lookup_or_create due to the empty string name given to indices.

onnxscript/onnxscript/optimizer/remove_unused.py

Line 14 in 69ae7f4

def remove_unused_optional_outputs(

[justinchuby]

I suggest removing all hacks. I will go fix what's necessary where the bug is. We are also moving to prefer trace_only=True for new functions so if you can include the flag in @torch_op that would be awesome.

[a-gardner1]

That would be awesome. The hacks are definitely getting out of hand. I'll wait for that fix so that I can continue to test with this locally.

[justinchuby]

Do you have a short script handy that will reproduce the error?

[a-gardner1]

if __name__ == '__main__':
    import logging
    import torch
    import numpy as np
    import onnx
    import onnxruntime as ort
    for i in range(16):
        sorted = bool(i & 1)
        return_inverse = bool((i & 2) > 1)
        return_counts = bool((i & 4) > 1)
        dim = 0 if bool((i & 8) > 1) else None

        print(
            f"Testing sorted={sorted}, return_inverse={return_inverse}, return_counts={return_counts}, dim={dim}"
        )

        def test_function(
                x: torch.Tensor,
                s: bool = sorted,
                ri: bool = return_inverse,
                rc: bool = return_counts,
                d: int | None = dim) -> Any:
            result = torch.unique(
                x,
                sorted=s,
                return_inverse=ri,
                return_counts=rc,
                dim=d)
            return result

        onnx_program = torch.onnx.dynamo_export(
            test_function,
            torch.arange(10),
            export_options=torch.onnx.ExportOptions(
                dynamic_shapes=True,
                diagnostic_options=torch.onnx.DiagnosticOptions(
                    verbosity_level=logging.DEBUG)))
        onnx_program.save("torch_unique.onnx")
        onnx_inputs = onnx_program.adapt_torch_inputs_to_onnx(torch.arange(10))
        onnx_outputs = onnx_program(*onnx_inputs)
        loaded_onnx_program = onnx.load("torch_unique.onnx")
        onnx.checker.check_model(loaded_onnx_program)
        ort_session = ort.InferenceSession("torch_unique.onnx")
        inputs = np.random.randint(0, 10, 10)
        print(f"Inputs: {inputs}")
        outputs = ort_session.run(None,
                                  {"l_x_": inputs})
        print(f"Outputs: {outputs}")
    print("Success")

[a-gardner1]

Oh, you should also test using the nightly release of PyTorch with the changes in pytorch/pytorch#126561.

[a-gardner1]

Is trace_only=True expected to require significant changes to the way one implements an op? It appears that enabling the flag breaks passing a value to op.ConstantOfShape and also breaks indexing a shape.

For example, op.ConstantOfShape([0], value=[0]) must become op.Cast(op.ConstantOfShape([0]), to=INT64.dtype), and output_size[dim] must become op.Slice(output_size, [dim], [dim+1]).

[justinchuby]

Your observation is correct. This may be the case because the gaps in implementation we have. Bridging the gaps is in our roadmap but is not the highest priority for the team.