| 148 | } |
| 149 | |
| 150 | DimsExprs QKVToContextPluginDynamic::getOutputDimensions( |
| 151 | int32_t outputIndex, DimsExprs const* inputs, int32_t /*nbInputs*/, IExprBuilder& exprBuilder) noexcept |
| 152 | { |
| 153 | // Input is BxSx3*N*H, output should be BxSxN*H |
| 154 | PLUGIN_ASSERT(outputIndex == 0); |
| 155 | // Copy over everything |
| 156 | DimsExprs output(inputs[IIDX]); |
| 157 | // Divide last dim by three |
| 158 | auto const* three = exprBuilder.constant(3); |
| 159 | output.d[HDIM] = exprBuilder.operation(DimensionOperation::kFLOOR_DIV, *inputs[IIDX].d[HDIM], *three); |
| 160 | return output; |
| 161 | } |
| 162 | bool QKVToContextPluginDynamic::supportsFormatCombination( |
| 163 | int32_t pos, PluginTensorDesc const* inOut, int32_t nbInputs, int32_t /*nbOutputs*/) noexcept |
| 164 | { |