class_docstr: Handles updating the state for a `lift_with_state` computation.
class_docstr_tok: ["Handles","updating","the","state","for","a","`","lift_with_state","`","computation","."]
docstr: Updates Haiku's internal state to the given state.
docstr_tok: ["Updates","Haiku","'s","internal","state","to","the","given","state","."]
code: def update(self, state: hk.State):
      """Updates Haiku's internal state to the given state."""
      frame = base.current_frame()
      for mod_name, bundle in state.items():
          if self._name is not None:
              mod_name = f"{self._name}/{mod_name}"
          for name, value in bundle.items():
              initial_pair = base.StatePair(value, value)
              initial = (
                  frame.state[mod_name].get(name, initial_pair).initial
              )
              frame.state[mod_name][name] = base.StatePair(
                  initial, value
              )
  
imports: from collections.abc import Mapping, MutableMapping
  import functools
  from typing import Any, Callable, Optional, TypeVar
  from haiku._src import base
  from haiku._src import data_structures
  from haiku._src import module
  from haiku._src import transform
  from haiku._src.typing import LiftingModuleType
label_desc: Use image node_id 4 for calling the LiftWithStateUpdater obj's underlying member method code with example usage: obj.update(state) without return types