mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 03:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Allow uncommitted single device PyArray in C++ pjit path.

Browse Source 
PiperOrigin-RevId: 481711690
This commit is contained in:
Kuangyuan Chen 2022-10-17 12:34:55 -07:00 committed by jax authors
parent 504b3c1b25
commit fd2f590b3b
1 changed files with 11 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
benchmarks/api_benchmark.py
View File

@ -696,6 +696,17 @@ def pjit_simple_4000_device(state):
pjit_simple_benchmark(
state, num_devices=4000, num_args=state.range(0), cpp_jit=state.range(1))
@google_benchmark.register
@jax_config.jax_array(True)
def simple_dispatch_array_pjit(state):
a = jax.device_put(1)
b = jax.device_put(2)
f = pjit_lib.pjit(operator.add)
f(a, b)
while state:
f(a, b)
if __name__ == "__main__":
google_benchmark.main()