Simple script that verifies JAX has access to GPU and can do basic ops (matrix multiply)

test_gpu_jax.py

import jax
import jax.numpy as jnp

if __name__ == '__main__':
    print("jax.devices()")
    print(jax.devices())

    a = jnp.asarray([[1.0, 2.0, 3.0], [4., 5., 6.]])
    b = jnp.asarray([[1.0, 2.0], [3.0, 4.0], [5., 6.]])

    print("Array 'a':")
    print(a)
    print("Array 'b':")
    print(b)

    print("where do the arrays live?")
    print(a.device_buffer.device())

    print("Result of jnp.dot(a,b)")
    c = jnp.dot(a,b)
    print(c)