void GroupQuantizeDecompressOp::setup()

in optimum/graphcore/custom_ops/group_quantize_decompress/group_quantize_decompress.cpp [50:92]

• 33 lines of code
• 7 McCabe index (conditional complexity)

void GroupQuantizeDecompressOp::setup() {
  auto xInfo = inInfo(0);
  auto groupScaleInfo = inInfo(1);
  auto groupBiasInfo = inInfo(2);

  // check expected shapes
  if (xInfo.rank() != 3) {
    throw error("GroupQuantizeDecompressOp::setup x should have rank 3");
  }

  if (groupScaleInfo.rank() != xInfo.rank()) {
    throw error(
        "GroupQuantizeDecompressOp::setup groupScale should same rank as x");
  }

  if (groupBiasInfo.rank() != xInfo.rank()) {
    throw error(
        "GroupQuantizeDecompressOp::setup groupBias should same rank as x");
  }

  if (groupScaleInfo.shape()[2] != 1) {
    throw error("GroupQuantizeDecompressOp::setup groupScale shape at last "
                "dimension should be 1");
  }

  if (groupBiasInfo.shape()[2] != 1) {
    throw error("GroupQuantizeDecompressOp::setup groupBias shape at last "
                "dimension should be 1");
  }

  if (groupScaleInfo.shape() != groupScaleInfo.shape()) {
    throw error("GroupQuantizeDecompressOp::setup groupScale and groupBias "
                "should have same shape");
  }

  auto nRows = xInfo.shape()[0];
  auto nGroups = xInfo.shape()[1];
  auto nIds = xInfo.shape()[2];

  // x decompressed
  outInfo(0) = TensorInfo(groupScaleInfo.data_type(),
                          Shape{nRows, nGroups * nIds * 4});
}