frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

broadcast op: it is an in-place op with both input and output set

Browse Source
This commit is contained in:
Yangqing Jia
2015-07-19 20:15:10 -07:00
parent 745d8ed969
commit 90affff039
2 changed files with 8 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
caffe2/mpi/broadcast_op.cc
View File

@ -12,7 +12,10 @@ class BroadcastOp final : public Operator<dtype, DeviceContext> {
USE_OPERATOR_BASE_FUNCTIONS;
BroadcastOp(const OperatorDef& operator_def, Workspace* ws)
: Operator<dtype, DeviceContext>(operator_def, ws),
root_(OperatorBase::template GetSingleArgument<int>("root", 0)) {}
root_(OperatorBase::template GetSingleArgument<int>("root", 0)) {
CHECK_EQ(operator_def.input(0), operator_def.output(0))
<< "Broadcast is an in-place operator.";
}
~BroadcastOp() {}
bool RunOnDevice() {
@ -25,8 +28,9 @@ class BroadcastOp final : public Operator<dtype, DeviceContext> {
protected:
int root_;
// Output: X. Note that X must have been initialized on root.
INPUT_OUTPUT_STATS(0, 0, 1, 1);
// Input: X. Output: X.
// Note that Broadcast works in-place by definition.
INPUT_OUTPUT_STATS(1, 1, 1, 1);
DISABLE_COPY_AND_ASSIGN(BroadcastOp);
};

1
caffe2/mpi/mpi_test.cc
View File

@ -23,6 +23,7 @@ const char kBcastNet[] =
" }"
" }"
" op {"
" input: \"X\""
" output: \"X\""
" type: \"Broadcast\""
" arg {"