# Path: t5x/state_utils.py
# def tensorstore_leaf(_, value):
# def flatten_state_dict(state_dict, keep_empty_nodes: bool = False):
# def get_name_tree(state_dict, keep_empty_nodes: bool = False):
# def intersect_state(
#     state_dict: Mapping[str, Any],
#     intersect_state_dict: Mapping[str, Any]) -> Mapping[str, Any]:
# def merge_state(state_dict: Mapping[str, Any],
#                 from_scratch_state: Mapping[str, Any]) -> Mapping[str, Any]:
# def apply_assignment_map(ckpt_optimizer_state,
#                          optimizer_state,
#                          assignment_map: Sequence[Tuple[str, Optional[str]]],
#                          require_all_rules_match: bool = True,
#                          *,
#                          is_resuming: bool = False):


import re
import numpy as np
from absl.testing import absltest
from absl.testing import parameterized
from t5x import state_utils
          }}),
  )
  def test_apply_assignment_map_partial_initialization(self,
                                                       ckpt_optimizer_state,
                                                       optimizer_state,
                                                       assignment_map,
                                                       expect_state):
    assignment_map = [(re.compile(k), v) for (k, v) in assignment_map]
    actual_state = state_utils.apply_assignment_map(
        ckpt_optimizer_state=ckpt_optimizer_state,
        optimizer_state=optimizer_state,
        assignment_map=assignment_map,
        require_all_rules_match=True)
    self.assertEqual(actual_state, expect_state)

  def test_get_name_tree(self):
    state_dict = {"a": {"b": {"c": 0, "d": {"e": 1}}}, "f": {}}

    self.assertEqual(
        state_utils.get_name_tree(state_dict),
        {"a": {
            "b": {
                "c": "a/b/c",
                "d": {
                    "e": "a/b/d/e"
                }
            }
        }})

    self.assertEqual(

# Path: t5x/state_utils.py
# def tensorstore_leaf(_, value):
# def flatten_state_dict(state_dict, keep_empty_nodes: bool = False):
# def get_name_tree(state_dict, keep_empty_nodes: bool = False):
# def intersect_state(
#     state_dict: Mapping[str, Any],
#     intersect_state_dict: Mapping[str, Any]) -> Mapping[str, Any]:
# def merge_state(state_dict: Mapping[str, Any],
#                 from_scratch_state: Mapping[str, Any]) -> Mapping[str, Any]:
# def apply_assignment_map(ckpt_optimizer_state,
#                          optimizer_state,
#                          assignment_map: Sequence[Tuple[str, Optional[str]]],
#                          require_all_rules_match: bool = True,
#                          *,
#                          is_resuming: bool = False):


# Path: t5x/state_utils_test.py
import re
import numpy as np
from absl.testing import absltest
from absl.testing import parameterized
from t5x import state_utils

          }}),
  )
  def test_apply_assignment_map_partial_initialization(self,
                                                       ckpt_optimizer_state,
                                                       optimizer_state,
                                                       assignment_map,
                                                       expect_state):
    assignment_map = [(re.compile(k), v) for (k, v) in assignment_map]
    actual_state = state_utils.apply_assignment_map(
        ckpt_optimizer_state=ckpt_optimizer_state,
        optimizer_state=optimizer_state,
        assignment_map=assignment_map,
        require_all_rules_match=True)
    self.assertEqual(actual_state, expect_state)

  def test_get_name_tree(self):
    state_dict = {"a": {"b": {"c": 0, "d": {"e": 1}}}, "f": {}}

    self.assertEqual(
        state_utils.get_name_tree(state_dict),
        {"a": {
            "b": {
                "c": "a/b/c",
                "d": {
                    "e": "a/b/d/e"
                }
            }
        }})

    self.assertEqual(

        state_utils.get_name_tree(state_dict, keep_empty_nodes=True), {