Skip to content

make torch.all and torch.any support all dtypes #44779

@rgommers

Description

@rgommers

🚀 Feature

Motivation

Current torch.all/torch.any and Tensor.all, Tensor.any only support bool/uint8. This is inconvenient when trying to write numpy-compatible code like

def somefunc(x):
    # x can be a numpy.ndarray or torch.Tensor instance
    if x.all():
        ....

It is also inconsistent with torch.logical_and et al., which work with all integer and float dtypes just fine.

Note that torch.all/any aren't documented as public API because they don't support all dtypes (see #7539 (comment)), that'd be nice to fix too.

Pitch

  1. Add support for all dtypes to torch.any/all.
  2. Document torch.all and torch.any as public API at https://pytorch.org/docs/master/torch.html#comparison-ops

Alternatives

The builtins all and any do work:

>>> t = torch.tensor([1, 2], dtype=torch.float64)                          
>>> all(t)                                                                 
True
>>> torch.all(t)                                                           
Traceback (most recent call last):
  File "<ipython-input-78-1c35bb87b542>", line 1, in <module>
    torch.all(t)
RuntimeError: all only supports torch.uint8 and torch.bool dtypes

however, they are orders of magnitude slower:

>>> t = torch.ones(10000, dtype=torch.bool)                                
>>> %timeit all(t)                                                         
10.9 ms ± 54.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
>>> %timeit torch.all(t)                                                   
8.74 µs ± 57.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Casting to bool first before using torch.any/all is the other alternative, but that's pretty annoying.

cc @mruberry @rgommers

Metadata

Metadata

Assignees

No one assigned

    Labels

    module: numpyRelated to numpy support, and also numpy compatibility of our operatorstriagedThis issue has been looked at a team member, and triaged and prioritized into an appropriate module

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions