mlir::FunctionType func_type = main_fn.getFunctionType();
  for (auto input_type : func_type.getInputs()) {
    tensorflow::TensorShape tensor_shape;
    xla::Shape xla_shape = xla::TypeToShape(input_type);
    TF_RETURN_IF_ERROR(tensorflow::TensorShape::BuildTensorShape(
        xla_shape.dimensions(), &tensor_shape));
    arg_shapes.emplace_back(tensor_shape);

    DataType dtype;
    TF_RETURN_IF_ERROR(ConvertToDataType(input_type, &dtype));

    Value result = reduction.getResult(0);

    // The mean op needs to divide by the product of the reduced dimensions.
    if (std::is_same<OpTy, TF::MeanOp>::value) {
      Value in_shape = rewriter.create<shape::ShapeOfOp>(loc, op.getInput());
      Value divisor_count = rewriter.create<arith::ConstantIndexOp>(loc, 1);
      for (size_t i = 0; i < input_shape.size(); ++i) {

    Arg<TF_Int64Tensor, [{2-D.  The `indices` of the second `SparseTensor`, size `[nnz, ndims]` Matrix.}]>:$b_indices,
    Arg<TF_NumberTensor, [{1-D.  The `values` of the second `SparseTensor`, size `[nnz]` Vector.}]>:$b_values,
    Arg<TF_Int64Tensor, [{1-D.  The `shape` of the second `SparseTensor`, size `[ndims]` Vector.}]>:$b_shape,