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

Internal change.

Browse Source 
PiperOrigin-RevId: 671583042
This commit is contained in:
Sebastian Bodenstein 2024-09-05 18:41:53 -07:00 committed by jax authors
parent 7d438601ae
commit e3b8177af3
2 changed files with 16 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
jaxlib/gpu/BUILD
View File

@ -14,7 +14,12 @@
# Shared CUDA/ROCM GPU kernels.
load("//jaxlib:jax.bzl", "cc_proto_library")
load(
"//jaxlib:jax.bzl",
"cc_proto_library",
"jax_visibility",
"xla_py_proto_library",
)
licenses(["notice"])
@ -72,3 +77,10 @@ cc_proto_library(
name = "triton_cc_proto",
deps = [":triton_proto"],
)
xla_py_proto_library(
name = "triton_py_pb2",
api_version = 2,
visibility = jax_visibility("triton_proto_py_users"),
deps = [":triton_proto"],
)

3
jaxlib/jax.bzl
View File

@ -297,3 +297,6 @@ def jax_generate_backend_suites(backends = []):
)
jax_test_file_visibility = []
def xla_py_proto_library(*args, **kw): # buildifier: disable=unused-variable
pass