mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
91797c0672
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/10946 ``` codemod -d . --extensions cc,cpp,cu,cuh,h caffe2/proto/caffe2.pb.h caffe2/proto/caffe2_pb.h ``` Reviewed By: houseroad Differential Revision: D9539945 fbshipit-source-id: 497d04720e8e7e61c05ffe1b23733d0cb774de7e
41 lines
823 B
C++
41 lines
823 B
C++
#include "caffe2/transforms/single_op_transform.h"
|
|
|
|
#include "caffe2/core/common.h"
|
|
#include "caffe2/core/logging.h"
|
|
#include "caffe2/core/net.h"
|
|
#include "caffe2/proto/caffe2_pb.h"
|
|
|
|
namespace caffe2 {
|
|
|
|
using transform::Graph;
|
|
|
|
bool SingleOpTransform::PatternRule(
|
|
const Graph& g,
|
|
const std::vector<int>& subgraph,
|
|
int idx) {
|
|
if (subgraph.size() == 0) {
|
|
return MatchOperator(g.node(idx).op);
|
|
}
|
|
return false;
|
|
}
|
|
|
|
bool SingleOpTransform::ValidatorRule(
|
|
const Graph& /*g*/,
|
|
const std::vector<int>& subgraph) {
|
|
if (subgraph.size() == 1) {
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
bool SingleOpTransform::ReplaceRule(
|
|
const std::vector<int>& subgraph,
|
|
Graph* g_ptr) {
|
|
CHECK(g_ptr);
|
|
auto& g = *g_ptr;
|
|
ReplaceOperator(&(g.node(subgraph[0]).op));
|
|
return true;
|
|
}
|
|
|
|
} // namespace caffe2
|