Python equivalent for R's expand grid

expand_grid.py

# Adapted from https://stackoverflow.com/questions/12130883/r-expand-grid-function-in-python
#
# Usage: 
# expand_grid([0, 1], [2,3,4])
#
# Output:
# array([[0, 2],
#       [1, 2],
#       [0, 3],
#       [1, 3],
#       [0, 4],
#       [1, 4]])

import numpy as np
from functools import reduce

def expand_grid(*arrs):
    ncols = len(arrs)
    nrows = reduce(lambda x, y: x * y, map(len, arrs), 1)
    
    return np.array(np.meshgrid(*arrs)).reshape(ncols, nrows).T