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

Load Python rules from rules_python

Browse Source 
PiperOrigin-RevId: 559789250
This commit is contained in:
Richard Levasseur 2023-08-24 10:07:25 -07:00 committed by jax authors
parent 665b176c2c
commit f891cbf64b
13 changed files with 19 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
jax/BUILD
View File

@ -29,6 +29,7 @@ load(
"pytype_library",
"pytype_strict_library",
)
load("@rules_python//python:defs.bzl", "py_library")
package(
default_applicable_licenses = [],

2
jax/_src/pallas/BUILD
View File

@ -11,10 +11,12 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
load(
"//jaxlib:jax.bzl",
"py_deps",
)
load("@rules_python//python:defs.bzl", "py_library")
package(
default_applicable_licenses = [],

2
jax/_src/pallas/mosaic/BUILD
View File

@ -13,11 +13,13 @@
# limitations under the License.
# Package for Mosaic-specific Pallas extensions
load(
"//jaxlib:jax.bzl",
"py_deps",
"py_library_providing_imports_info",
)
load("@rules_python//python:defs.bzl", "py_library")
package(
default_applicable_licenses = [],

2
jax/_src/pallas/triton/BUILD
View File

@ -13,12 +13,14 @@
# limitations under the License.
# Package for Triton-specific Pallas extensions
load(
"//jaxlib:jax.bzl",
"py_deps",
"py_library_providing_imports_info",
"pytype_strict_library",
)
load("@rules_python//python:defs.bzl", "py_library")
package(
default_applicable_licenses = [],

1
jax/experimental/jax2tf/BUILD
View File

@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
load("@rules_python//python:defs.bzl", "py_library")
load(
"//jaxlib:jax.bzl",
"jax2tf_deps",

2
jax/experimental/jax2tf/tests/back_compat_testdata/BUILD
View File

@ -12,6 +12,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
load("@rules_python//python:defs.bzl", "py_library")
licenses(["notice"])
package(

2
jax/experimental/jax2tf/tests/flax_models/BUILD
View File

@ -12,6 +12,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
load("@rules_python//python:defs.bzl", "py_library")
# Note: these examples were imported May 26, 2022 and may be out of sync.
licenses(["notice"])

1
jax/tools/BUILD
View File

@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
load("@rules_python//python:defs.bzl", "py_library")
load(
"//jaxlib:jax.bzl",
"py_deps",

1
jaxlib/cuda/BUILD
View File

@ -14,6 +14,7 @@
# NVIDIA CUDA kernels
load("@rules_python//python:defs.bzl", "py_library")
load(
"//jaxlib:jax.bzl",
"cuda_library",

1
jaxlib/mosaic/BUILD
View File

@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
load("@rules_python//python:defs.bzl", "py_library")
load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
licenses(["notice"])

2
jaxlib/mosaic/python/BUILD
View File

@ -13,8 +13,10 @@
# limitations under the License.
# Mosaic Python bindings
load("@llvm-project//mlir:tblgen.bzl", "gentbl_filegroup")
load("//jaxlib:jax.bzl", "py_deps")
load("@rules_python//python:defs.bzl", "py_library")
gentbl_filegroup(
name = "tpu_python_gen_raw",

1
tests/BUILD
View File

@ -12,6 +12,7 @@
# 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",

1
tests/pallas/BUILD
View File

@ -12,6 +12,7 @@
# 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_test",