loc, tensorflow::GetTypeFromTFTensorShape({}, shape_dtype), size,
        input_size);
    auto if_cond = rewriter.create<TF::GreaterOp>(
        loc, tensorflow::GetTypeFromTFTensorShape({}, rewriter.getI1Type()),
        size_diff, scalar_zero);

    // Build the argument/result types for if branch function.
    auto input_shape = rewriter.create<TF::ShapeOp>(

      cond = b.create<mhlo::DynamicBroadcastInDimOp>(
          tensorflow::GetTypeFromTFTensorShape(result_type.getShape(),
                                               b.getI1Type()),
          cond, result_extents,
          GetI64ElementsAttrForSeq(0, cond_type.getRank(), &b));
    }
    Value select = b.create<mhlo::SelectOp>(