MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT / getOutputDimensions

Method getOutputDimensions

plugin/flattenConcat/flattenConcat.cpp:90–130  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

88}
89
90Dims FlattenConcat::getOutputDimensions(int32_t index, Dims const* inputs, int32_t nbInputDims) noexcept
91{
92 try
93 {
94 PLUGIN_ASSERT(nbInputDims >= 1);
95 PLUGIN_ASSERT(index == 0);
96
97 mNumInputs = nbInputDims;
98 mCopySize.resize(mNumInputs);
99 mInputConcatAxis.resize(mNumInputs);
100 int32_t outputConcatAxis = 0;
101
102 for (int32_t i = 0; i < nbInputDims; ++i)
103 {
104 int32_t flattenInput = 0;
105 PLUGIN_ASSERT(inputs[i].nbDims == 3);
106 if (mConcatAxisID != 1)
107 {
108 PLUGIN_ASSERT(inputs[i].d[0] == inputs[0].d[0]);
109 }
110 if (mConcatAxisID != 2)
111 {
112 PLUGIN_ASSERT(inputs[i].d[1] == inputs[0].d[1]);
113 }
114 if (mConcatAxisID != 3)
115 {
116 PLUGIN_ASSERT(inputs[i].d[2] == inputs[0].d[2]);
117 }
118 flattenInput = inputs[i].d[0] * inputs[i].d[1] * inputs[i].d[2];
119 outputConcatAxis += flattenInput;
120 }
121
122 return Dims3(mConcatAxisID == 1 ? outputConcatAxis : 1, mConcatAxisID == 2 ? outputConcatAxis : 1,
123 mConcatAxisID == 3 ? outputConcatAxis : 1);
124 }
125 catch (std::exception const& e)
126 {
127 caughtError(e);
128 }
129 return Dims{};
130}
131
132int32_t FlattenConcat::initialize() noexcept
133{

Callers

nothing calls this directly

Calls 3

Dims3Class · 0.85
caughtErrorFunction · 0.85
resizeMethod · 0.45

Tested by

no test coverage detected