• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1load("//tensorflow:tensorflow.bzl", "tf_py_test")
2load("//tensorflow/compiler/mlir/tfr:build_defs.bzl", "gen_op_libraries")
3load("//tensorflow/tsl/platform/default:distribute.bzl", "distribute_py_test")
4
5package(
6    default_visibility = [
7        ":friends",
8    ],
9    licenses = ["notice"],
10)
11
12package_group(
13    name = "friends",
14    packages = [
15        "//tensorflow/compiler/mlir/tfr/...",
16    ],
17)
18
19gen_op_libraries(
20    name = "mnist_ops",
21    src = "ops_defs.py",
22    deps = [
23        "//tensorflow:tensorflow_py",
24    ],
25)
26
27tf_py_test(
28    name = "mnist_ops_test",
29    size = "small",
30    srcs = ["mnist_ops_test.py"],
31    data = [":mnist_ops_mlir"],
32    python_version = "PY3",
33    srcs_version = "PY3",
34    tags = [
35        "no_pip",
36        "no_windows",  # TODO(b/170752141)
37        "nomac",  # TODO(b/170752141)
38    ],
39    deps = [
40        ":mnist_ops",
41        ":mnist_ops_py",
42        "//tensorflow:tensorflow_py",
43        "//tensorflow/compiler/mlir/tfr:test_utils",
44    ],
45)
46
47py_library(
48    name = "mnist_train",
49    srcs = ["mnist_train.py"],
50    data = [":mnist_ops_mlir"],
51    srcs_version = "PY3",
52    deps = [
53        ":mnist_ops",
54        ":mnist_ops_py",
55        "//tensorflow:tensorflow_py",
56        "//tensorflow/python:framework",
57        "@absl_py//absl/flags",
58    ],
59)
60
61distribute_py_test(
62    name = "mnist_train_test",
63    size = "medium",
64    srcs = ["mnist_train_test.py"],
65    data = [":mnist_ops_mlir"],
66    disable_v3 = True,  # Not needed. Save some resources and test time.
67    python_version = "PY3",
68    tags = [
69        "no_cuda_asan",  # Not needed, and there were issues with timeouts.
70        "no_oss",  # Avoid downloading mnist data set in oss.
71        "nomultivm",  # Not needed. Save some resources and test time.
72        "notap",  # The test is too long to run as part of llvm presubmits (b/173661843).
73        "notpu",  # Takes too long (b/192305423)
74        "notsan",  # Not needed, and there were issues with timeouts.
75    ],
76
77    # TODO(b/175056184): Re-enable xla_enable_strict_auto_jit once the issues
78    # with GPU and the MLIR bridge are worked out.
79    xla_enable_strict_auto_jit = False,
80    deps = [
81        ":mnist_train",
82        "//tensorflow/python:client_testlib",
83        "//tensorflow/python:extra_py_tests_deps",
84        "//tensorflow/python:is_mlir_bridge_test_true",
85        "//tensorflow/python/distribute:combinations",
86        "//tensorflow/python/distribute:strategy_combinations",
87        "//tensorflow/python/distribute:test_util",
88        "@absl_py//absl/testing:parameterized",
89    ],
90)
91