From c5c3664d2b5542de042765e25b371ba43e7485ac Mon Sep 17 00:00:00 2001 From: Colin Gaffney Date: Mon, 18 May 2026 10:13:51 -0700 Subject: [PATCH] Rename `save/load_pytree` to `save/load`. Eliminate most user-facing "pytree" terminology in favor of "state" as a more specific term. Add `deprecations.py` for handling deprecated public functions. PiperOrigin-RevId: 917299781 --- integration_tests/orbax_safetensors_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/integration_tests/orbax_safetensors_test.py b/integration_tests/orbax_safetensors_test.py index 9096d38..b3c6ef1 100644 --- a/integration_tests/orbax_safetensors_test.py +++ b/integration_tests/orbax_safetensors_test.py @@ -59,7 +59,7 @@ def load_nested_safetensors( with ocp.Context( checkpoint_layout=ocp.options.CheckpointLayout.SAFETENSORS ): - meta = ocp.pytree_metadata(path) + meta = ocp.metadata(path) sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec()) flat_abstract = { k: jax.ShapeDtypeStruct(shape=v.shape, dtype=v.dtype, sharding=sharding) @@ -67,7 +67,7 @@ def load_nested_safetensors( } with ocp.Context(checkpoint_layout=ocp.options.CheckpointLayout.SAFETENSORS): - flat_tree = ocp.load_pytree(path, abstract_pytree=flat_abstract) + flat_tree = ocp.load(path, abstract_state=flat_abstract) if rename_rules: rename_transform = renaming.rename_by_regex(rename_rules)