pymoo.optimize.minimize giving outputs of different dimensions every time

Question

I have a following optimization code to optimize a function that calculates in arrays:

import numpy as np
from pymoo.algorithms.moo.nsga2 import NSGA2
from pymoo.factory import get_problem
from pymoo.optimize import minimize
from scipy.optimize import minimize as min_scipy
from pymoo.core.problem import Problem
from pymoo.visualization.scatter import Scatter

problem = get_problem("zdt1")


gamma = 0.5        # MC efficiency
alpha = 0.5        # Price Sensitivity
beta = 0.5         # Coefficient for 
M = 7              # Market Size
N = 3              # N Competitors
Wmax = 5           # Upper Ceiling of discretized technology space

My function is the following:

def profit_array(x, w_j, w_mj, gamma, M, N):
    P = np.reshape(x, (Wmax, Wmax))
    numer = np.exp(-alpha * np.array(P) + beta * np.array(w_j))
    denom = 1 + np.array(N) * np.exp(-alpha * np.array(P) + beta * np.array(w_mj)) + numer
    share = np.divide(np.array(numer), np.array(denom))
    mc = np.exp(-(gamma*np.array(w_j)))
    pi = (np.array(P) - np.array(mc))*np.array(share) * np.array(M)
    negpi = np.reshape(-pi, (1, Wmax*Wmax))
    return negpi


class ProfitMax2(Problem):
    def __init__(self, gamma, M, N, Wmax):
        self.gamma = gamma
        self.M = M
        self.N = N
        self.Wmax = Wmax
        self.n_var = self.Wmax ** 2  # Total number of decision variables
        
        self.Wmat = np.array([[1]*self.Wmax])
        for i in range(2, self.Wmax+1):
            add = np.array([[i] * self.Wmax])
            self.Wmat = np.concatenate([self.Wmat, add])

        self.Wtrans = np.transpose(self.Wmat)

        # Define bounds for decision variables (P)
        xl = np.zeros(self.n_var)  # Lower bound for all elements of P
        xu = np.ones(self.n_var) * 10  # Upper bound for all elements of P

        super().__init__(n_var=self.n_var, n_obj=1, n_constr=0, xl=xl, xu=xu)

    def _evaluate(self, x, out, *args, **kwargs):
        # Reshape x into a 5x5 matrix (P)
        # Calculate profit using P, Wmat, and Wtrans
        profit = profit_array(x, self.Wmat, self.Wtrans, self.gamma, self.M, self.N)

        # Set the objective value (minimize negative profit)
        out["F"] = profit

# Create a problem with Wmax
testprob = ProfitMax2(gamma=0.5, M=7, N=3, Wmax=5)
algorithm = NSGA2()

# Generate an initial x array with the correct number of elements
x_init = np.asmatrix(np.random.uniform(0, 10, testprob.n_var))  # Generate random values within bounds

res = minimize(testprob, algorithm, termination=('n_gen', 100), x0=x_init)

When I run the code, it gives me an error of "cannot reshape array of size 2500 into shape (5,5)", which is weird since my initial x is an array of 25. I manually ran the profit_array function by directly using x_init with Wmat and Wtrans generated separately, and the function works. So I have no idea why it suddenly bloats to size 2500 when put into an algorithm.

One thing I noticed that could be related to the problem: I tested a scalar version of my optimization just to check, which is the following:

def profit_sub_scalar(P, w_j, w_mj, gamma, M, N):
    numer = np.exp(-alpha * P + beta * w_j)
    denom = 1 + N * np.exp(-alpha * P + beta * w_mj) + numer
    share = numer / denom
    mc = np.exp(-gamma * w_j)
    pi = (P - mc) * share * M
    negpi = -pi
    return negpi

class ProfitMax_Scalar(Problem):
    def __init__(self, gamma, M, N, w_j, w_mj):
        self.gamma = gamma
        self.M = M
        self.N = N
        self.w_j = w_j
        self.w_mj = w_mj
        self.n_var = 1  # Total number of decision variables

        # Define bounds for decision variables (P)
        xl = 0
        xu = 10
        
        super().__init__(n_var=self.n_var, n_obj=1, n_constr=0, xl=xl, xu = xu)

    def _evaluate(self, x, out, *args, **kwargs):
        # Calculate profit using P, Wmat, and Wtrans
        profit = profit_sub_scalar(x, self.w_j, self.w_mj, self.gamma, self.M, self.N)

        # Set the objective value (minimize negative profit)
        out["F"] = profit

# Create a problem with Wmax
testprob = ProfitMax_Scalar(gamma=0.5, M=7, N=3, w_j = 3, w_mj = 4)
algorithm = NSGA2()

x_init = 0

res = minimize(testprob, algorithm, termination=('n_gen', 100), x0=x_init)

Here the scalar code runs fine, but when I check the result (res), the outputs are arrays of different size every time I run the code! Sometimes it's an array of 4 (res.X = [5.5486, 5.5486, 5.5486, 5.5486]), and other times it's an array of 40 with each element the same value (5.5486). I believe that this is related to why my code is generating error when calculated in arrays, but I have no idea why the pymoo package does this and how I should edit my code to fix this weirdness.

Any kind of help is much appreciated. Thank you!