Utility to set value in a pytree (and so in state)

PiperOrigin-RevId: 615524997

Author: vroulet

docs/api/utilities.rst

@@ -106,6 +106,7 @@ Tree
     tree_ones_like
     tree_random_like
     tree_scalar_mul
+    tree_set
     tree_sub
     tree_sum
     tree_vdot
@@ -155,6 +156,10 @@ Tree scalar multiply
 ~~~~~~~~~~~~~~~~~~~~
 .. autofunction:: tree_scalar_mul
 
+Set values in a tree
+~~~~~~~~~~~~~~~~~~~~
+.. autofunction:: tree_set
+
 Tree subtract
 ~~~~~~~~~~~~~
 .. autofunction:: tree_sub

optax/tree_utils/__init__.py

@@ -17,6 +17,7 @@
 from optax.tree_utils._state_utils import tree_get
 from optax.tree_utils._state_utils import tree_get_all_with_path
 from optax.tree_utils._state_utils import tree_map_params
+from optax.tree_utils._state_utils import tree_set
 
 from optax.tree_utils._tree_math import tree_add
 from optax.tree_utils._tree_math import tree_add_scalar_mul

optax/tree_utils/_state_utils.py

@@ -175,23 +175,33 @@ def tree_get_all_with_path(
       (``path_to_value``, ``value``). Here ``value`` is one entry of the state
       that corresponds to the ``key``, and ``path_to_value`` is a path returned
       by :func:`jax.tree_util.tree_flatten_with_path`.
+
+  Raises:
+    ValueError: If the input tree is flat, i.e., it is not a tuple/list/dict.
   """
-  values_with_path_found = []
+  found_values_with_path = []
   tree_flatten_with_path, _ = jax.tree_util.tree_flatten_with_path(tree)
+  if not tree_flatten_with_path or not tree_flatten_with_path[0][0]:
+    raise ValueError(
+        "The input tree cannot be flat, i.e., it must be a tuple/list/dict."
+    )
   for path, val in tree_flatten_with_path:
     key_leaf = _convert_jax_key_fn(path[-1])
     if key_leaf == key:
-      values_with_path_found.append((path, val))
-  return values_with_path_found
+      found_values_with_path.append((path, val))
+  return found_values_with_path
 
 
 def tree_get(tree: base.PyTree, key: Any, default: Optional[Any] = None) -> Any:
   """Extract a value from leaves of a pytree matching a given key.
 
   Search in the leaves of a pytree for a specific ``key`` (which can be a key
   from a dictionary or a name from a NamedTuple).
+
   If no leaves in the tree have the required ``key`` returns ``default``.
 
+  Raises a ``KeyError`` if multiple values of ``key`` are found in ``tree``.
+
   .. seealso:: :func:`optax.tree_utils.tree_get_all_with_path`
 
   Examples:
@@ -217,14 +227,68 @@ def tree_get(tree: base.PyTree, key: Any, default: Optional[Any] = None) -> Any:
 
   Raises:
     KeyError: If multiple values of ``key`` are found in ``tree``.
+    ValueError: If the input tree is flat, i.e., it is not a tuple/list/dict.
   """
-  values_with_path_found = tree_get_all_with_path(tree, key)
-  if len(values_with_path_found) > 1:
+  found_values_with_path = tree_get_all_with_path(tree, key)
+  if len(found_values_with_path) > 1:
     raise KeyError(f"Found multiple values for '{key}' in {tree}.")
-  elif not values_with_path_found:
+  elif not found_values_with_path:
     return default
   else:
-    return values_with_path_found[0][1]
+    return found_values_with_path[0][1]
+
+
+def tree_set(tree: base.PyTree, **kwargs: Any) -> base.PyTree:
+  """Creates a copy of tree with some leaves replaced as specified by kwargs.
+
+  Raises a ``KeyError`` if some keys in ``kwargs`` are not present in the tree.
+
+  Examples:
+    >>> import jax.numpy as jnp
+    >>> import optax
+    >>> params = jnp.array([1., 2., 3.])
+    >>> opt = optax.inject_hyperparams(optax.adam)(learning_rate=1.)
+    >>> state = opt.init(params)
+    >>> new_state = optax.tree_utils.tree_set(state, learning_rate=2.)
+    >>> lr = optax.tree_utils.tree_get(new_state, 'learning_rate')
+    >>> print(lr)
+    2.0
+
+  Args:
+    tree: pytree whose values are to be replaced.
+    **kwargs: dictionary of keys with values to replace in the tree.
+
+  Returns:
+    new_tree
+      new pytree with the same structure as tree. For each leaf whose
+      key/name matches a key in ``**kwargs``, their values are set by the
+      corresponding value in ``**kwargs``.
+
+  Raises:
+    KeyError: If no values of some key in ``**kwargs`` are found in ``tree``.
+    ValueError: If the input tree is flat, i.e., it is not a tuple/list/dict.
+  """
+  tree_flatten_with_path, _ = jax.tree_util.tree_flatten_with_path(tree)
+  if not tree_flatten_with_path or not tree_flatten_with_path[0][0]:
+    raise ValueError(
+        "The input tree cannot be flat, i.e., it must be a tuple/list/dict."
+    )
+  key_leaves = [
+      _convert_jax_key_fn(path[-1]) for path, _ in tree_flatten_with_path
+  ]
+  if (left_keys := set(kwargs) - set(key_leaves)):
+    left_keys_str = " nor ".join({f"'{key}'" for key in left_keys})
+    raise KeyError(f"Found no value for {left_keys_str} in {tree}.")
+
+  def _replace(path, value):
+    """Replace a value in tree if key from path matches some key in kwargs."""
+    key_leaf = _convert_jax_key_fn(path[-1])
+    if key_leaf in kwargs:
+      return kwargs[key_leaf]
+    else:
+      return value
+
+  return jax.tree_util.tree_map_with_path(_replace, tree)
 
 
 @jax.tree_util.register_pytree_node_class

optax/tree_utils/_state_utils_test.py

@@ -244,6 +244,12 @@ def test_map_non_params_to_none(self):
   def test_tree_get_all_with_path(self):
     params = jnp.array([1.0, 2.0, 3.0])
 
+    with self.subTest('Test with flat tree'):
+      tree = ()
+      self.assertRaises(ValueError, _state_utils.tree_get, tree, 'foo')
+      tree = jnp.array([1.0, 2.0, 3.0])
+      self.assertRaises(ValueError, _state_utils.tree_get, tree, 'foo')
+
     with self.subTest('Test with single value in state'):
       key = 'count'
       opt = transform.scale_by_adam()
@@ -253,8 +259,8 @@ def test_tree_get_all_with_path(self):
       self.assertEqual(values_found, expected_result)
 
     with self.subTest('Test with no value in state'):
-      key = 'count'
-      opt = alias.sgd(learning_rate=1.0)
+      key = 'apple'
+      opt = alias.adam(learning_rate=1.0)
       state = opt.init(params)
       values_found = _state_utils.tree_get_all_with_path(state, key)
       self.assertEmpty(values_found)
@@ -318,7 +324,7 @@ def test_tree_get(self):
 
     with self.subTest('Test jitted tree_get'):
       opt = _inject.inject_hyperparams(alias.sgd)(
-          learning_rate=lambda x: 1/(x+1)
+          learning_rate=lambda x: 1 / (x + 1)
       )
       state = opt.init(params)
 
@@ -327,10 +333,64 @@ def get_learning_rate(state):
         return _state_utils.tree_get(state, 'learning_rate')
 
       for i in range(4):
-         # we simply update state, we don't care about updates.
+        # we simply update state, we don't care about updates.
         _, state = opt.update(params, state)
         lr = get_learning_rate(state)
-        self.assertEqual(lr, 1/(i+1))
+        self.assertEqual(lr, 1 / (i + 1))
+
+  def test_tree_set(self):
+    params = jnp.array([1.0, 2.0, 3.0])
+
+    with self.subTest('Test with flat tree'):
+      tree = ()
+      self.assertRaises(ValueError, _state_utils.tree_get, tree, 'foo')
+      tree = jnp.array([1.0, 2.0, 3.0])
+      self.assertRaises(ValueError, _state_utils.tree_get, tree, 'foo')
+
+    with self.subTest('Test modifying an injected hyperparam'):
+      opt = _inject.inject_hyperparams(alias.adam)(learning_rate=1.0)
+      state = opt.init(params)
+      new_state = _state_utils.tree_set(state, learning_rate=2.0, b1=3.0)
+      lr = _state_utils.tree_get(new_state, 'learning_rate')
+      self.assertEqual(lr, 2.0)
+
+    with self.subTest('Test modifying an attribute of the state'):
+      opt = _inject.inject_hyperparams(alias.adam)(learning_rate=1.0)
+      state = opt.init(params)
+      new_state = _state_utils.tree_set(state, learning_rate=2.0, b1=3.0)
+      b1 = _state_utils.tree_get(new_state, 'b1')
+      self.assertEqual(b1, 3.0)
+
+    with self.subTest('Test modifying a value not present in the state'):
+      opt = _inject.inject_hyperparams(alias.adam)(learning_rate=1.0)
+      state = opt.init(params)
+      self.assertRaises(KeyError, _state_utils.tree_set, state, ema=2.0)
+
+    with self.subTest('Test jitted tree_set'):
+
+      @jax.jit
+      def set_learning_rate(state, lr):
+        return _state_utils.tree_set(state, learning_rate=lr)
+
+      modified_state = state
+      lr = 1.0
+      for i in range(4):
+        modified_state = set_learning_rate(modified_state, lr / (i + 1))
+        # we simply update state, we don't care about updates.
+        _, modified_state = opt.update(params, modified_state)
+        modified_lr = _state_utils.tree_get(modified_state, 'learning_rate')
+        self.assertEqual(modified_lr, lr / (i + 1))
+
+    with self.subTest('Test modifying several values at once'):
+      opt = combine.chain(
+          alias.adam(learning_rate=1.0), alias.adam(learning_rate=1.0)
+      )
+      state = opt.init(params)
+      new_state = _state_utils.tree_set(state, count=2.0)
+      values_found = _state_utils.tree_get_all_with_path(new_state, 'count')
+      self.assertLen(values_found, 2)
+      for _, value in values_found:
+        self.assertEqual(value, 2.0)
 
 
 def _fake_params():