jax

Modified memory space export test to run on ROCm (for some tests).

#698

Merged