# 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
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for state_utils."""




class StateUtilsTest(parameterized.TestCase):

  @parameterized.parameters(
      dict(
          state_dict={"a": {
              "b": 2,
              "c": 3
          }},
          intersect_state_dict={
              "a": {
                  "b": 4
              },
              "d": 5
          },
          expect_state={"a": {
              "b": 2
          }}))
  def test_intersect_state(self, state_dict, intersect_state_dict,
                           expect_state):

# 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

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for state_utils."""




class StateUtilsTest(parameterized.TestCase):

  @parameterized.parameters(
      dict(
          state_dict={"a": {
              "b": 2,
              "c": 3
          }},
          intersect_state_dict={
              "a": {
                  "b": 4
              },
              "d": 5
          },
          expect_state={"a": {
              "b": 2
          }}))
  def test_intersect_state(self, state_dict, intersect_state_dict,
                           expect_state):

    actual_state = state_utils.intersect_state(state_dict, intersect_state_dict)