Supplying varying number of input arguments for np.meshgrid

I have a function that uses np.meshgrid to get the matrix form of supplied co-ordinates. I have a parameter dim that determines what dimension I am working with and needs to return an array with dim dimension along axis 1. I have attached an MWE below.

import numpy as np

def func(dim=2):
    a = [1, 2]
    return np.array(np.meshgrid([a]*dim)).T.reshape(-1, dim)

func()
"""
returns
array([[1, 2],
       [1, 2]])
"""

However my expected output is

array([[1, 1],
       [1, 2],
       [2, 1],
       [2, 2]])

, which is obtained by replacing return np.array(np.meshgrid([a]*dim)).T.reshape(-1, dim) with return np.array(np.meshgrid(a, a)).T.reshape(-1, dim). Notice that I passed in (a, a) as the parameters to np.meshgrid, because dim=2. If dim=3, the input would be (a, a, a) and so on.

How can I achieve this? If any other function can do this, I am open to that as well. Thanks.

Answer

Look at the difference between these:

def print_args(*args):
    print(args)

a = [1, 2]
print_args([a] * 2)
# ([[1, 2], [1, 2]],)

print_args(*[a] * 2)
# ([1, 2], [1, 2])

In the first one you are passing a list of lists. In the second one you are unpacking (the first * in func) the list so each sub-list is a positional argument.

def func(dim=2):
    a = [1, 2]
    return np.array(np.meshgrid(*[a] * dim)).T.reshape(-1, dim)

func()
# array([[1, 1],
#        [1, 2],
#        [2, 1],
#        [2, 2]])