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)))).
//