[JIT] Batch op. batch_sum return a mask which is on cpu

[Strasser-Pablo]

🐛 Bug

To Reproduce

Steps to reproduce the behavior:

1. Create a computation graph with a sum operation and another operation which is on gpu.

TensorIterator expected type CUDAByteType but got CPUByteType (check_type_conversions at ../aten/src/ATen/native/TensorIterator.cpp:547)
frame #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&) + 0x6a (0x7f860917d8fa in /usr/local/lib/python3.5/dist-packages/torch/lib/libc10.so)
frame #1: at::TensorIterator::check_type_conversions() + 0x24d (0x7f8668082b2d in /usr/local/lib/python3.5/dist-packages/torch/lib/libcaffe2.so)
frame #2: at::TensorIterator::Builder::build() + 0x2ef (0x7f8668087d3f in /usr/local/lib/python3.5/dist-packages/torch/lib/libcaffe2.so)
frame #3: at::TensorIterator::binary_op(at::Tensor&, at::Tensor const&, at::Tensor const&) + 0x309 (0x7f8668089589 in /usr/local/lib/python3.5/dist-packages/torch/lib/libcaffe2.so)
frame #4: at::native::mul_out(at::Tensor&, at::Tensor const&, at::Tensor const&) + 0x95 (0x7f8667f284e5 in /usr/local/lib/python3.5/dist-packages/torch/lib/libcaffe2.so)
frame #5: at::native::mul(at::Tensor const&, at::Tensor const&) + 0x43 (0x7f8667f29553 in /usr/local/lib/python3.5/dist-packages/torch/lib/libcaffe2.so)
frame #6: at::TypeDefault::mul(at::Tensor const&, at::Tensor const&) const + 0x4d (0x7f86682a601d in /usr/local/lib/python3.5/dist-packages/torch/lib/libcaffe2.so)
frame #7: torch::autograd::VariableType::mul(at::Tensor const&, at::Tensor const&) const + 0x302 (0x7f8609808842 in /usr/local/lib/python3.5/dist-packages/torch/lib/libtorch.so.1)
frame #8: <unknown function> + 0x54df7b (0x7f86098d3f7b in /usr/local/lib/python3.5/dist-packages/torch/lib/libtorch.so.1)
frame #9: <unknown function> + 0x5c0e43 (0x7f8609946e43 in /usr/local/lib/python3.5/dist-packages/torch/lib/libtorch.so.1)
frame #10: torch::jit::InterpreterState::run(std::vector<c10::IValue, std::allocator<c10::IValue> >&) + 0x31 (0x7f8609942651 in /usr/local/lib/python3.5/dist-packages/torch/lib/libtorch.so.1)
frame #11: torch::jit::GraphExecutor::run(std::vector<c10::IValue, std::allocator<c10::IValue> >&) + 0x1c4 (0x7f8609924a74 in /usr/local/lib/python3.5/dist-packages/torch/lib/libtorch.so.1)
frame #12: ./MainRandomizeRnnDistance() [0x47fd12]
frame #13: ./MainRandomizeRnnDistance() [0x45da0d]
frame #14: __libc_start_main + 0xf0 (0x7f85ffef3830 in /lib/x86_64-linux-gnu/libc.so.6)
frame #15: ./MainRandomizeRnnDistance() [0x459fe9]

Expected behavior

Having no problem.

Environment

Collecting environment information...
PyTorch version: 1.0.0a0+5510196
Is debug build: No
CUDA used to build PyTorch: 9.0.176

OS: Ubuntu 16.04.5 LTS
GCC version: (GCC) 8.2.0
CMake version: version 3.12.0

Python version: 3.5
Is CUDA available: Yes
CUDA runtime version: 9.0.176
GPU models and configuration: GPU 0: GeForce GTX 1070
Nvidia driver version: 410.78
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.7.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_static_v7.a

Versions of relevant libraries:
[pip] Could not collect
[conda] Could not collect

Additional context

I think this can be solved by the following change

@torch.jit.script
def batch_sum(data, mask, dims):
    data = data * mask.type_as(data)
    for _ in range(dims.size(0)):
        data = data.sum(1)
    mask = torch.ones([data.size(0)], dtype=torch.uint8).type_as(mask)
    #mask = torch.ones([data.size(0)], dtype=torch.uint8)
    dims = dims[:0]  # empty tensor
    return data, mask, dims

[eellison]

Looks like a real issue. we'll look into this

[jamesr66a]

Hi @Strasser-Pablo, do you think you could provide a minimal repro script for this issue? I'd like to add it as a test case along with the fix

[Strasser-Pablo]

Hi, Sorry for the delay.

Here is a basic python script showing this issue and another similar issue with unsqueeze.

import torch
import torch.jit
import torch.jit.batchop

