out[j] = node_map_->GetNode(const_name);
if (out[j] == nullptr) {
out[j] = graph_->add_node();
- *out[j] = CreateNodeDef(const_name, TensorValue(&value));
+ TF_RETURN_IF_ERROR(
+ CreateNodeDef(const_name, TensorValue(&value), out[j]));
out[j]->set_device(node.device());
node_map_->AddNode(const_name, out[j]);
string ctrl_dep =
value.vec<int64>()(i) = i;
}
}
- *reduction_indices = CreateNodeDef(const_name, TensorValue(&value));
+ TF_RETURN_IF_ERROR(
+ CreateNodeDef(const_name, TensorValue(&value), reduction_indices));
reduction_indices->set_device(node->device());
string ctrl_dep =
AddControlDependency(node->input(1), graph_, node_map_.get());
} // namespace
// static
-NodeDef ConstantFolding::CreateNodeDef(const string& name,
- const TensorValue& tensor) {
- NodeDef node;
- node.set_name(name);
- node.set_op("Const");
+Status ConstantFolding::CreateNodeDef(const string& name,
+ const TensorValue& tensor,
+ NodeDef* node) {
+ node->set_name(name);
+ node->set_op("Const");
AttrValue attr_type;
attr_type.set_type(tensor->dtype());
- node.mutable_attr()->insert({"dtype", attr_type});
+ node->mutable_attr()->insert({"dtype", attr_type});
AttrValue attr_tensor;
TensorProto* t = attr_tensor.mutable_tensor();
bool optimized = false;
+ size_t encoded_size;
// Use the packed representation whenever possible to avoid generating large
// graphdefs. Moreover, avoid repeating the last values if they're equal.
if (tensor->NumElements() > 4) {
} \
if (last_index < kint32max) { \
optimized = true; \
+ encoded_size = (last_index + 1) * sizeof(NAME); \
t->mutable_##NAME##_val()->Reserve(last_index + 1); \
t->mutable_##NAME##_val()->AddNAlreadyReserved(last_index + 1); \
val_ptr = tensor->flat<TYPE>().data(); \
tensor->shape().AsProto(t->mutable_tensor_shape());
} else {
tensor->AsProtoTensorContent(t);
+ encoded_size = t->tensor_content().size();
+ }
+ node->mutable_attr()->insert({"value", attr_tensor});
+
+ if (encoded_size < 10 * 1024 * 1024) {
+ return Status::OK();
}
- node.mutable_attr()->insert({"value", attr_tensor});
- return node;
+ return errors::InvalidArgument(
+ strings::StrCat("Can't fold ", name, ", its size would be too large"));
}
Status ConstantFolding::EvaluateNode(const NodeDef& node,
return Status(error::INVALID_ARGUMENT, "Expected at least one output.");
}
+ outputs->resize(output_tensors.size());
for (size_t i = 0; i < output_tensors.size(); i++) {
string node_name = OptimizedNodeName(node, "-folded");
if (output_tensors.size() > 1) {
node_name = strings::StrCat(node_name, "-", i);
}
if (output_tensors[i].tensor) {
- outputs->push_back(CreateNodeDef(node_name, output_tensors[i]));
+ TF_RETURN_IF_ERROR(
+ CreateNodeDef(node_name, output_tensors[i], &outputs->at(i)));
} else {
// Create an empty NodeDef to identify dead outputs (e.g. the output of a
// switch that's not selected by the switch predicate).
- outputs->push_back(NodeDef());
+ outputs->at(i) = NodeDef();
}
}
return Status::OK();