-#include <c10/util/Exception.h>
#include <torch/csrc/jit/passes/onnx/peephole.h>
+#include <c10/util/Exception.h>
#include <c10/util/Optional.h>
// Determine whether `from` can broadcast to `to`, and if so at which
// position. `from` must be a suffix of `to`, except that any
// occurences of 1 in `from` are treated as wildcards.
-c10::optional<size_t> fusibleExpandTo(at::IntArrayRef from, at::IntArrayRef to) {
+c10::optional<size_t> fusibleExpandTo(
+ at::IntArrayRef from,
+ at::IntArrayRef to) {
if (from.size() > to.size()) {
return c10::nullopt;
}
// always have this information (because expands are only ever traced,
// not generated from symbolic), but if for some reason we don't
// have it, we need to skip.
- if (!unexpanded_input->isTensor() || !n->output()->isTensor())
+ if (!unexpanded_input->isCompleteTensor() ||
+ !n->output()->isCompleteTensor())
continue;
// Not all broadcasts are supported by ONNX broadcast.
Node* gather_indices = graph->create(onnx::Constant, 1);
gather_indices->insertBefore(n);
- gather_indices->t_(attr::value, autograd::make_variable(at::scalar_to_tensor(at::Scalar(1))));
+ gather_indices->t_(
+ attr::value,
+ autograd::make_variable(at::scalar_to_tensor(at::Scalar(1))));
Node* batch_size = graph->create(onnx::Gather, 1);
batch_size->insertBefore(n);
}
}
-static void fuseSplitListUnpack(Block *b) {
- for(auto it = b->nodes().begin(), end = b->nodes().end(); it != end; ++it) {
- for (auto *child_block : it->blocks()) {
+static void fuseSplitListUnpack(Block* b) {
+ for (auto it = b->nodes().begin(), end = b->nodes().end(); it != end; ++it) {
+ for (auto* child_block : it->blocks()) {
fuseSplitListUnpack(child_block);
}
- if (it->kind() == prim::ListUnpack && it->input()->node()->kind() == onnx::Split) {
+ if (it->kind() == prim::ListUnpack &&
+ it->input()->node()->kind() == onnx::Split) {
auto origSplitNode = it->input()->node();
- Node * splitNode = b->owningGraph()->create(onnx::Split, it->outputs().size());
- for (size_t i=0; i<splitNode->outputs().size(); ++i) {
+ Node* splitNode =
+ b->owningGraph()->create(onnx::Split, it->outputs().size());
+ for (size_t i = 0; i < splitNode->outputs().size(); ++i) {
splitNode->outputs()[i]->copyMetadata(it->outputs()[i]);
}
splitNode->copyAttributes(*origSplitNode);
}
void removeMaxPoolUnusedOutput(Block* b) {
- for (auto it = b->nodes().begin(), end = b->nodes().end(); it != end; ++it) {
+ for (auto it = b->nodes().begin(), end = b->nodes().end(); it != end; ++it) {
auto n = *it;
for (auto* child_block : n->blocks()) {
removeMaxPoolUnusedOutput(child_block);
return node;
})
.VS(copyMetadata)
- .VS(isTensor)
+ .VS(isCompleteTensor)
.VS(requires_grad)
.def("toIValue", [](Value& n) { return toIValue(&n); })
.def("type", [](Value& v) { return v.type(); });
[](Block& b) {
return py::make_iterator(b.outputs().begin(), b.outputs().end());
})
- .def(
- "returnNode",
- [](Block& b) {
- return b.return_node();
- })
- .def(
- "paramNode",
- [](Block& b) {
- return b.param_node();
- });
+ .def("returnNode", [](Block& b) { return b.return_node(); })
+ .def("paramNode", [](Block& b) { return b.param_node(); });
#define NS(name) def(#name, &Node ::name)
py::class_<Node, std::unique_ptr<Node, py::nodelete>>(m, "Node")
if _is_value(size):
shape = size
else:
- if self.isTensor():
+ if self.isCompleteTensor():
self_sizes = self.type().sizes()
if self_sizes and len(size) == 2 and self_sizes[0] == size[0]:
return g.op("Flatten", self, axis_i=1)
def type_as(g, self, other):
- if self.isTensor() and other.isTensor() and self.type().scalarType() == other.type().scalarType():
+ if self.isCompleteTensor() and other.isCompleteTensor() and self.type().scalarType() == other.type().scalarType():
return self
- if other.isTensor():
+ if other.isCompleteTensor():
other_type_name = other.type().scalarType()
return g.op("Cast", self, to_i=cast_pytorch_to_onnx[other_type_name])
else:
repeats = g.op("Constant", value_t=torch.LongTensor(repeats))
const_repeats = _maybe_get_const(repeats, 'is')
- if self.isTensor() and not _is_value(const_repeats):
+ if self.isCompleteTensor() and not _is_value(const_repeats):
sizes = self.type().sizes()
diff_dims = len(const_repeats) - len(sizes)
if diff_dims > 0: