# Path: vmoe/utils.py
# def partialclass(cls: Type[Any], *base_args, **base_kwargs):
#     def __init__(self, *args, **kwargs):
#   def __init__(self, *iterators):
#   def __iter__(self):
#   def __next__(self) -> Tuple[Any, ...]:
# def safe_map(f: Callable[..., Any], *iterables) -> Iterator[Any]:
# def safe_zip(*iterables) -> Iterator[Tuple[Any, ...]]:
# def tree_rngs_split(rngs, num_splits=2):
# def tree_shape_dtype_struct(tree):
#   def fn(x):
#   class _NewClass(cls):
# class SafeZipIteratorError(RuntimeError):
# class SafeZipIterator:


import dataclasses
from absl.testing import absltest
from vmoe import utils
class UtilsTest(absltest.TestCase):

  def test_partialclass(self):

    @dataclasses.dataclass
    class A:
      a: int
      b: int

    cls = utils.partialclass(A, a=1)
    self.assertTrue(issubclass(cls, A))
    obj = cls(b=2)
    self.assertIsInstance(obj, A)
    self.assertEqual(obj.a, 1)
    self.assertEqual(obj.b, 2)

  def test_safe_map(self):
    self.assertEqual(
        list(utils.safe_map(lambda x, y: x + y, [1, 3], [2, 4])),
        [1 + 2, 3 + 4])

  def test_safe_zip(self):
    self.assertIsInstance(utils.safe_zip([], []), utils.SafeZipIterator)
    self.assertEqual(list(utils.safe_zip([], [])), [])
    self.assertEqual(list(utils.safe_zip([1])), [(1,)])
    self.assertEqual(list(utils.safe_zip([1, 2], ['a', 'b'])),
                     [(1, 'a'), (2, 'b')])
    with self.assertRaises(utils.SafeZipIteratorError):
      list(utils.safe_zip([], ['a'], [1.]))
    with self.assertRaises(utils.SafeZipIteratorError):

# Path: vmoe/utils.py
# def partialclass(cls: Type[Any], *base_args, **base_kwargs):
#     def __init__(self, *args, **kwargs):
#   def __init__(self, *iterators):
#   def __iter__(self):
#   def __next__(self) -> Tuple[Any, ...]:
# def safe_map(f: Callable[..., Any], *iterables) -> Iterator[Any]:
# def safe_zip(*iterables) -> Iterator[Tuple[Any, ...]]:
# def tree_rngs_split(rngs, num_splits=2):
# def tree_shape_dtype_struct(tree):
#   def fn(x):
#   class _NewClass(cls):
# class SafeZipIteratorError(RuntimeError):
# class SafeZipIterator:


# Path: vmoe/utils_test.py
import dataclasses
from absl.testing import absltest
from vmoe import utils

class UtilsTest(absltest.TestCase):

  def test_partialclass(self):

    @dataclasses.dataclass
    class A:
      a: int
      b: int

    cls = utils.partialclass(A, a=1)
    self.assertTrue(issubclass(cls, A))
    obj = cls(b=2)
    self.assertIsInstance(obj, A)
    self.assertEqual(obj.a, 1)
    self.assertEqual(obj.b, 2)

  def test_safe_map(self):
    self.assertEqual(
        list(utils.safe_map(lambda x, y: x + y, [1, 3], [2, 4])),
        [1 + 2, 3 + 4])

  def test_safe_zip(self):
    self.assertIsInstance(utils.safe_zip([], []), utils.SafeZipIterator)
    self.assertEqual(list(utils.safe_zip([], [])), [])
    self.assertEqual(list(utils.safe_zip([1])), [(1,)])
    self.assertEqual(list(utils.safe_zip([1, 2], ['a', 'b'])),
                     [(1, 'a'), (2, 'b')])
    with self.assertRaises(utils.SafeZipIteratorError):
      list(utils.safe_zip([], ['a'], [1.]))
    with self.assertRaises(utils.SafeZipIteratorError):

      list(utils.safe_zip([1], [], [1.]))