Add indexing='ij' to jnp.meshgrid call for vmap output to replicate elementwise evaluation

[mattmcd]

In the An Introduction to JAX page the Writing Vectorized Code section uses the default 'xy' cartesian indexing for meshgrid

lecture-jax/lectures/jax_intro.md

Line 614 in 6275b79

x_mesh, y_mesh = jnp.meshgrid(x, y)

This results in output shape (n_y, n_x) for the x_mesh, y_mesh, and z_mesh arrays, while the z_vmap returned from f_vec has shape (n_x, n_y).

Testing with

n = 6000
x = jnp.linspace(-2, 2, n)
y = jnp.linspace(-4, 4, 2*n)

Default xy cartesian indexing

Changing the call to x_grid, y_grid = jnp.meshgrid(x, y, indexing='ij') results in both versions agreeing.

ij matrix indexing

In this example there's no difference since x and y are the same shape and the function f is symmetric but in general this can be a subtle bug that catches people out.

[jstac]

Thanks @mattmcd , much appreciated! This is my mistake.

@mmcky Could you please supervise getting this fixed? Note, however, #267 --- it should be fixed in that version and this lecture should be removed.

Please also check https://python-programming.quantecon.org/numpy_vs_numba_vs_jax.html for the same issue.

[mattmcd]

I forgot to mention but in case it's useful for more context, this behaviour is the same as the MATLAB meshgrid (xy indexing for 2D plots) vs ndgrid (ij indexing for multidimensional arrays) which is where I first encountered it.

[jstac]

Thanks @mattmcd . I always felt like this did the opposite of what I wanted.