@torch.jit.batch()
def foo(a,b):
   return torch.sum(a).unsqueeze(1)+b

device = torch.device("cuda")

a=torch.randn(10,device=device)
b=torch.randn(10,device=device)
result=foo(a,b)

This will generate the following error:

Traceback (most recent call last):
  File "test_python.py", line 13, in <module>
    result=foo(a,b)
  File "/opt/conda/lib/python3.6/site-packages/torch/jit/__init__.py", line 755, in wrapper
    res = res_mod(*new_args)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 489, in __call__
    result = self.forward(*input, **kwargs)
RuntimeError: 
Expected object of backend CUDA but got backend CPU for sequence element 1 in sequence argument at position #1 'tensors' (checked_tensor_list_unwrap at /opt/conda/conda-bld/pytorch_1544199946412/work/aten/src/ATen/Utils.h:87)
frame #0: c10::Error::Error(c10::SourceLocation, std::string const&) + 0x45 (0x7ff602109cc5 in /opt/conda/lib/python3.6/site-packages/torch/lib/libc10.so)
frame #1: <unknown function> + 0x11f856f (0x7ff605a9856f in /opt/conda/lib/python3.6/site-packages/torch/lib/libcaffe2_gpu.so)
frame #2: at::CUDAByteType::_th_cat(c10::ArrayRef<at::Tensor>, long) const + 0xac (0x7ff605a9e33c in /opt/conda/lib/python3.6/site-packages/torch/lib/libcaffe2_gpu.so)
frame #3: at::native::cat(c10::ArrayRef<at::Tensor>, long) + 0xa4 (0x7ff602935464 in /opt/conda/lib/python3.6/site-packages/torch/lib/libcaffe2.so)
frame #4: at::TypeDefault::cat(c10::ArrayRef<at::Tensor>, long) const + 0x4f (0x7ff602b12cef in /opt/conda/lib/python3.6/site-packages/torch/lib/libcaffe2.so)
frame #5: torch::autograd::VariableType::cat(c10::ArrayRef<at::Tensor>, long) const + 0x1e8 (0x7ff5ff635898 in /opt/conda/lib/python3.6/site-packages/torch/lib/libtorch.so.1)
frame #6: <unknown function> + 0x57d818 (0x7ff5ff775818 in /opt/conda/lib/python3.6/site-packages/torch/lib/libtorch.so.1)
frame #7: <unknown function> + 0x680af8 (0x7ff5ff878af8 in /opt/conda/lib/python3.6/site-packages/torch/lib/libtorch.so.1)
frame #8: torch::jit::InterpreterState::run(std::vector<c10::IValue, std::allocator<c10::IValue> >&) + 0x22 (0x7ff5ff873e52 in /opt/conda/lib/python3.6/site-packages/torch/lib/libtorch.so.1)
frame #9: <unknown function> + 0x65e95c (0x7ff5ff85695c in /opt/conda/lib/python3.6/site-packages/torch/lib/libtorch.so.1)
frame #10: <unknown function> + 0x3ceff2 (0x7ff642b4cff2 in /opt/conda/lib/python3.6/site-packages/torch/lib/libtorch_python.so)
frame #11: <unknown function> + 0x3aba53 (0x7ff642b29a53 in /opt/conda/lib/python3.6/site-packages/torch/lib/libtorch_python.so)
frame #12: <unknown function> + 0x111e36 (0x7ff64288fe36 in /opt/conda/lib/python3.6/site-packages/torch/lib/libtorch_python.so)
<omitting python frames>
frame #37: __libc_start_main + 0xf0 (0x7ff653212830 in /lib/x86_64-linux-gnu/libc.so.6)
:
operation failed in interpreter:
@torch.jit.script
def batch_unsqueeze(data, mask, dims, dim_):
    if int(dim_) < 0:
        dim_ = dim_ + data.dim() + 1
    dim = int(dim_)
    # if dim == 0:
    #     raise ValueError("cannot do unsqueeze along batch_dim")
    data = data.unsqueeze(dim)
    mask = mask.unsqueeze(dim)
    dims = torch.cat((dims[:dim], torch.zeros([1], dtype=torch.uint8), dims[dim:dims.size(0)]))
           ~~~~~~~~~ <--- HERE
    return data, mask, dims

This can be solved by replacing the batch unsqueeze version with:

@torch.jit.script
def batch_unsqueeze(data, mask, dims, dim_):
    if int(dim_) < 0:
        dim_ = dim_ + data.dim() + 1
    dim = int(dim_)
    # if dim == 0:
    #     raise ValueError("cannot do unsqueeze along batch_dim")
    data = data.unsqueeze(dim)
    mask = mask.unsqueeze(dim)
    dims = torch.cat((dims[:dim], torch.zeros([1], dtype=torch.uint8).type_as(dims), dims[dim:dims.size(0)]))
    return data, mask, dims

