[WIP] Adding dtype argument to the Unary Ops for dtype promotion (testing on expm1 function) #33063
Conversation
krshrimali
changed the title
expm1 function)
💊 CircleCI build failures summary and remediationsAs of commit 7aac2ac:
Detailed failure analysisOne may explore the probable reasons each build failed interactively on the Dr. CI website. 🕵️ 2 new failures recognized by patternsThe following build failures do not appear to be due to upstream breakage:
|
| @@ -1210,6 +1210,10 @@ | |||
| supports_named_tensor: True | |||
| variants: function, method | |||
|
|
|||
| - func: expm1.dtype(Tensor self, ScalarType dtype) -> Tensor | |||
There was a problem hiding this comment.
does this work as an additional optional argument on the existing signature? e.g. func: expm1(Tensor self, *, ScalarType? dtype=None)
There was a problem hiding this comment.
Unfortunately, no! I have been trying this for a while now, but it always hits a roadblock, and eventually we'll have to change the functions name (not overloading) and maybe even add more helper functions. This seems messy to me at first.
There was a problem hiding this comment.
As-written this allows calls like expm1(self, dtype) (positionally) which wouldn't be allowed by the equivalent numpy api.
There was a problem hiding this comment.
This can be made keyword only argument (by putting the argument after * in native_functions.yaml) without needing an optional argument on the existing signature. I'll make a commit with a keyword-only approach and ping once done.
I'm still trying with optional argument, but it's gonna change a lot of signatures in the yaml file. I shall post the exact errors I get once am done with the possible combinations of modifications I've in my mind.
|
Some additional comments on using
I don't see, so far to my knowledge, any possible breaking to the current flow with using this approach. Although, I may be wrong and am open to any possible challenges to this approach. |
| @@ -1210,7 +1210,7 @@ | |||
| supports_named_tensor: True | |||
| variants: function, method | |||
|
|
|||
| - func: expm1.dtype(Tensor self, ScalarType dtype) -> Tensor | |||
| - func: expm1.dtype(Tensor self, *, ScalarType dtype) -> Tensor | |||
There was a problem hiding this comment.
@nairbv - This enables keyword-only argument feature for expm1 and tests OK. Would like to hear your comments on this.
- This enables
dtypehandling, but as I mentioned before - this does not change any promotion logic for the ops, except allowing the dtype promotion for supported ops with supported dtypes.
For example:
x = torch.tensor(2.3, dtype=torch.float16, device="cuda")
output = torch.expm1(x, dtype=torch.float32)
print(output, output.dtype)
Outputs: tensor(8.9820, device="cuda:0"), torch.float32 (expected)
While:
torch.expm1(x, torch.float32) # not kwonly friendlygives an error (expected):
TypeError: expm1() received an invalid combination of arguments - got (Tensor, torch.dtype), but expected one of:
* (Tensor input, Tensor out)
didn't match because some of the arguments have invalid types: (Tensor, torch.dtype)
* (Tensor input, torch.dtype dtype)
didn't match because some of the arguments have invalid types: (Tensor, torch.dtype)
- Trying to promote
fp16tointdtype also throws an error (expected):RuntimeError: result type Half can't be cast to the desired output type Int.
- The
out=...handling for dtype promotion is absent for now. This can be one of the TODOs, depending on your comments & user requirements.
For example, doing: torch.expm1(x, out=torch.zeros((), dtype=torch.float32)) throws a RuntimeError: RuntimeError: expected dtype Float but got dtype Half. This is because the function called when out=... is passed does not go through our call chain with dtype argument.
Will appreciate your review & comments on this!
There was a problem hiding this comment.
This enables keyword-only argument feature for expm1 and tests OK. Would like to hear your comments on this.
Yes I think the , *, syntax is the correct function definition, though I think we'll want ScalarType? dtype=None so that it is optional. This matches torch.mean and prod, as well as numpy.expm1.
This enables dtype handling, but as I mentioned before - this does not change any promotion logic for the ops, except allowing the dtype promotion for supported ops with supported dtypes.
I think we'll expect the behavior to conform to how these options are used in other functions and/or in numpy.
In the docs for ops like torch.prod/sum/mean:
dtype (torch.dtype, optional) – the desired data type of returned tensor. If specified, the input tensor is casted to dtype before the operation is performed. This is useful for preventing data type overflows. Default: None.
TypeError: expm1() received an invalid combination of arguments - got (Tensor, torch.dtype)
This looks correct. dtype as a positional arg shouldn't be supported here.
doing: torch.expm1(x, out=torch.zeros((), dtype=torch.float32)) throws a RuntimeError
I think we'll need to support this. We'll have to ensure we've set appropriate flags to make the correct checks in validate_dtype etc in TensorIterator.
There was a problem hiding this comment.
@nairbv - Thanks for the comments. Binary Ops use PROMOTE flag by default and is there any reason why Unary Ops should not use the PROMOTE flag by default? As far as we can tell, the current solution works fine for tensors on both CPU and CUDA devices.
There was a problem hiding this comment.
Binary ops iterate through a list containing their two operands to determine a common output type (using type promotion rules implemented in at::result_type).
You're modifying a unary op (only one operand, no list), so it seems unclear what behavior you're trying to control through use of this flag.
aten/src/ATen/native/UnaryOps.cpp
Outdated
| @@ -95,7 +95,11 @@ Tensor& ceil_out(Tensor& result, const Tensor& self) { return unary_op_impl_out( | |||
| Tensor ceil(const Tensor& self) { return unary_op_impl(self, at::ceil_out); } | |||
| Tensor& ceil_(Tensor& self) { return unary_op_impl_(self, at::ceil_out); } | |||
|
|
|||
| Tensor& expm1_out(Tensor& result, const Tensor& self) { return unary_op_impl_out(result, self, expm1_stub); } | |||
| Tensor& expm1_out(Tensor& result, const Tensor& self) { | |||
| if (result.defined()) return unary_op_impl_out(result, self, expm1_stub, /*strategy_promote=*/ true); | |||
There was a problem hiding this comment.
This enables dtype prmotion support for out=... arguments to the expm1 function. Sample outputs:
input: torch.expm1(x /*dtype=torch.float16*/, out = torch.zeros((), dtype=torch.float32, device="cuda")).dtype
output: float32
input: torch.expm1(x /*dtype=torch.float16*/, out = torch.zeros((), dtype=torch.double, device="cuda")).dtype
output: float64
There was a problem hiding this comment.
result.defined() is always going to be true at this point, I think.
| auto iter = TensorIterator(); | ||
| iter.set_check_mem_overlap(check_mem_overlap); | ||
| iter.add_output(out); | ||
| iter.add_input(a); | ||
| iter.num_outputs_ = 1; | ||
| if(promoting == true) iter.promote_common_dtype(); |
There was a problem hiding this comment.
promote_common_dtype implies we're finding the common dtype across all operands and setting the output dtype based on that. These are unary ops, so this doesn't look right.
Also, nit, don't need to check ==true when using a boolean in a condition.
This PR is currently WIP and is intended for discussion on adding dtype argument for the type promotion support of Unary Ops. This PR:
expm1overload withdtypeargument.TensorIterator::unary_opto have promoting argument (default to false,CommonDTypeStrategy::CHECK) - if true, then changes the flag toPROMOTE.unary_op_implandunary_op_impl_out).The objective is to allow the Unary Ops to have a
dtypeargument, to allow the internal promotion logic to promote the dtype if required (and allowed). The question is, if these changes have any unexpected (and unwanted) effect on the current flow of PyTorch.cc: @mcarilli @nairbv