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

Fix the mesh_axes of benchmarks to be Pspecs

Browse Source 
PiperOrigin-RevId: 437932954
This commit is contained in:
Yash Katariya 2022-03-28 21:50:25 -07:00 committed by jax authors
parent d69c7b3c21
commit a68b0f3a0a
1 changed files with 10 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
benchmarks/gda_benchmark.py
View File

@ -19,19 +19,20 @@ import google_benchmark
import jax
from jax._src import test_util as jtu
from jax._src.util import prod
from jax.interpreters.sharded_jit import PartitionSpec as P
from jax.experimental import global_device_array as gda
import numpy as np
mesh_shapes_axes = [
((256, 8), ["x", "y"]),
((256, 8), [None]),
((256, 8), ["x"]),
((256, 8), ["y"]),
((256, 8), [("x", "y")]),
((128, 8), ["x", "y"]),
((4, 2), ["x", "y"]),
((16, 4), ["x", "y"]),
((16, 4), [("x", "y")]),
((256, 8), P("x", "y")),
((256, 8), P(None)),
((256, 8), P("x")),
((256, 8), P("y")),
((256, 8), P(("x", "y"))),
((128, 8), P("x", "y")),
((4, 2), P("x", "y")),
((16, 4), P("x", "y")),
((16, 4), P(("x", "y"))),
]