Skip to content

Instantly share code, notes, and snippets.

@Steboss89

Steboss89/jax_concatenate_matrix.py

Created October 21, 2022 13:59
Show Gist options
  • Save Steboss89/82d4da2e08d2bfcbb410b7fc8c151895 to your computer and use it in GitHub Desktop.
Save Steboss89/82d4da2e08d2bfcbb410b7fc8c151895 to your computer and use it in GitHub Desktop.
Download ZIP
Transpose the product matrix to a SIZE x SIZE matrix
Raw
jax_concatenate_matrix.py
c = f(a, c)
# remember to block until ready
c[0][0].block_until_ready()
# convert back to original matrix size
c_arr = np.hstack(np.concatenate(np.array(c), axis=1)).reshape(SIZE,SIZE)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment