jax.numpy.meshgrid

jax.numpy.meshgrid(*xi, copy=True, sparse=False, indexing='xy')

Return a list of coordinate matrices from coordinate vectors.

LAX-backend implementation of numpy.meshgrid().

The JAX version of this function may in some cases return a copy rather than a view of the input.

Original docstring below.

Make N-D coordinate arrays for vectorized evaluations of N-D scalar/vector fields over N-D grids, given one-dimensional coordinate arrays x1, x2,…, xn.

Changed in version 1.9: 1-D and 0-D cases are allowed.

Parameters:

• indexing ({'xy', 'ij'}, optional) – Cartesian (‘xy’, default) or matrix (‘ij’) indexing of output. See Notes for more details.

• sparse (bool, optional) –

  If True the shape of the returned coordinate array for dimension i is reduced from (N1, ..., Ni, ... Nn) to (1, ..., 1, Ni, 1, ..., 1). These sparse coordinate grids are intended to be use with Broadcasting. When all coordinates are used in an expression, broadcasting still leads to a fully-dimensonal result array.

  Default is False.

• copy (bool, optional) – If False, a view into the original arrays are returned in order to conserve memory. Default is True. Please note that sparse=False, copy=False will likely return non-contiguous arrays. Furthermore, more than one element of a broadcast array may refer to a single memory location. If you need to write to the arrays, make copies first.

• xi (Union[Array, ndarray, bool_, number, bool, int, float, complex])

Returns:

X1, X2,…, XN – For vectors x1, x2,…, xn with lengths Ni=len(xi), returns (N1, N2, N3,..., Nn) shaped arrays if indexing=’ij’ or (N2, N1, N3,..., Nn) shaped arrays if indexing=’xy’ with the elements of xi repeated to fill the matrix along the first dimension for x1, the second for x2 and so on.

Return type:

list of ndarrays