return InferShapeForCaseRegion(case_region);

  if (auto while_op = dyn_cast<WhileOp>(op))
    return InferShapeForWhile(
        while_op, while_op.body_function().getFunctionType().getResults());

  if (auto while_region = dyn_cast<WhileRegionOp>(op))
    return InferShapeForWhile(
        while_region,
        while_region.getBody().front().getTerminator()->getOperandTypes());