auto input_type = operand0.getType().cast<ShapedType>();

  // Check axis bounds.
  if (input_type.hasRank()) {
    int32_t axis_value = op.getAxis();
    if (axis_value < 0) axis_value += input_type.getRank() + 1;
    if (axis_value < 0 || axis_value >= input_type.getRank() + 1)
      return op.emitOpError()
             << "op attribute 'axis' should be in range [-rank - 1, rank + 1), "

  if (!matchPattern(mean_op.getAxis(), m_Constant(&axis_const)))
    return failure();
  if (axis_const.size() != 2) return failure();
  auto axis_values = axis_const.getValues<APInt>();
  int i = 1;
  for (auto axis_value : axis_values) {
    if (axis_value != i++) return failure();
  }

  auto output = mean_op.getOutput();
  auto output_type = mlir::dyn_cast_or_null<RankedTensorType>(output.getType());

        prev_input_ty = input_ty;
      }
      expanded_inputs.push_back(
          rewriter.create<ExpandDimsOp>(loc, inferred_ty, input, axis_value));
    }

    rewriter.replaceOpWithNewOp<ConcatV2Op>(op, op.getType(), expanded_inputs,
                                            axis_value);
    return success();
  }
};

// Lowers SpaceToBatchND by reducing to reshape(transpose(reshape(pad(input)))).
//

        given:
        file("src/main/cpp/main.cpp") << """
            #include "foo.h"

            int main(int argc, char** argv) {
                return EXIT_VALUE;
            }
        """

        and:
        buildFile << """
            apply plugin: 'cpp-application'

            def headerDirectory = objects.directoryProperty()

  auto shape = llvm::to_vector<2>(indices_ty.getShape());
  if (axis_val == -1) axis_val = shape.size();

  int64_t depth_val = ShapedType::kDynamic;
  DenseIntElementsAttr depth_attr;
  if (matchPattern(depth, m_Constant(&depth_attr)) &&
      depth_attr.getNumElements() == 1)
    depth_val = (*depth_attr.begin()).getSExtValue();
  shape.insert(shape.begin() + axis_val, depth_val);