Add DLRM-V2 with sparsecore.

Hilly12 · Google-ML-Automation · commit 88635ec52d76 · 2025-04-22T20:38:52.000-07:00
PiperOrigin-RevId: 750431182
diff --git a/jax_tpu_embedding/sparsecore/lib/flax/BUILD b/jax_tpu_embedding/sparsecore/lib/flax/BUILD
@@ -18,6 +18,7 @@ package(
     default_applicable_licenses = ["//:license"],
     default_visibility = [
         "//jax_tpu_embedding/sparsecore:__subpackages__",
+        pypi_requirement("recml:__subpackages__"),
         "//video/youtube/discovery/lrm/plum:__subpackages__",
     ],
 )
diff --git a/jax_tpu_embedding/sparsecore/lib/nn/BUILD b/jax_tpu_embedding/sparsecore/lib/nn/BUILD
@@ -18,6 +18,7 @@ package(
     default_applicable_licenses = ["//:license"],
     default_visibility = [
         "//jax_tpu_embedding/sparsecore:__subpackages__",
+        pypi_requirement("recml:__subpackages__"),
         "//video/youtube/discovery/lrm/plum:__subpackages__",
         "//video/youtube/discovery/ml_infra/jax:__subpackages__",
     ],
diff --git a/jax_tpu_embedding/sparsecore/utils/BUILD b/jax_tpu_embedding/sparsecore/utils/BUILD
@@ -16,7 +16,10 @@ load("//third_party/bazel/python:pytype.bzl", "pytype_strict_library")
 
 package(
     default_applicable_licenses = ["//:license"],
-    default_visibility = ["//jax_tpu_embedding/sparsecore:__subpackages__"],
+    default_visibility = [
+        "//jax_tpu_embedding/sparsecore:__subpackages__",
+        pypi_requirement("recml:__subpackages__"),
+    ],
 )
 
 pytype_strict_library(

Original file line number	Diff line number	Diff line change
`@@ -18,6 +18,7 @@ package(`
`18`	`18`	`default_applicable_licenses = ["//:license"],`
`19`	`19`	`default_visibility = [`
`20`	`20`	`"//jax_tpu_embedding/sparsecore:__subpackages__",`
	`21`	`+ pypi_requirement("recml:__subpackages__"),`
`21`	`22`	`"//video/youtube/discovery/lrm/plum:__subpackages__",`
`22`	`23`	`],`
`23`	`24`	`)`