code: def matmul(A, B):
      return jnp.matmul(A, B, preferred_element_type=jnp.float32)
  
imports: import time
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: matmul(A, B) and returns: jnp