ONE - On-device Neural Engine
Loading...
Searching...
No Matches
moco::ShapeGraphBuilder Class Referencefinal

GraphBuilder for Shape node. More...

#include <Shape.h>

Collaboration diagram for moco::ShapeGraphBuilder:

Public Member Functions

bool validate (const tensorflow::NodeDef &) const override
 
void build (const tensorflow::NodeDef &, GraphBuilderContext *) const override
 
- Public Member Functions inherited from moco::GraphBuilder
virtual ~GraphBuilder ()
 

Detailed Description

GraphBuilder for Shape node.

Definition at line 28 of file Shape.h.

Member Function Documentation

◆ build()

void moco::ShapeGraphBuilder::build ( const tensorflow::NodeDef &  node,
GraphBuilderContext context 
) const
overridevirtual

Implements moco::GraphBuilder.

Definition at line 68 of file Shape.cpp.

69{
70 assert(context != nullptr);
71
72 loco::Graph *graph = context->graph();
73 SymbolTable *tensor_names = context->tensor_names();
74 UpdateQueue *updates = context->updates();
75
76 // create TF dialect Shape node
77 auto tf_shape = graph->nodes()->create<TFShape>();
78 tf_shape->name(node.name());
79
80 if (plier::tf::has_attrs(node, {"out_type"}))
81 {
83 // TODO Support other dtype like S64
84 assert(dtype == loco::DataType::S32);
85
86 tf_shape->dtype(dtype);
87 }
88 else
89 {
90 // Set to S32, TF-documented default value for 'out_type'
91 tf_shape->dtype(loco::DataType::S32);
92 }
93
94 TensorName output_name(node.name(), 0);
95 tensor_names->enroll(output_name, tf_shape);
96
97 auto update = std::make_unique<ShapeGraphUpdate>(tf_shape, TensorName(node.input(0)));
98 updates->enroll(std::move(update));
99}
A neural network graph.
Definition Graph.h:161
Class to store and query loco::Node* with string name key.
void enroll(const TensorName &tensor_name, loco::Node *node)
Registers a name with corresponding loco::Node *.
Class to store GraphUpdate objects.
void enroll(std::unique_ptr< GraphUpdate > &&update)
Registers GraphUpdate objects.
FeatureShapeUpdater update(loco::FeatureShape &feature_shape)
bool has_attrs(const tensorflow::NodeDef &node, const std::vector< std::string > &attr_names)
Definition Convert.cpp:35
tensorflow::DataType get_datatype_attr(const tensorflow::NodeDef &node, const std::string &attr_name)
Definition Convert.cpp:43
loco::DataType as_loco_datatype(const tensorflow::DataType dtype)
Definition Convert.cpp:123
NodeName name(void) const
Definition TFNodeDecl.h:50

References plier::tf::as_loco_datatype(), moco::SymbolTable::enroll(), moco::UpdateQueue::enroll(), plier::tf::get_datatype_attr(), moco::GraphBuilderContext::graph(), plier::tf::has_attrs(), moco::TFNode::name(), moco::GraphBuilderContext::tensor_names(), moco::update(), and moco::GraphBuilderContext::updates().

◆ validate()

bool moco::ShapeGraphBuilder::validate ( const tensorflow::NodeDef &  node) const
overridevirtual

Implements moco::GraphBuilder.

Definition at line 60 of file Shape.cpp.

61{
62 if (node.input_size() != 1)
63 return false;
64
65 return plier::tf::has_attrs(node, {"T"});
66}

References plier::tf::has_attrs().


The documentation for this class was generated from the following files: