Internal change only

PiperOrigin-RevId: 715058507
jax-ml · Feb 11, 2025 · 92397f3 · 92397f3
1 parent 70e3316
commit 92397f3
Show file tree

Hide file tree

Showing 4 changed files with 0 additions and 2,734 deletions.
diff --git a/jax_tpu_embedding/sparsecore/lib/core/primitives/tests/BUILD b/jax_tpu_embedding/sparsecore/lib/core/primitives/tests/BUILD
@@ -46,88 +46,6 @@ pytype_strict_contrib_test(
     ],
 )
 
-pytype_strict_contrib_test(
-    name = "sparse_dense_matmul_csr_with_mini_batching_test",
-    srcs = [
-        "sparse_dense_matmul_csr_with_mini_batching_test.py",
-    ],
-    args = [
-        # TPU chip configuration,
-        "--logtostderr",
-    ],
-    env = {
-        "XLA_FLAGS": "--xla_dump_to=sponge",
-    },
-    tags = [
-        "exclusive-if-local",
-        "requires-tpu",
-    ],
-    deps = [
-        "//jax_tpu_embedding/sparsecore/lib/core:input_preprocessing",
-        "//jax_tpu_embedding/sparsecore/lib/core/primitives:sparse_dense_matmul_csr_with_mini_batching",
-        "//third_party/jax:tpu_support",
-        pypi_requirement("absl/testing:absltest"),
-        pypi_requirement("einops"),
-        pypi_requirement("jax"),
-        pypi_requirement("numpy"),
-    ],
-)
-
-pytype_strict_contrib_test(
-    name = "sparse_dense_matmul_grad_with_sgd_with_mini_batching_test",
-    srcs = [
-        "sparse_dense_matmul_grad_with_sgd_with_mini_batching_test.py",
-    ],
-    args = [
-        # TPU chip configuration,
-        "--logtostderr",
-    ],
-    env = {
-        "XLA_FLAGS": "--xla_dump_to=sponge",
-    },
-    tags = [
-        "exclusive-if-local",
-        "requires-tpu",
-    ],
-    deps = [
-        "//jax_tpu_embedding/sparsecore/lib/core/primitives:sparse_dense_matmul_grad_with_sgd_with_mini_batching",
-        "//third_party/jax:tpu_support",
-        pypi_requirement("absl/logging"),
-        pypi_requirement("absl/testing:absltest"),
-        pypi_requirement("einops"),
-        pypi_requirement("jax"),
-        pypi_requirement("numpy"),
-    ],
-)
-
-pytype_strict_contrib_test(
-    name = "sparse_dense_matmul_grad_with_adagrad_with_mini_batching_test",
-    srcs = [
-        "sparse_dense_matmul_grad_with_adagrad_with_mini_batching_test.py",
-    ],
-    args = [
-        # TPU chip configuration,
-        "--logtostderr",
-    ],
-    env = {
-        "XLA_FLAGS": "--xla_dump_to=sponge",
-    },
-    tags = [
-        "exclusive-if-local",
-        "requires-tpu",
-    ],
-    deps = [
-        "//jax_tpu_embedding/sparsecore/lib/core/primitives:sparse_dense_matmul_grad_with_adagrad_with_mini_batching",
-        "//third_party/jax:tpu_support",
-        pypi_requirement("absl/logging"),
-        pypi_requirement("absl/testing:absltest"),
-        pypi_requirement("absl/testing:parameterized"),
-        pypi_requirement("einops"),
-        pypi_requirement("jax"),
-        pypi_requirement("numpy"),
-    ],
-)
-
 pytype_strict_contrib_test(
     name = "sparse_dense_matmul_grad_with_sgd_test",
     srcs = [