code: def add_imports(doctest_namespace):
      import jax
      import numpy
  
      doctest_namespace["jax"] = jax
      doctest_namespace["lax"] = jax.lax
      doctest_namespace["jnp"] = jax.numpy
      doctest_namespace["np"] = numpy
  
label_desc: Use image node_id 1 for calling a global function with example usage: add_imports(doctest_namespace) without return types