Qiao Zhang zhangqiaorjc

Google Brain: JAX (2020.9 - now), Tensorflow (2018.7-2020.8)

zhangqiaorjc / spmd_pipeline_example.py

Last active November 25, 2021 22:26 — forked from mattjj/spmd_pipeline_example.py

	import itertools as it

	import jax
	import jax.numpy as jnp

	jax.config.update('jax_enable_x64', True)
	jax.config.update('jax_platform_name', 'cpu')

	L = num_stages = 5
	N = batch_size = 6

zhangqiaorjc / zhangqiaorjc_gpu_jax.py

Last active August 24, 2021 23:06 — forked from mattjj/zhangqiaorjc_gpu_jax.py

	# [email protected]
	import atexit
	import functools

	from absl import app
	from absl import flags
	from absl import logging

	import jax
	from jax.lib import xla_extension as xc