Utilizing scipy.optimize.minimize with multiple variables of different shapes

Question

I am curious is there is a straightforward method for utilizing scipy.optimize.minimize with multiple variables that take different shapes. For example, let's take a look at a matrix decomposition problem.

I apologize, but I will be using latex here in the hope that one day SO will implement it.

We can deconstruct the matrix $ A_{n \times m} $ into two matrices $ W_{k \times n} $ and $ H_{k \times m} s.t. A \approx W^TH $

There are numerous methods for solving for W and H, but let this just serve as an example problem.

We could solve this problem with scipy.optimize.minimize by first defining a cost function, and perhaps the first and second derivatives of that function, then initializing W and H and using minimize to calculate the values of W and H that minimize the function. Let's do that:

def f(x, *args):
    A, w_shape, w_size, h_shape = args
    W = x[:w_size].reshape(w_shape)
    H = x[w_size:].reshape(h_shape)
    sse = 0.5 * np.sum((A - W.T.dot(H))**2)
    return sse

def f_prime(x, *args):
    A, w_shape, w_size, h_shape = args
    W = x[:w_size].reshape(w_shape)
    H = x[w_size:].reshape(h_shape)
    gw = H.dot(H.T).dot(W) - H.dot(A.T)
    gh = W.dot(W.T).dot(H) - W.dot(A)
    return np.concatenate([gw.flatten(), gh.flatten()])

A = np.array([13, 18, 17, 24, 21, 30]).reshape(3,2)
W = np.random.rand(2,3)
H = np.random.rand(2,2)
x0 = np.concatenate([W.flatten(), H.flatten()])
out = so.minimize(f, x0, args=(A, W.shape, W.size, H.shape), jac=f_prime, method='BFGS')
print(out['fun'])
2.0172245004128515e-12

So from the output of the minimize function it looks like we successfully calculated a W and H that approximates A (quite well). We can check the reconstruction as well.

Wf = out['x'][:W.size].reshape(W.shape)
Hf = out['x'][W.size:].reshape(H.shape)
print(Wf.T.dot(Hf))
array([[ 13.00000041,  18.00000054],
       [ 17.00000062,  23.99999976],
       [ 21.00000084,  29.99999844]]

Great! But as you can see I implemented some tricks for passing a single flattened array into the x0 argument and then split it up into W and H within each function. This example is pretty simple, but in more complex problems this can get a little annoying/ugly. Moreover, it feels hacky.

There is documentation for using minimize for multiple variables (see Scipy lecture notes: 2.7. Mathematical optimization: finding minima of functions), just not with multiple arrays of different shapes. There are also SO questions along this line like Multiple variables in SciPy's optimize.minimize, but again no mention of the variable being arrays.

Therefore, I am wondering if there is a more elegant method for passing multiple arrays as arguments to minimize. Something akin to:

minimize(f, [W.flatten(), H.flatten()], jac=f_prime, args=...)

Which doesn't work, but seems much more straightforward.