pytorch
da10ccd3 - Implements cpu_kernel_multiple_outputs and torch.frexp (#51097)

Commit
3 years ago
Implements cpu_kernel_multiple_outputs and torch.frexp (#51097) Summary: Close https://github.com/pytorch/pytorch/issues/51108 Related https://github.com/pytorch/pytorch/issues/38349 This PR implements the `cpu_kernel_multiple_outputs` to support returning multiple values in a CPU kernel. ```c++ auto iter = at::TensorIteratorConfig() .add_output(out1) .add_output(out2) .add_input(in1) .add_input(in2) .build(); at::native::cpu_kernel_multiple_outputs(iter, [=](float a, float b) -> std::tuple<float, float> { float add = a + b; float mul = a * b; return std::tuple<float, float>(add, mul); } ); ``` The `out1` will equal to `torch.add(in1, in2)`, while the result of `out2` will be `torch.mul(in1, in2)`. It helps developers implement new torch functions that return two tensors more conveniently, such as NumPy-like functions [divmod](https://numpy.org/doc/1.18/reference/generated/numpy.divmod.html?highlight=divmod#numpy.divmod) and [frexp](https://numpy.org/doc/stable/reference/generated/numpy.frexp.html#numpy.frexp). This PR adds `torch.frexp` function to exercise the new functionality provided by `cpu_kernel_multiple_outputs`. Pull Request resolved: https://github.com/pytorch/pytorch/pull/51097 Reviewed By: albanD Differential Revision: D26982619 Pulled By: heitorschueroff fbshipit-source-id: cb61c7f2c79873ab72ab5a61cbdb9203531ad469
Parents
Loading