Skip to content

Commit

Permalink
Fix dependencies on orbax/checkpoint/_src/multihost and orbax/checkpo…
Browse files Browse the repository at this point in the history
…int/_src/metadata as these two BUILDs are now added.

PiperOrigin-RevId: 714781886
  • Loading branch information
liangyaning33 authored and Orbax Authors committed Jan 13, 2025
1 parent 1b3c0f1 commit d956793
Showing 1 changed file with 17 additions and 1 deletion.
18 changes: 17 additions & 1 deletion checkpoint/orbax/checkpoint/_src/path/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,11 @@ py_library(
py_library(
name = "step",
srcs = ["step.py"],
deps = [
"//checkpoint/orbax/checkpoint/_src/metadata:checkpoint",
"//checkpoint/orbax/checkpoint/_src/metadata:step_metadata_serialization",
"//checkpoint/orbax/checkpoint/_src/multihost",
],
)

py_test(
Expand All @@ -25,6 +30,8 @@ py_test(
deps = [
":atomicity",
":step",
"//checkpoint/orbax/checkpoint/_src/metadata:checkpoint",
"//checkpoint/orbax/checkpoint/_src/metadata:step_metadata_serialization",
],
)

Expand Down Expand Up @@ -58,6 +65,9 @@ py_library(
":path",
":step",
":utils",
"//checkpoint/orbax/checkpoint/_src/metadata:checkpoint",
"//checkpoint/orbax/checkpoint/_src/metadata:step_metadata_serialization",
"//checkpoint/orbax/checkpoint/_src/multihost",
],
)

Expand All @@ -68,12 +78,14 @@ py_test(
":atomicity",
":atomicity_types",
":step",
"//checkpoint/orbax/checkpoint/_src/multihost",
],
)

py_library(
name = "atomicity_types",
srcs = ["atomicity_types.py"],
deps = ["//checkpoint/orbax/checkpoint/_src/metadata:checkpoint"],
)

py_library(
Expand All @@ -89,10 +101,14 @@ py_library(
py_library(
name = "format_utils",
srcs = ["format_utils.py"],
deps = ["//checkpoint/orbax/checkpoint/_src/metadata:checkpoint"],
)

py_test(
name = "format_utils_test",
srcs = ["format_utils_test.py"],
deps = [":format_utils"],
deps = [
":format_utils",
"//checkpoint/orbax/checkpoint/_src/metadata:checkpoint",
],
)

0 comments on commit d956793

Please sign in to comment.