code: def print_environment_info(return_string: bool = False) -> str | None:
      """Returns a string containing local environment & JAX installation information.
  
      This is useful information to include when asking a question or filing a bug.
  
      Args: return_string (bool) : if True, return the string rather than printing
      to stdout.
      """
      # TODO(jakevdp): should we include other info, e.g. jax.config.values?
      python_version = sys.version.replace("\n", " ")
      with np.printoptions(threshold=4, edgeitems=2):
          devices_short = str(np.array(xla_bridge.devices())).replace(
              "\n", ""
          )
      info = textwrap.dedent(
          f"""\
    jax:    {version.__version__}
    jaxlib: {lib.version_str}
    numpy:  {np.__version__}
    python: {python_version}
    jax.devices ({xla_bridge.device_count()} total, {xla_bridge.local_device_count()} local): {devices_short}
    process_count: {xla_bridge.process_count()}"""
      )
      nvidia_smi = try_nvidia_smi()
      if nvidia_smi:
          info += "\n\n$ nvidia-smi\n" + nvidia_smi
      if return_string:
          return info
      else:
          return print(info)
  
file_path: jax/jax/_src/environment_info.py
imports: from __future__ import annotations
  import subprocess
  import sys
  import textwrap
  from jax import version
  from jax._src import lib
  from jax._src import xla_bridge
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: print_environment_info(return_string) and returns: info, print