[JIT] Add BatchTensor class #8922
Conversation
ChunliF
requested review from
apaszke,
colesbury,
ezyang,
gchanan,
soumith and
zdevito
as code owners
There was a problem hiding this comment.
@ezyang has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
| this->dims = dims; | ||
| } | ||
|
|
||
| BatchTensor::BatchTensor(std::vector<at::Tensor> datalist, std::vector<bool> dims) { |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| data_item += datalist[i]; | ||
| mask_item.fill_(1); | ||
| } | ||
| this->dims = dims; |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| } | ||
| this->data = data; | ||
| this->mask = mask; | ||
| this->dims = dims; |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| this->dims = dims; | ||
| } | ||
|
|
||
| std::vector<at::Tensor> BatchTensor::examples() { |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
There was a problem hiding this comment.
@ezyang has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
| .def(py::init<at::Tensor, at::Tensor, at::Tensor>()) | ||
| .def(py::init<std::vector<at::Tensor>, at::Tensor>()) | ||
| .def("examples", &BatchTensor::examples) | ||
| .def("get_data", &BatchTensor::getData) |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
test/test_jit.py
Outdated
| else: | ||
| super(TestBatched, self).assertEqual(x, y, prec, message, allow_inf) | ||
|
|
||
| def mb_rand(self, *dims): |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
test/test_jit.py
Outdated
|
|
||
|
|
||
| class TestBatched(TestCase): | ||
| def assertEqual(self, xs, ybs, prec=None, message='', allow_inf=False): |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| } | ||
|
|
||
| public: | ||
| at::Tensor data; |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
There was a problem hiding this comment.
@ezyang has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
|
@zdevito Just want to confirm that we agreed that we won't try to make BatchTensor wind it's way through all of the existing machinery? |
|
@ezyang -- yes we are going to make it a separate type in python, and represent it as a tuple in the script. In the future when we can define new types more easily we can make it more integrated but it is far to complicated to do it correct right now. |
There was a problem hiding this comment.
@ChunliF has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
|
Sorry to send another email to everyone's inbox. Are we decided that the Python |
|
@ssnl good point, it most definitely should be hidden. |
Summary: Add BatchTensor class - construct from data, mask, dims or construct from list of tensors - can return a list of tensors from an BatchTensor class next step: do IR level transformation and operators Closes pytorch#8922 Differential Revision: D8668986 Pulled By: ChunliF fbshipit-source-id: 8b24d2a9f46a3b42dbb397e99e9e059dfb2b326e
7 participants
Add BatchTensor class
next step: do IR level transformation and operators