Status BlocksparseConvShape()

in src/blocksparse_conv_op.cc [27:155]
115 lines of code
14 McCabe index (conditional complexity)

Status BlocksparseConvShape(InferenceContext* ctx)
{
  int32 mode, C, K;
  std::vector<int32> DHW, MPQ, dimF;

  TF_RETURN_IF_ERROR(ctx->GetAttr("mode",   &mode));
  TF_RETURN_IF_ERROR(ctx->GetAttr("C",      &C));
  TF_RETURN_IF_ERROR(ctx->GetAttr("K",      &K));
  TF_RETURN_IF_ERROR(ctx->GetAttr("DHW",    &DHW));
  TF_RETURN_IF_ERROR(ctx->GetAttr("MPQ",    &MPQ));
  TF_RETURN_IF_ERROR(ctx->GetAttr("dimF",   &dimF));

  // fprop: NKMPQ = KCTRS . NCDHW
  if (mode == 0)
  {
    // a=F, b=I, c=O
    ShapeHandle I;
    DimensionHandle dh;
    TF_RETURN_IF_ERROR(ctx->WithRankAtLeast(ctx->input(7), 3, &I));
    TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(I, 1), C, &dh));

    int32 rankI = ctx->Rank(I);
    if (rankI == 5)
    {
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(I, 2), DHW[0], &dh));
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(I, 3), DHW[1], &dh));
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(I, 4), DHW[2], &dh));
      ctx->set_output(0, ctx->MakeShape({
        ctx->Dim(I, 0), ctx->MakeDim(K), ctx->MakeDim(MPQ[0]), ctx->MakeDim(MPQ[1]), ctx->MakeDim(MPQ[2]) // NKMPQ
      }));
    }
    else if (rankI == 4)
    {
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(I, 2), DHW[1], &dh));
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(I, 3), DHW[2], &dh));
      ctx->set_output(0, ctx->MakeShape({
        ctx->Dim(I, 0), ctx->MakeDim(K), ctx->MakeDim(MPQ[1]), ctx->MakeDim(MPQ[2]) // NKPQ
      }));
    }
    else if (rankI == 3)
    {
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(I, 2), DHW[2], &dh));
      ctx->set_output(0, ctx->MakeShape({
        ctx->Dim(I, 0), ctx->MakeDim(K), ctx->MakeDim(MPQ[2])  // NKQ
      }));
    }
    else
      return errors::InvalidArgument("BlocksparseConv requires an input rank between 3 and 5: ", rankI);
  }
  // bprop: NCDHW = KCTRS . NKMPQ
  else if (mode == 1)
  {
    // a=F, b=I, c=O
    ShapeHandle I;
    DimensionHandle dh;
    TF_RETURN_IF_ERROR(ctx->WithRankAtLeast(ctx->input(7), 3, &I));
    TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(I, 1), K, &dh));

    int32 rankI = ctx->Rank(I);
    if (rankI == 5)
    {
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(I, 2), MPQ[0], &dh));
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(I, 3), MPQ[1], &dh));
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(I, 4), MPQ[2], &dh));
      ctx->set_output(0, ctx->MakeShape({
        ctx->Dim(I, 0), ctx->MakeDim(C), ctx->MakeDim(DHW[0]), ctx->MakeDim(DHW[1]), ctx->MakeDim(DHW[2]) // NCDHW
      }));
    }
    else if (rankI == 4)
    {
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(I, 2), MPQ[1], &dh));
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(I, 3), MPQ[2], &dh));
      ctx->set_output(0, ctx->MakeShape({
        ctx->Dim(I, 0), ctx->MakeDim(C), ctx->MakeDim(DHW[1]), ctx->MakeDim(DHW[2]) // NCHW
      }));
    }
    else if (rankI == 3)
    {
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(I, 2), MPQ[2], &dh));
      ctx->set_output(0, ctx->MakeShape({
        ctx->Dim(I, 0), ctx->MakeDim(C), ctx->MakeDim(DHW[2]) // NCW
      }));
    }
    else
      return errors::InvalidArgument("BlocksparseConv requires an input rank between 3 and 5: ", rankI);
  }
  // updat: KCTRS = NKMPQ . NCDHW
  else if (mode == 2)
  {
    // a=E, b=I, c=U
    ShapeHandle E, I;
    DimensionHandle dh;
    TF_RETURN_IF_ERROR(ctx->WithRankAtLeast(ctx->input(6), 3, &E));
    int32 rankI = ctx->Rank(E);
    TF_RETURN_IF_ERROR(ctx->WithRank(ctx->input(7), rankI, &I));

    if (rankI == 5)
    {
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(E, 2), MPQ[0], &dh));
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(E, 3), MPQ[1], &dh));
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(E, 4), MPQ[2], &dh));
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(I, 2), DHW[0], &dh));
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(I, 3), DHW[1], &dh));
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(I, 4), DHW[2], &dh));
    }
    else if (rankI == 4)
    {
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(E, 2), MPQ[1], &dh));
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(E, 3), MPQ[2], &dh));
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(I, 2), DHW[1], &dh));
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(I, 3), DHW[2], &dh));
    }
    else if (rankI == 3)
    {
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(E, 2), MPQ[2], &dh));
      TF_RETURN_IF_ERROR(ctx->WithValue(ctx->Dim(I, 2), DHW[2], &dh));
    }
    else
      return errors::InvalidArgument("BlocksparseConv requires an input rank between 3 and 5: ", rankI);


    std::vector<DimensionHandle> dims;
    for (std::vector<int32>::iterator it = dimF.begin() ; it != dimF.end(); ++it)
      dims.emplace_back(ctx->MakeDim(*it));

    ctx->set_output(0, ctx->MakeShape( dims ));
  }
  return Status::OK();
}