<jupyter_start><jupyter_text>pyhpc-benchmarks @ Google ColabTo run all benchmarks, you need to switch the runtime type to match the corresponding section (CPU, TPU, GPU).<jupyter_code>!rm -rf pyhpc-benchmarks; git clone https://github.com/dionhaefner/pyhpc-benchmarks.git
%cd pyhpc-benchmarks
# check CPU model
!lscpu |grep 'Model name'<jupyter_output><empty_output><jupyter_text>CPU<jupyter_code>!pip install -U -q numba jax tensorflow theano
!taskset -c 0 python run.py benchmarks/equation_of_state/
!taskset -c 0 python run.py benchmarks/isoneutral_mixing/
!taskset -c 0 python run.py benchmarks/turbulent_kinetic_energy/<jupyter_output><empty_output><jupyter_text>TPUMake sure to set accelerator to "TPU" before executing this.<jupyter_code>!pip install -U -q "jax>=0.1.54" "jaxlib>=0.1.37"
os.environ['JAX_BACKEND_TARGET'] = "grpc://" + os.environ['COLAB_TPU_ADDR']
os.environ['JAX_BACKEND_TARGET']
!python run.py benchmarks/equation_of_state -b jax -b numpy --device tpu
!python run.py benchmarks/isoneutral_mixing -b jax -b numpy --device tpu
!python run.py benchmarks/turbulent_kinetic_energy -b jax -b numpy --device tpu<jupyter_output><empty_output><jupyter_text>GPUMake sure to set accelerator to "GPU" before executing this.<jupyter_code># get GPU model
!nvidia-smi -L
!pip install -U -q jax tensorflow-gpu
!for backend in jax tensorflow pytorch cupy; do CUDA_VISIBLE_DEVICES="0" python run.py benchmarks/equation_of_state/ --device gpu -b $backend -b numpy; done
!for backend in jax pytorch cupy; do CUDA_VISIBLE_DEVICES="0" python run.py benchmarks/isoneutral_mixing/ --device gpu -b $backend -b numpy; done
!CUDA_VISIBLE_DEVICES="0" python run.py benchmarks/turbulent_kinetic_energy/ --device gpu -b jax -b numpy<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10813887307911212, 'fraction_numerical': 0.00796812749003984, 'mean_word_length': 5.787644787644788, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 19, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '628254', 'license': 'Unlicense', 'n_tokens_mistral': 665, 'n_tokens_neox': 637, 'n_words': 172, 'path': 'results/pyhpc-benchmarks-colab.ipynb', 'repo_name': 'alexnick83/pyhpc-benchmarks', 'size': '1757'}