Then the following error is left:

Traceback (most recent call last):
  File "test_python.py", line 13, in <module>
    result=foo(a,b)
  File "/opt/conda/lib/python3.6/site-packages/torch/jit/__init__.py", line 755, in wrapper
    res = res_mod(*new_args)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 489, in __call__
    result = self.forward(*input, **kwargs)
RuntimeError: 
expected type CPUByteType but got CUDAByteType (compute_types at /opt/conda/conda-bld/pytorch_1544199946412/work/aten/src/ATen/native/TensorIterator.cpp:134)
frame #0: c10::Error::Error(c10::SourceLocation, std::string const&) + 0x45 (0x7f4f3c2aacc5 in /opt/conda/lib/python3.6/site-packages/torch/lib/libc10.so)
frame #1: at::TensorIterator::compute_types() + 0x36b (0x7f4f3cac67eb in /opt/conda/lib/python3.6/site-packages/torch/lib/libcaffe2.so)
frame #2: at::TensorIterator::Builder::build() + 0x46 (0x7f4f3cac84e6 in /opt/conda/lib/python3.6/site-packages/torch/lib/libcaffe2.so)
frame #3: at::TensorIterator::binary_op(at::Tensor&, at::Tensor const&, at::Tensor const&) + 0x2c4 (0x7f4f3cac8e54 in /opt/conda/lib/python3.6/site-packages/torch/lib/libcaffe2.so)
frame #4: at::native::mul_out(at::Tensor&, at::Tensor const&, at::Tensor const&) + 0x9d (0x7f4f3c9aea5d in /opt/conda/lib/python3.6/site-packages/torch/lib/libcaffe2.so)
frame #5: at::native::mul(at::Tensor const&, at::Tensor const&) + 0x29 (0x7f4f3c9af2c9 in /opt/conda/lib/python3.6/site-packages/torch/lib/libcaffe2.so)
frame #6: at::TypeDefault::mul(at::Tensor const&, at::Tensor const&) const + 0x5d (0x7f4f3cc9663d in /opt/conda/lib/python3.6/site-packages/torch/lib/libcaffe2.so)
frame #7: torch::autograd::VariableType::mul(at::Tensor const&, at::Tensor const&) const + 0x1bf (0x7f4f3989bd8f in /opt/conda/lib/python3.6/site-packages/torch/lib/libtorch.so.1)
frame #8: <unknown function> + 0x5e6f2b (0x7f4f3997ff2b in /opt/conda/lib/python3.6/site-packages/torch/lib/libtorch.so.1)
frame #9: <unknown function> + 0x680af8 (0x7f4f39a19af8 in /opt/conda/lib/python3.6/site-packages/torch/lib/libtorch.so.1)
frame #10: torch::jit::InterpreterState::run(std::vector<c10::IValue, std::allocator<c10::IValue> >&) + 0x22 (0x7f4f39a14e52 in /opt/conda/lib/python3.6/site-packages/torch/lib/libtorch.so.1)
frame #11: <unknown function> + 0x65e95c (0x7f4f399f795c in /opt/conda/lib/python3.6/site-packages/torch/lib/libtorch.so.1)
frame #12: <unknown function> + 0x3ceff2 (0x7f4f7ccedff2 in /opt/conda/lib/python3.6/site-packages/torch/lib/libtorch_python.so)
frame #13: <unknown function> + 0x3aba53 (0x7f4f7cccaa53 in /opt/conda/lib/python3.6/site-packages/torch/lib/libtorch_python.so)
frame #14: <unknown function> + 0x111e36 (0x7f4f7ca30e36 in /opt/conda/lib/python3.6/site-packages/torch/lib/libtorch_python.so)
<omitting python frames>
frame #39: __libc_start_main + 0xf0 (0x7f4f8d3b3830 in /lib/x86_64-linux-gnu/libc.so.6)
:
operation failed in interpreter:
@torch.jit.script
def batch_add(data1, mask1, dims1, data2, mask2, dims2, alpha_):
    alpha = float(alpha_)
    data = torch.add(data1, data2, alpha=alpha)
    mask = mask1 * mask2
           ~~~~~~~~~~~~~ <--- HERE
    dims = dims1.__or__(dims2)
    return data, mask, dims

Which is finally solved with:

@torch.jit.script
def batch_sum(data, mask, dims):
    data = data * mask.type_as(data)
    for _ in range(dims.size(0)):
        data = data.sum(1)
    mask = torch.ones([data.size(0)], dtype=torch.uint8).type_as(mask)
    dims = dims[:0]  # empty tensor
    return data, mask, dims

I don't think this is the only instance of this problem. The problem mainly come when tensor are created inside the function and an output is generate which is no more on the gpu which will lead to the error to the next operator trying to use it.

I hope this information help you.