frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-11-03 23:45:05 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add division overload with rounding_mode selection (#51706)

Browse Source 
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/51706

Pull Request resolved: https://github.com/pytorch/pytorch/pull/50280

As mentioned in gh-43874, this adds a `rounding_mode={'true', 'trunc', 'floor'}`
argument so `torch.div` can be used as a replacement for `floor_divide` during
the transitional period.

I've included dedicated kernels for truncated and floor division which
aren't strictly necessary for float, but do perform significantly better (~2x) than
doing true division followed by a separate rounding kernel.

Note: I introduce new overloads for `aten::div` instead of just adding a default
`rounding_mode` because various JIT passes rely on the exact operator schema.

Test Plan: Imported from OSS

Reviewed By: ngimel

Differential Revision: D26123271

Pulled By: mruberry

fbshipit-source-id: 51a83717602114597ec9c4d946e35a392eb01d46
This commit is contained in:
Peter Bell
2021-02-04 13:00:41 -08:00
committed by Facebook GitHub Bot
parent 949ab213dd
commit b150f150ba
20 changed files with 600 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
torch/_tensor_docs.py
View File

@ -1194,25 +1194,25 @@ See :func:`torch.dist`
""")
add_docstr_all('div', r"""
div(value) -> Tensor
div(value, *, rounding_mode='true') -> Tensor
See :func:`torch.div`
""")
add_docstr_all('div_', r"""
div_(value) -> Tensor
div_(value, *, rounding_mode='true') -> Tensor
In-place version of :meth:`~Tensor.div`
""")
add_docstr_all('divide', r"""
divide(value) -> Tensor
divide(value, *, rounding_mode='true') -> Tensor
See :func:`torch.divide`
""")
add_docstr_all('divide_', r"""
divide_(value) -> Tensor
divide_(value, *, rounding_mode='true') -> Tensor
In-place version of :meth:`~Tensor.divide`
""")