wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
tensorflow--tensorflow/tensorflow/compiler/tests/BUILD
T
wehub-resource-sync 8a852e4b4e
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details
chore: import upstream snapshot with attribution
2026-07-13 12:14:16 +08:00

3021 lines
91 KiB
Python
Raw Permalink Blame History

load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.bzl", "py_test")
load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_cuda_cc_test")
load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library")
load("//tensorflow/compiler/tests:build_combined_defs.bzl", "tf_xla_combined_py_test")
load(
"//tensorflow/compiler/tests:build_defs.bzl",
"generate_backend_suites",
"tf_xla_py_strict_test",
# copybara:uncomment_begin(google-only)
# "tpu_backends",
# copybara:uncomment_end
)
load(
"//tensorflow/core/platform:build_config_root.bzl",
"tf_cuda_tests_tags",
)
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
load("//tensorflow/python/tpu:tpu.bzl", "tpu_py_strict_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [":internal"],
licenses = ["notice"],
)
package_group(
name = "internal",
includes = [
"//tensorflow/compiler/tf2xla:internal",
],
)
package_group(
name = "friends",
includes = [
"//tensorflow/compiler/tf2xla:friends",
],
packages = [
# To pass open source testing in the pip Kokoros.
"//platforms/xla/tests/neural_nets",
],
)
generate_backend_suites()
py_library(
name = "xla_test",
testonly = 1,
srcs = ["xla_test.py"],
visibility = [":friends"],
deps = [
"//tensorflow/core:protos_all_py",
"//tensorflow/python/client:session",
"//tensorflow/python/compiler/xla:compiler_py",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:random_seed",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:flags",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/tpu:tpu_py",
"//third_party/py/numpy",
],
)
py_library(
name = "test_utils",
testonly = 1,
srcs = [
"__init__.py",
"test_utils.py",
],
strict_deps = True,
deps = [
"//third_party/py/numpy",
"@six_archive//:six",
],
)
py_test(
name = "xla_test_test",
size = "small",
srcs = ["xla_test_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
deps = [
":xla_test",
"//tensorflow/python/platform:client_testlib",
],
)
#LINT.IfChange(combined_tests)
# If you add a new tf_xla_py_strict_test please either add the test file to one of the combined test
# targets that matches in all tags and other settings or add a new combined test target.
tf_xla_combined_py_test(
name = "combined_ops_test_a",
size = "medium",
timeout = "long",
package = "tensorflow.compiler.tests",
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
tests = [
# go/keep-sorted start
":adadelta_test_lib",
":adagrad_test_lib",
":bincount_op_test_lib",
":bucketize_op_test_lib",
":clustering_test_lib",
":const_arg_test_lib",
":conv2d_test_lib",
":conv3d_test_lib",
":data_format_ops_test_lib",
":dynamic_slice_ops_test_lib",
":dynamic_stitch_test_lib",
# go/keep-sorted end
],
)
tf_xla_combined_py_test(
name = "combined_ops_test_b",
size = "medium",
timeout = "long",
package = "tensorflow.compiler.tests",
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
tests = [
# go/keep-sorted start
":extract_image_patches_op_test_lib",
":fake_quant_ops_test_lib",
":fifo_queue_test_lib",
":ftrl_test_lib",
":function_test_lib",
":fused_batchnorm_test_lib",
":gather_nd_op_test_lib",
":lrn_ops_test_lib",
":manip_ops_test_lib",
# go/keep-sorted end
],
)
tf_xla_combined_py_test(
name = "combined_ops_test_c",
size = "medium",
timeout = "long",
package = "tensorflow.compiler.tests",
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
tests = [
# go/keep-sorted start
":matrix_inverse_op_test_lib",
":matrix_solve_op_test_lib",
":momentum_test_lib",
":nary_ops_test_lib",
":nullary_ops_test_lib",
":placeholder_test_lib",
":pooling_ops_3d_test_lib",
":pooling_ops_test_lib",
":proximal_adagrad_test_lib",
":proximal_gradient_descent_test_lib",
# go/keep-sorted end
],
)
tf_xla_combined_py_test(
name = "combined_ops_test_d",
size = "medium",
timeout = "long",
package = "tensorflow.compiler.tests",
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
tests = [
# go/keep-sorted start
":quantized_ops_test_lib",
":random_ops_test_lib",
":reduce_ops_test_lib",
":reduce_window_test_lib",
":reverse_ops_test_lib",
":rmsprop_test_lib",
":searchsorted_op_test_lib",
":spacetobatch_op_test_lib",
":sparse_to_dense_op_test_lib",
# go/keep-sorted end
],
)
tf_xla_combined_py_test(
name = "combined_ops_test_e",
size = "medium",
timeout = "long",
exec_properties = {
"cpp_link.mem": "16g",
},
package = "tensorflow.compiler.tests",
tags = [
"no_cuda_asan", # times out in individual tests
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"noasan", # times out consistently from 2023-08-24
],
tests = [
# go/keep-sorted start
":adagrad_da_test_lib",
":adam_test_lib",
":argminmax_test_lib",
":listdiff_op_test_lib",
":slice_ops_test_lib",
":unary_ops_test_lib",
# go/keep-sorted end
],
)
tf_xla_combined_py_test(
name = "combined_ops_test_f",
size = "medium",
timeout = "long",
# copybara:uncomment_begin(google-only)
# disabled_backends = tpu_backends(),
# copybara:uncomment_end
exec_properties = {
"cpp_link.mem": "16g",
},
package = "tensorflow.compiler.tests",
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
tests = [
# go/keep-sorted start
":add_n_test_lib",
":cond_test_lib",
":float_ops_test_lib",
":while_test_lib",
# go/keep-sorted end
],
)
tf_xla_combined_py_test(
name = "combined_ops_test_g",
size = "medium",
timeout = "long",
exec_properties = {
"cpp_link.mem": "16g",
},
package = "tensorflow.compiler.tests",
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
tests = [
# go/keep-sorted start
":cast_test_lib",
# go/keep-sorted end
],
)
#LINT.ThenChange(:individual_tests)
#LINT.IfChange(individual_tests)
tf_xla_py_strict_test(
name = "adadelta_test",
size = "medium",
srcs = ["adadelta_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/training:adadelta",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "adagrad_test",
size = "small",
srcs = ["adagrad_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap", # Handled in the combined test above.
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/training:adagrad",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "adagrad_da_test",
size = "small",
srcs = ["adagrad_da_test.py"],
tags = [
"no_cuda_asan", # times out
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/training:adagrad_da",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "adam_test",
size = "small",
srcs = ["adam_test.py"],
tags = [
"no_cuda_asan", # times out
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/training:adam",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "add_n_test",
size = "small",
srcs = ["add_n_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:list_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
],
)
tf_xla_py_strict_test(
name = "argminmax_test",
size = "small",
srcs = ["argminmax_test.py"],
tags = [
"no_cuda_asan", # times out
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"noasan", # times out consistently from 2023-08-24
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "binary_ops_test",
size = "medium",
srcs = ["binary_ops_test.py"],
shard_count = 5,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly", # Times out frequently in fastbuild mode.
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_gen",
"//tensorflow/python/ops:bitwise_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:math_ops_gen",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:nn_ops_gen",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "complex_div_test",
size = "medium",
srcs = ["complex_div_test.py"],
enabled_backends = [
"cpu",
"gpu",
"gpu_a100",
"gpu_h100",
],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly", # Times out frequently in fastbuild mode.
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops_gen",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "bucketize_op_test",
size = "small",
srcs = ["bucketize_op_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
],
)
tf_xla_py_strict_test(
name = "categorical_op_test",
size = "small",
srcs = ["categorical_op_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:random_seed",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:stateless_random_ops",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "cholesky_op_test",
size = "medium",
srcs = ["cholesky_op_test.py"],
disabled_backends = ["gpu_h100"], # TODO: b/362907482 - Re-enable after fixing.
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"no_rocm",
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:linalg_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "cond_test",
size = "small",
srcs = ["cond_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/client:session",
"//tensorflow/python/compiler/xla",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:cond",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:control_flow_switch_case",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:tensor_array_ops",
"//tensorflow/python/platform:client_testlib",
],
)
tf_xla_py_strict_test(
name = "self_adjoint_eig_op_test",
size = "medium",
srcs = ["self_adjoint_eig_op_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:linalg_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "searchsorted_op_test",
size = "small",
timeout = "moderate",
srcs = ["searchsorted_op_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "svd_op_test",
size = "medium",
srcs = ["svd_op_test.py"],
disabled_backends = [
# TODO(b/129396575): Fails on CPU.
"cpu",
"cpu_ondemand",
],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"no_rocm",
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:linalg_ops",
"//tensorflow/python/ops:linalg_ops_gen",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "matrix_inverse_op_test",
size = "small",
timeout = "moderate",
srcs = ["matrix_inverse_op_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:linalg_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "matrix_solve_op_test",
size = "small",
timeout = "moderate",
srcs = ["matrix_solve_op_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/ops:linalg_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/platform:sysconfig",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "matrix_triangular_solve_op_test",
size = "small",
timeout = "moderate",
srcs = ["matrix_triangular_solve_op_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:linalg_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "clustering_test",
size = "small",
srcs = ["clustering_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "concat_ops_test",
size = "medium",
srcs = ["concat_ops_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"noasan", # Timed out on 2023-07-12
],
deps = [
":xla_test",
"//tensorflow/python:pywrap_sanitizers",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_gen",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:gradients_impl",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "conv2d_test",
size = "medium",
srcs = ["conv2d_test.py"],
shard_count = 10,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":test_utils",
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:nn_ops_gen",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "tensor_float_32_test",
size = "medium",
srcs = ["tensor_float_32_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
use_xla_device = False, # Uses tf.function(jit_compile=True)
deps = [
":xla_test",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/platform:sysconfig",
"//tensorflow/python/platform:test",
],
)
tf_xla_py_strict_test(
name = "conv3d_test",
size = "medium",
srcs = ["conv3d_test.py"],
shard_count = 5,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":test_utils",
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:gradient_checker",
"//tensorflow/python/ops:nn_grad",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:nn_ops_gen",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "depthwise_conv_op_test",
size = "medium",
srcs = ["depthwise_conv_op_test.py"],
shard_count = 5,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"no_rocm",
"optonly", # Times out frequently in fastbuild mode.
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:nn_grad",
"//tensorflow/python/ops:nn_impl",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "dynamic_slice_ops_test",
size = "small",
srcs = ["dynamic_slice_ops_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/compiler/tf2xla/python:xla",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "einsum_op_test",
size = "medium",
srcs = ["einsum_op_test.py"],
enabled_backends = [
"cpu",
"gpu",
"gpu_a100",
"gpu_h100",
],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:special_math_ops",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "reshape_op_test",
size = "small",
srcs = ["reshape_op_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:test",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "dynamic_stitch_test",
size = "small",
srcs = ["dynamic_stitch_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:data_flow_ops",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "extract_image_patches_op_test",
size = "small",
srcs = ["extract_image_patches_op_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "eager_test",
size = "medium",
srcs = ["eager_test.py"],
tags = [
"multi_and_single_gpu",
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
deps = [
":xla_test",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/eager:backprop",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:indexed_slices",
"//tensorflow/python/framework:ops",
"//tensorflow/python/layers",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:cond",
"//tensorflow/python/ops:embedding_ops",
"//tensorflow/python/ops:functional_ops",
"//tensorflow/python/ops:init_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:random_ops_gen",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/platform:test",
"//tensorflow/python/training:adam",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "fifo_queue_test",
size = "medium",
srcs = ["fifo_queue_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python:extra_py_tests_deps",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:data_flow_ops",
"//tensorflow/python/platform:client_testlib",
],
)
tf_xla_py_strict_test(
name = "fft_test",
size = "medium",
srcs = ["fft_test.py"],
shard_count = 12,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"no_rocm",
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python:extra_py_tests_deps",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:gradients_impl",
"//tensorflow/python/ops/signal",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "slice_ops_test",
size = "medium",
srcs = ["slice_ops_test.py"],
tags = [
"no_cuda_asan", # times out
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:test",
],
)
tf_xla_py_strict_test(
name = "ftrl_test",
size = "medium",
srcs = ["ftrl_test.py"],
shard_count = 8,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/training:adagrad",
"//tensorflow/python/training:ftrl",
"//tensorflow/python/training:gradient_descent",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "ftrl_ops_test",
size = "medium",
srcs = ["ftrl_ops_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:training_ops_gen",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "function_test",
size = "small",
srcs = ["function_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:function",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "image_ops_test",
size = "small",
timeout = "long",
srcs = ["image_ops_test.py"],
enabled_backends = [
"cpu",
"gpu",
"gpu_a100",
"gpu_h100",
],
shard_count = 10,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly", # Times out frequently in fastbuild mode.
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:image_ops",
"//tensorflow/python/ops:image_ops_gen",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "listdiff_op_test",
size = "small",
srcs = ["listdiff_op_test.py"],
tags = [
"no_cuda_asan", # times out
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "lrn_ops_test",
size = "medium",
srcs = ["lrn_ops_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/ops:nn_ops_gen",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "manip_ops_test",
size = "small",
srcs = ["manip_ops_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:manip_ops",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "matrix_band_part_test",
size = "medium",
timeout = "long",
srcs = ["matrix_band_part_test.py"],
tags = [
"no_aarch64", # TODO(b/315533266)
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "matrix_diag_ops_test",
size = "medium",
timeout = "long",
srcs = ["matrix_diag_ops_test.py"],
shard_count = 4,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
deps = [
":xla_test",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_gen",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "momentum_test",
size = "small",
srcs = ["momentum_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/training:momentum",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "nary_ops_test",
size = "small",
srcs = ["nary_ops_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "nullary_ops_test",
size = "small",
srcs = ["nullary_ops_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "pooling_ops_test",
size = "medium",
srcs = ["pooling_ops_test.py"],
shard_count = 20,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:nn_ops_gen",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "pooling_ops_3d_test",
size = "medium",
srcs = ["pooling_ops_3d_test.py"],
shard_count = 20,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:nn_ops_gen",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "proximal_adagrad_test",
size = "medium",
srcs = ["proximal_adagrad_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap", # Handled in the combined test above.
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/training:adagrad",
"//tensorflow/python/training:proximal_adagrad",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "proximal_gradient_descent_test",
size = "medium",
srcs = ["proximal_gradient_descent_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap", # Handled in the combined test above.
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/training:gradient_descent",
"//tensorflow/python/training:proximal_gradient_descent",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "qr_op_test",
size = "medium",
srcs = ["qr_op_test.py"],
disabled_backends = [
# Test is very slow on CPU.
"cpu",
"cpu_ondemand",
],
shard_count = 5,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"nozapfhahn",
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:linalg_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "unstack_test",
size = "medium",
srcs = ["unstack_test.py"],
disabled_backends = [
# TODO(b/149750262): timeout on CPU.
"cpu",
"cpu_ondemand",
],
shard_count = 1,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "random_ops_test",
size = "medium",
srcs = ["random_ops_test.py"],
shard_count = 10,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops/distributions:special_math",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "reduce_ops_test",
size = "medium",
srcs = ["reduce_ops_test.py"],
shard_count = 5,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "reduce_window_test",
size = "small",
srcs = ["reduce_window_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/compiler/tf2xla/python:xla",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:function",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "reverse_ops_test",
size = "medium",
srcs = ["reverse_ops_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "reverse_sequence_op_test",
size = "medium",
srcs = ["reverse_sequence_op_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
# copybara:uncomment_begin(google-only)
# tf_xla_py_strict_test(
# name = "reverse_sequence_op_args_test",
# size = "medium",
# srcs = ["reverse_sequence_op_args_test.py"],
# main = "reverse_sequence_op_args_test.py",
# tags = [
# "no_pip",
# "optonly",
# ],
# deps = [
# ":xla_test",
# "//tensorflow/compiler/jit:xla_cpu_jit", # DisableOnExport
# "//tensorflow/python/compat:v2_compat",
# "//tensorflow/python/eager:def_function",
# "//tensorflow/python/framework:errors",
# "//tensorflow/python/ops:array_ops",
# "//tensorflow/python/platform:client_testlib",
# ],
# )
# copybara:uncomment_end
tf_xla_py_strict_test(
name = "rmsprop_test",
size = "small",
srcs = ["rmsprop_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap", # Handled in the combined test above.
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/training:rmsprop",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "scan_ops_test",
size = "medium",
timeout = "long",
srcs = ["scan_ops_test.py"],
shard_count = 6,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"nozapfhahn",
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "segment_reduction_ops_test",
size = "medium",
srcs = ["segment_reduction_ops_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/client:device_lib",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "spacetobatch_op_test",
size = "medium",
srcs = ["spacetobatch_op_test.py"],
shard_count = 3,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_gen",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "sparse_to_dense_op_test",
size = "medium",
srcs = ["sparse_to_dense_op_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:sparse_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "stack_ops_test",
size = "small",
srcs = ["stack_ops_test.py"],
enabled_backends = [
"cpu",
"gpu",
"gpu_a100",
"gpu_h100",
],
tags = [
"config-cuda-only",
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
use_xla_device = False,
deps = [
":xla_test",
"//tensorflow/python/compiler/xla",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:data_flow_ops_gen",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "stateful_random_ops_test",
size = "medium",
srcs = ["stateful_random_ops_test.py"],
disabled_backends = [
"gpu",
"gpu_a100",
"gpu_h100",
],
shard_count = 10,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/client",
"//tensorflow/python/client:device_lib",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/kernel_tests/random:util",
"//tensorflow/python/ops:random_ops_util",
"//tensorflow/python/ops:stateful_random_ops",
"//tensorflow/python/ops:stateful_random_ops_gen",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:flags",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "stateless_random_ops_test",
size = "medium",
srcs = ["stateless_random_ops_test.py"],
disabled_backends = [
"gpu",
"gpu_a100",
"gpu_h100",
],
shard_count = 10,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"nozapfhahn", # Times out under coverage
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/client:device_lib",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/kernel_tests/random:util",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:random_ops_util",
"//tensorflow/python/ops:stateless_random_ops",
"//tensorflow/python/ops:stateless_random_ops_v2_gen",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "stochastic_cast_op_test",
size = "medium",
srcs = ["stochastic_cast_op_test.py"],
# TODO(b/232442915): Enable GPUs.
disabled_backends = [
"gpu",
"gpu_a100",
"gpu_h100",
],
shard_count = 20,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/kernel_tests/random:util",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:standard_ops",
"//tensorflow/python/ops:stateless_random_ops",
"//tensorflow/python/ops:stochastic_cast_op",
"//tensorflow/python/ops:tensor_array_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "tensor_array_ops_test",
size = "medium",
srcs = ["tensor_array_ops_test.py"],
enabled_backends = [
"cpu",
"gpu",
"gpu_a100",
"gpu_h100",
],
tags = [
"config-cuda-only",
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"v1only",
],
use_xla_device = False,
deps = [
":xla_test",
"//tensorflow/python/compiler/xla",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:control_flow_util",
"//tensorflow/python/ops:data_flow_ops_gen",
"//tensorflow/python/ops:gradients_impl",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:tensor_array_grad",
"//tensorflow/python/ops:tensor_array_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "tensor_list_ops_test",
size = "small",
srcs = ["tensor_list_ops_test.py"],
# TensorList ops are only implemented on CPU.
enabled_backends = ["cpu"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:list_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "ternary_ops_test",
size = "medium",
srcs = ["ternary_ops_test.py"],
shard_count = 4,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
deps = [
":xla_test",
"//tensorflow/python:extra_py_tests_deps",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:math_ops_gen",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "unary_ops_test",
size = "medium",
srcs = ["unary_ops_test.py"],
shard_count = 20,
tags = [
"no_cuda_asan", # times out
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"noasan", #times out
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:bitwise_ops",
"//tensorflow/python/ops:functional_ops_gen",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:nn_ops_gen",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
"@six_archive//:six",
],
)
tf_xla_py_strict_test(
name = "float_ops_test",
size = "medium",
srcs = ["float_ops_test.py"],
tags = [
"no_cuda_asan", # times out
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"noasan", #times out
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "cast_test",
size = "medium",
srcs = ["cast_test.py"],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "fused_batchnorm_test",
size = "medium",
srcs = ["fused_batchnorm_test.py"],
shard_count = 5,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":test_utils",
":xla_test",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:gradient_checker",
"//tensorflow/python/ops:nn",
"//tensorflow/python/ops:nn_ops_gen",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "variable_ops_test",
size = "small",
timeout = "moderate",
srcs = ["variable_ops_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:init_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:state_ops",
"//tensorflow/python/ops:state_ops_gen",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:test",
"//tensorflow/python/training:gradient_descent",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "while_test",
size = "small",
srcs = ["while_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/compiler/tf2xla/python:xla",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:function",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:gradients_impl",
"//tensorflow/python/ops:map_fn",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "case_test",
size = "small",
srcs = ["case_test.py"],
disabled_backends = ["cpu_ondemand"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
use_xla_device = False, # Uses tf.function(jit_compile=True)
deps = [
":xla_test",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:control_flow_switch_case",
"//tensorflow/python/ops:image_ops",
"//tensorflow/python/ops:io_ops",
"//tensorflow/python/platform:client_testlib",
],
)
tf_xla_py_strict_test(
name = "cast_ops_test",
size = "small",
srcs = ["cast_ops_test.py"],
disabled_backends = ["cpu_ondemand"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
use_xla_device = False, # Uses tf.function(jit_compile=True)
deps = [
":xla_test",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:image_ops",
"//tensorflow/python/ops:io_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/platform:client_testlib",
],
)
tf_xla_py_strict_test(
name = "gather_test",
size = "medium",
srcs = ["gather_test.py"],
tags = [
"no_pip",
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:flags",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "gather_nd_op_test",
size = "medium",
srcs = ["gather_nd_op_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "scatter_nd_op_test",
size = "medium",
srcs = ["scatter_nd_op_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "sort_ops_test",
size = "medium",
srcs = ["sort_ops_test.py"],
shard_count = 50,
# Times out in fastbuild mode.
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"nozapfhahn",
"optonly",
],
deps = [
":xla_test",
"//tensorflow/compiler/tf2xla/python:xla",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:function",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:sort_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "data_format_ops_test",
size = "small",
srcs = ["data_format_ops_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "xla_device_test",
size = "medium",
srcs = ["xla_device_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:control_flow_ops_gen",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "fake_quant_ops_test",
size = "medium",
srcs = ["fake_quant_ops_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_gen",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "placeholder_test",
size = "small",
srcs = ["placeholder_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
],
deps = [
":xla_test",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:test",
],
)
tf_xla_py_strict_test(
name = "quantized_ops_test",
size = "medium",
srcs = ["quantized_ops_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap", # Handled in the combined test above.
],
deps = [
":xla_test",
"//tensorflow/compiler/tf2xla/python:xla",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:bitwise_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "xla_ops_test",
size = "medium",
srcs = ["xla_ops_test.py"],
disabled_backends = [
"gpu",
"gpu_a100",
"gpu_h100",
],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
deps = [
":xla_test",
"//tensorflow/compiler/tf2xla/ops:gen_xla_ops",
"//tensorflow/compiler/tf2xla/python:xla",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:function",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:random_ops_util",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
"@xla//xla:xla_data_proto_py",
],
)
tf_xla_py_strict_test(
name = "xla_custom_call_ops_test",
size = "small",
srcs = ["xla_custom_call_ops_test.py"],
disabled_backends = [
"gpu",
"gpu_a100",
"gpu_h100",
],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap", # Handled in the combined test above.
],
use_xla_device = False, # Uses tf.function(jit_compile=True)
deps = [
":xla_test",
"//tensorflow/compiler/tf2xla/python:xla",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/platform:client_testlib",
],
)
tf_xla_py_strict_test(
name = "runtime_shape_check_test",
size = "small",
srcs = ["runtime_shape_check_test.py"],
disabled_backends = [
"cpu",
"cpu_ondemand",
],
tags = [
"no_pip",
"notap",
],
use_xla_device = False,
deps = [
":xla_test",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:client_testlib",
],
)
tf_xla_py_strict_test(
name = "conv_node_name_test",
size = "medium",
srcs = ["conv_node_name_test.py"],
tags = [
"no_oss", # TODO(b/148108508): Re-enable this test in OSS.
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
deps = [
":xla_test",
"//tensorflow/python/framework:ops",
"//tensorflow/python/layers",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "tridiagonal_solve_ops_test",
size = "medium",
srcs = ["tridiagonal_solve_ops_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:gradients",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops/linalg:linalg_impl",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "tridiagonal_matmul_ops_test",
size = "medium",
srcs = ["tridiagonal_matmul_ops_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:gradient_checker_v2",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:stateless_random_ops",
"//tensorflow/python/ops/linalg:linalg_impl",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tf_xla_py_strict_test(
name = "special_math_test",
size = "medium",
srcs = ["special_math_test.py"],
shard_count = 5,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python:extra_py_tests_deps",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/ops:gradient_checker_v2",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:math_ops_gen",
"//tensorflow/python/ops:random_ops_gen",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/flags",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "repeat_op_test",
size = "medium",
srcs = ["repeat_op_test.py"],
shard_count = 1,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python:extra_py_tests_deps",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:client_testlib",
],
)
tf_xla_py_strict_test(
name = "image_ops_jit_compile_test",
size = "medium",
srcs = ["image_ops_jit_compile_test.py"],
disabled_backends = [
"cpu_ondemand",
],
shard_count = 1,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly",
],
use_xla_device = False,
deps = [
":xla_test",
"//tensorflow/python:extra_py_tests_deps",
"//tensorflow/python/eager:backprop",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:image_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
],
)
tf_xla_py_strict_test(
name = "ensure_shape_op_test",
size = "medium",
srcs = ["ensure_shape_op_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:check_ops",
"//tensorflow/python/platform:client_testlib",
],
)
tf_xla_py_strict_test(
name = "where_op_test",
size = "small",
srcs = ["where_op_test.py"],
enabled_backends = [
"cpu",
"gpu",
"gpu_a100",
"gpu_h100",
],
tags = [
"no_pip",
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/tpu:tpu_py",
],
)
tf_xla_py_strict_test(
name = "where_op_tpu_test",
size = "small",
srcs = ["where_op_test.py"],
disabled_backends = [
"cpu",
"cpu_ondemand",
"gpu",
"gpu_a100",
"gpu_h100",
],
main = "where_op_test.py",
tags = [
"no_pip",
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/tpu:tpu_py",
],
)
tf_xla_py_strict_test(
name = "const_arg_test",
size = "small",
srcs = ["const_arg_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap", # Handled in the combined test above.
],
deps = [
":xla_test",
"//tensorflow/compiler/tf2xla/python:xla",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:test",
],
)
tf_xla_py_strict_test(
name = "sharding_util_ops_test",
srcs = ["sharding_util_ops_test.py"],
disabled_backends = [
"cpu",
"cpu_ondemand",
"gpu",
"gpu_a100",
"gpu_h100",
],
shard_count = 10,
tags = [
"notap",
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/client:session",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:tpu_ops_gen",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tpu_py_strict_test(
name = "approx_topk_test",
srcs = ["approx_topk_test.py"],
disable_experimental = False,
deps = [
"//tensorflow/python/eager:backprop",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/eager:test",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:variables",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "xla_call_module_test",
size = "small",
srcs = ["xla_call_module_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
use_xla_device = False, # Uses tf.function(jit_compile=True)
deps = [
":xla_test",
"//tensorflow/compiler/mlir/stablehlo",
"//tensorflow/compiler/tf2xla/ops:gen_xla_ops",
"//tensorflow/compiler/tf2xla/python:xla",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:function",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "xla_call_module_no_platform_check_test",
size = "small",
srcs = ["xla_call_module_test.py"],
# cpu_ondemand overrides the TF_XLA_FLAGS
disabled_backends = ["cpu_ondemand"],
env = {"TF_XLA_FLAGS": "--tf_xla_call_module_disabled_checks=platform"},
main = "xla_call_module_test.py",
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
use_xla_device = False, # Uses tf.function(jit_compile=True)
deps = [
":xla_test",
"//tensorflow/compiler/mlir/stablehlo",
"//tensorflow/compiler/tf2xla/ops:gen_xla_ops",
"//tensorflow/compiler/tf2xla/python:xla",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:function",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "xla_call_module_no_shape_assertions_check_test",
size = "small",
srcs = ["xla_call_module_test.py"],
disabled_backends = ["cpu_ondemand"], # cpu_ondemand overrides the TF_XLA_FLAGS
env = {"TF_XLA_FLAGS": "--tf_xla_call_module_disabled_checks=shape_assertions"},
main = "xla_call_module_test.py",
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
use_xla_device = False, # Uses tf.function(jit_compile=True)
deps = [
":xla_test",
"//tensorflow/compiler/mlir/stablehlo",
"//tensorflow/compiler/tf2xla/ops:gen_xla_ops",
"//tensorflow/compiler/tf2xla/python:xla",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:function",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "bincount_op_test",
size = "small",
srcs = ["bincount_op_test.py"],
shard_count = 1,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap", # Handled in the combined test above.
],
deps = [
":xla_test",
"//tensorflow/python/compat:v2_compat",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:errors",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops_gen",
"//tensorflow/python/platform:test",
],
)
tf_xla_py_strict_test(
name = "unique_ops_test",
size = "small",
srcs = ["unique_ops_test.py"],
enabled_backends = [
"cpu",
"gpu",
"gpu_a100",
"gpu_h100",
],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
deps = [
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_gen",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
tpu_py_strict_test(
name = "mean_op_test",
size = "small",
srcs = ["mean_op_test.py"],
disable_experimental = False,
tags = [
"notsan", # timesout
],
deps = [
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
],
)
tf_xla_py_strict_test(
name = "xla_dump_to_test",
size = "medium",
srcs = ["xla_dump_to_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"optonly",
],
deps = [
":xla_test",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:test",
"//third_party/py/numpy",
],
)
# copybara:uncomment_begin(google-only)
# tf_xla_py_strict_test(
# name = "xla_dump_to_sponge_test",
# size = "medium",
# srcs = ["xla_dump_to_sponge_test.py"],
# tags = [
# "optonly",
# ],
# deps = [
# ":xla_test",
# "//third_party/py/numpy",
# "//tensorflow/python/ops:array_ops",
# "//tensorflow/python/ops:math_ops",
# "//tensorflow/python/platform:test",
# ],
# )
# copybara:uncomment_end
#LINT.ThenChange(:combined_tests)
cuda_py_strict_test(
name = "xla_device_gpu_test",
size = "small",
srcs = ["xla_device_gpu_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
xla_enable_strict_auto_jit = False,
xla_enabled = True,
deps = [
"//tensorflow/python/client:session",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "jit_test",
size = "medium",
srcs = ["jit_test.py"],
#shard_count = 5,
tags = [
"no_cuda_asan", # Times out.
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
xla_enable_strict_auto_jit = False,
xla_enabled = True,
deps = [
":test_utils",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/client:session",
"//tensorflow/python/compiler/xla:compiler_py",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:function",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:cond",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:gradients_impl",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
cuda_py_strict_test(
name = "async_comp_test",
size = "medium",
srcs = ["async_comp_test.py"],
shard_count = 1,
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
xla_enable_strict_auto_jit = False,
xla_enabled = True,
deps = [
"//tensorflow/core:protos_all_py",
"//tensorflow/python/client:session",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:function",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "dense_layer_test",
size = "medium",
srcs = ["dense_layer_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
xla_enable_strict_auto_jit = False,
xla_enabled = True,
deps = [
":test_utils",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/compiler/xla:compiler_py",
"//tensorflow/python/framework:ops",
"//tensorflow/python/layers",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
cc_library(
name = "randomized_tests_library",
testonly = 1,
srcs = ["randomized_tests.cc"],
deps = [
"//tensorflow/compiler/jit",
"//tensorflow/compiler/jit:common",
"//tensorflow/compiler/jit:flags_headers",
"//tensorflow/core:core_cpu",
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:tensorflow_opensource",
"//tensorflow/core:test",
"//tensorflow/core:testlib",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:fixed_array",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
"@xla//xla:xla_data_proto_cc",
"@xla//xla/tsl/platform:status",
],
)
tf_cuda_cc_test(
name = "randomized_tests",
size = "medium",
args = ["--tf_xla_test_use_mlir=false"],
shard_count = 20,
# This test is randomized, so only run it if explicitly requested.
tags = [
"manual",
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
] + tf_cuda_tests_tags(),
deps = [":randomized_tests_library"],
)
tf_cuda_cc_test(
name = "randomized_tests_mlir",
size = "medium",
args = ["--tf_xla_test_use_mlir=true"],
shard_count = 20,
# This test is randomized, so only run it if explicitly requested.
tags = [
"manual",
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"notap",
] + tf_cuda_tests_tags(),
deps = [":randomized_tests_library"],
)
# Create a deterministic version of randomized_tests_mlir with fixed seed.
# This can be used in presubmit checks as it is no longer randomized.
tf_cuda_cc_test(
name = "randomized_tests_mlir_seeded",
size = "medium",
args = [
"--tf_xla_random_seed=200839030",
"--tf_xla_test_use_mlir=true",
"--tf_xla_test_device=GPU:0",
],
shard_count = 20,
tags = [
"config-cuda-only",
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"no_rocm", # ROCmSoftwarePlatform #958
"noasan", # TODO(b/201651800)
"requires-gpu-nvidia",
] + tf_cuda_tests_tags(),
deps = [":randomized_tests_library"],
)
# Create a deterministic version of randomized_tests with fixed seed.
# This can be used in presubmit checks as it is no longer randomized.
tf_cuda_cc_test(
name = "randomized_tests_seeded",
size = "medium",
args = [
"--tf_xla_random_seed=200839030",
"--tf_xla_test_use_mlir=false",
"--tf_xla_test_device=GPU:0",
],
shard_count = 20,
tags = [
"config-cuda-only",
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
"no_rocm", # ROCmSoftwarePlatform #958
"noasan", # TODO(b/201651800)
"requires-gpu-nvidia",
] + tf_cuda_tests_tags(),
deps = [":randomized_tests_library"],
)
tf_cuda_cc_test(
name = "unary_ops_composition_test",
srcs = ["unary_ops_composition_test.cc"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
] + tf_cuda_tests_tags(),
deps = [
"//tensorflow/cc:cc_ops",
"//tensorflow/compiler/jit",
"//tensorflow/compiler/jit:flags",
"//tensorflow/compiler/jit:xla_kernel_creator",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"//tensorflow/core/kernels:ops_testutil",
"@com_google_absl//absl/log:check",
],
)
py_library(
name = "lstm",
testonly = 1,
srcs = ["lstm.py"],
strict_deps = True,
deps = [
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:variable_v1",
"@six_archive//:six",
],
)
cuda_py_strict_test(
name = "lstm_test",
srcs = ["lstm_test.py"],
tags = [
"no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip
],
xla_enable_strict_auto_jit = False,
xla_enabled = True,
deps = [
":lstm",
":xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:gradients_impl",
"//tensorflow/python/ops:init_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
# An example of ahead-of-time compilation using tfcompile. The
# lstm_layer_inference.pbtxt file was generated by running lstm_test
# --dump_graph_dir, and the config file was written by hand.
#
# Run the following to build a minimal benchmark of the computation on Android:
# $ bazel build -c opt --cxxopt='-std=c++11' --linkopt='-lm' \
# --cpu=armeabi-v7a \
# --host_crosstool_top=@bazel_tools//tools/cpp:toolchain \
# --crosstool_top=//external:android/crosstool \
# //tensorflow/compiler/tests:lstm_layer_inference_benchmark
#
# Currently the resulting binary size is ~190KB
tf_library(
name = "lstm_layer_inference",
testonly = 1,
config = "lstm_layer_inference.config.pbtxt",
cpp_class = "LSTMLayerInference",
graph = "lstm_layer_inference.pbtxt",
tags = ["manual"],
tfcompile_flags = ["--xla_cpu_multi_thread_eigen=false"],
)
cuda_py_strict_test(
name = "const_test",
size = "small",
srcs = ["const_test.py"],
xla_enable_strict_auto_jit = False,
xla_enabled = True,
deps = [
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tpu_py_strict_test(
name = "giant_const_op_test",
srcs = [
"giant_const_op_test.py",
],
disable_experimental = True,
# TODO(b/188995810): Add an optimization in MLIR importer to not
# materialize giant splat constants.
deps = [
"//tensorflow/python/distribute:tpu_strategy",
"//tensorflow/python/distribute/cluster_resolver:tpu_cluster_resolver_py",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/eager:remote",
"//tensorflow/python/eager:test",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/platform:flags",
"//third_party/py/numpy",
],
)
Reference in New Issue View Git Blame Copy Permalink