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

[JAX] Export shard_map_test for testing on additional JAX backends

Browse Source 
PiperOrigin-RevId: 567522898
This commit is contained in:
Hyeontaek Lim 2023-09-21 22:52:00 -07:00 committed by jax authors
parent 1d8f2ac08f
commit f0bde75dd3
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tests/BUILD
View File

@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
load("@rules_python//python:defs.bzl", "py_test")
load(
"//jaxlib:jax.bzl",
"jax_generate_backend_suites",
@ -21,6 +20,7 @@ load(
"py_deps",
"pytype_test",
)
load("@rules_python//python:defs.bzl", "py_test")
licenses(["notice"])
@ -1207,6 +1207,7 @@ exports_files(
"pmap_test.py",
"pjit_test.py",
"python_callback_test.py",
"shard_map_test.py",
"transfer_guard_test.py",
],
visibility = jax_test_file_visibility,