8000 __array__ and __array_wrap__ · Issue #2935 · pytorch/pytorch · GitHub
[go: up one dir, main page]

Skip to content

__array__ and __array_wrap__ #2935

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
kohr-h opened this issue Oct 2, 2017 · 3 comments · Fixed by #2945
Closed

__array__ and __array_wrap__ #2935

kohr-h opened this issue Oct 2, 2017 · 3 comments · Fixed by #2945

Comments

@kohr-h
Copy link
Contributor
kohr-h commented Oct 2, 2017

To improve Numpy interoperability, torch tensors could implement the array interface in a quite simple manner:

def __array__(self, dtype=None):
    if dtype is None:
        return self.cpu().numpy()
    else:
        return self.cpu().numpy().astype(dtype, copy=False)


def __array_wrap__(self, array):
    if array.ndim == 0:
        if array.dtype.kind == 'b':
            return bool(array)
        elif array.dtype.kind in ('i', 'u'):
            return int(array)
        elif array.dtype.kind == 'f':
            return float(array)
        elif array.dtype.kind == 'c':
            return complex(array)
        else:
            raise RuntimeError
    else:
        if array.dtype == bool:
            # Workaround, torch has no built-in bool tensor
            cls = torch.ByteTensor
            array = array.astype('uint8')
        else:
            cls = _tensor_cls(array.dtype, use_cuda=self.is_cuda)
        return cls(array)


def _tensor_cls(dtype, use_cuda):
    mapping = {
        np.dtype('float16'): 'Half',
        np.dtype('float32'): 'Float',
        np.dtype('float64'): 'Double',
        np.dtype('int8'): 'Char',
        np.dtype('int16'): 'Short',
        np.dtype('int32'): 'Int',
        np.dtype('int64'): 'Long',
        np.dtype('uint8'): 'Byte',
    }
    try:
        name = mapping[np.dtype(dtype)]
    except KeyError:
        raise ValueError('dtype {} not supported'.format(dtype))
    if use_cuda:
        return getattr(torch.cuda, name)
    else:
        return getattr(torch, name)

It would enable simple things like numpy.asarray to turn tensors into Numpy arrays, and with that make basically all of Numpy applicable to torch tensors (with some restrictions). Is something like that of interest?

@apaszke
Copy link
Contributor
apaszke commented Oct 2, 2017

Yes that looks great! Note that it's much simpler to use torch.from_numpy, which will create a Tensor of appropriate type for you

@kohr-h
Copy link
Contributor Author
kohr-h commented Oct 2, 2017

Alright, I'll make a PR then.

Note that it's much simpler to use torch.from_numpy, which will create a Tensor of appropriate type for you

You're right, I don't know why I overlooked this.

Regarding placement, my suggestion is to define the generic methods like above in torch/__init__.py and then just set the attribute in each class.
Question1 regarding HalfTensor: it doesn't implement numpy() so this code will crash with that tensor class. Should we let it crash or rather leave the attribute off the class? An argument for the former is that if it's left out, numpy.asarray will still silently work but yield a (largely useless) object array with 1 element (the tensor).
Question2 regarding bools: Okay to use ByteTensor? Again, the alternative is to just let it crash.

@apaszke
Copy link
Contributor
apaszke commented Oct 2, 2017
  1. We coud add a numpy() for half as well
  2. Yes, that's fine

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging a pull request may close this issue.

2 participants
0