# 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
# 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 utils."""



def create_fake_parameters():
  qkv = 20
  p_length = 2
  embed_dim = 200
  layers = 4
  parameters = {}
  for i in range(layers):
    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']

# 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

# 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 utils."""



def create_fake_parameters():
  qkv = 20
  p_length = 2
  embed_dim = 200
  layers = 4
  parameters = {}
  for i in range(layers):
    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))