This an unfair example of trying to make the numpy readable:

    n = len(a)
    M = np.exp(1j*k*(np.tile(a,(n,1))**2 + np.tile(a.reshape(n,1),(1,n))**2))

This has a common issue in functional programming and there's an easy trick to fix it. Change the big one liner by naming things:

    n = len(a)
    A = np.tile(a,(n,1))
    A_T = np.tile(a.reshape(n,1),(1,n))
    M = np.exp(1j*k*(A**2 + A_T**2))

Much easier to read now, and it even exposes an error (the sqrt is missing!). Before we fix that, let's make the obvious simplification of using the .T for transpose:

    n = len(a)
    A = np.tile(a,(n,1))
    M = np.exp(1j*k*(A**2 + A.T**2))

Then let's use numpy's broadcasting to be more idiomatic:

    A = a[:, np.newaxis]
    M = np.exp(1j*k*(A**2 + A.T**2))

Now let's fix the bug:

    A = a[:, np.newaxis]
    M = np.exp(1j*k*np.sqrt(A**2 + A.T**2))

or even

    A = a[:, np.newaxis]
    root_sum_squares = np.sqrt(A**2 + A.T**2)
    M = np.exp(1j*k*root_sum_squares)

I think the last two of these are fairer comparisons.

--------

And to pitch the hype train for jax: if you want to fuse the loops and/or get gpu/tpu for free, just swap np for jnp:

    A = a[:, jnp.newaxis]
    M = jnp.exp(1j*k*jnp.sqrt(A**2 + A.T**2))