# Path: prompt_tuning/train/utils.py
# def match_any(regexes: Sequence[str]) -> Callable[[str, Any], bool]:
#   def _match_any(path, _):
# def np_save(path: str, arr) -> None:
#   def __init__(
#       self, traversals_and_optimizers: Sequence[Tuple[traverse_util.Traversal,
#                                                       optim.OptimizerDef]]):
#   def __init__(self, *args, save_paths: Sequence[str], **kwargs):
#   def save_numpy(
#       self, targets: flax.core.frozen_dict.FrozenDict, step: int) -> None:
#   def save(self,
#            train_state: train_state_lib.TrainState,
#            state_transformation_fns: Sequence[
#                checkpoints.SaveStateTransformationFn] = (),
#            *,
#            concurrent_gb: int = 128):
# class MultiOptimizer(optim.MultiOptimizer):
# class Checkpointer(checkpoints.Checkpointer):


import inspect
import operator
import os
import flax
import numpy as np
from unittest import mock
from absl.testing import absltest
from flax import traverse_util
from prompt_tuning.train import utils
    for proj in ('query', 'key', 'value', 'out'):
      parameters[f'encoder/layers_{i}/attention/'
                 f'{proj}/kernel'] = np.random.rand(qkv, qkv),
  parameters['encoder/prompt/prompt'] = np.random.rand(p_length, embed_dim)
  parameters = flax.core.freeze(traverse_util.unflatten_dict(
      {tuple(k.split('/')): v for k, v in parameters.items()}))
  return parameters


class UtilsTest(absltest.TestCase):

  def test_match_any(self):
    regexes = ['.*', 'example']
    self.assertTrue(utils.match_any(regexes)('example', None))

  def test_match_any_closure(self):
    regexes = ['.*']
    self.assertTrue(callable(utils.match_any(regexes)))

  def test_match_any_full_mach(self):
    regexes = ['dogs']
    # This shouldn't match because we have extra text at the start.
    parameter = 'i/love/dogs'
    match_any = utils.match_any(regexes)
    self.assertFalse(match_any(parameter, None))

  def test_np_save(self):
    write_array = np.reshape(np.arange(20), (4, -1))
    test_file = self.create_tempfile('written_array.npy')
    test_file_name = test_file.full_path

# Path: prompt_tuning/train/utils.py
# def match_any(regexes: Sequence[str]) -> Callable[[str, Any], bool]:
#   def _match_any(path, _):
# def np_save(path: str, arr) -> None:
#   def __init__(
#       self, traversals_and_optimizers: Sequence[Tuple[traverse_util.Traversal,
#                                                       optim.OptimizerDef]]):
#   def __init__(self, *args, save_paths: Sequence[str], **kwargs):
#   def save_numpy(
#       self, targets: flax.core.frozen_dict.FrozenDict, step: int) -> None:
#   def save(self,
#            train_state: train_state_lib.TrainState,
#            state_transformation_fns: Sequence[
#                checkpoints.SaveStateTransformationFn] = (),
#            *,
#            concurrent_gb: int = 128):
# class MultiOptimizer(optim.MultiOptimizer):
# class Checkpointer(checkpoints.Checkpointer):


# Path: prompt_tuning/train/utils_test.py
import inspect
import operator
import os
import flax
import numpy as np
from unittest import mock
from absl.testing import absltest
from flax import traverse_util
from prompt_tuning.train import utils

    for proj in ('query', 'key', 'value', 'out'):
      parameters[f'encoder/layers_{i}/attention/'
                 f'{proj}/kernel'] = np.random.rand(qkv, qkv),
  parameters['encoder/prompt/prompt'] = np.random.rand(p_length, embed_dim)
  parameters = flax.core.freeze(traverse_util.unflatten_dict(
      {tuple(k.split('/')): v for k, v in parameters.items()}))
  return parameters


class UtilsTest(absltest.TestCase):

  def test_match_any(self):
    regexes = ['.*', 'example']
    self.assertTrue(utils.match_any(regexes)('example', None))

  def test_match_any_closure(self):
    regexes = ['.*']
    self.assertTrue(callable(utils.match_any(regexes)))

  def test_match_any_full_mach(self):
    regexes = ['dogs']
    # This shouldn't match because we have extra text at the start.
    parameter = 'i/love/dogs'
    match_any = utils.match_any(regexes)
    self.assertFalse(match_any(parameter, None))

  def test_np_save(self):
    write_array = np.reshape(np.arange(20), (4, -1))
    test_file = self.create_tempfile('written_array.npy')
    test_file_name = test_file.full_path

    utils.np_save(test_file_name, write_array)