Orbax no longer supports lazy restoration of checkpoints.

[copybara-service]

Orbax no longer supports lazy restoration of checkpoints.

If our checkpoint is {'a': 2, 'b': 4}, and we wish to 'lazily' load the tree such that only 'a' is materialized, equivalent functionality can be achieved using the following:

ckptr = ocp.PyTreeCheckpointer()
tree = {'a': <dummy_value>, 'b': <dummy_value>}
ckptr.restore(path, item=tree, transforms={'b': ocp.Transform(use_fallback=True)})
{'a': 2, 'b': <dummy_value>}

OR

ckptr = ocp.PyTreeCheckpointer()
tree = {'a': <dummy_value>}
# Returns a tree just containing 'a', materialized.
ckptr.restore(path, item=tree, transforms={})
{'a': 2}

Previously this could have also been done with the lazy feature:

ckptr = ocp.PyTreeCheckpointer()
ckptr.restore(path, restore_args={'a': RestoreArgs(), 'b': RestoreArgs(lazy=True)})
{'a': 2, 'b': LazyValue()}