# Path: fedjax/core/serialization.py
# def save_state(state, path):
# def load_state(path):
# def _ndarray_to_bytes(arr):
# def _dtype_from_name(name):
# def _ndarray_from_bytes(data):
# def _bytes_ndarray_to_bytes(x):
# def _object_ndarray_from_bytes(data):
# def _msgpack_ext_pack(x):
# def _msgpack_ext_unpack(code, data):
# def msgpack_serialize(pytree):
# def msgpack_deserialize(encoded_pytree):
# class _MsgpackExtType(enum.IntEnum):
# 
# Path: fedjax/core/sqlite_federated_data.py
# def decompress_and_deserialize(data: bytes):
#   def new(
#       path: str,
#       parse_examples: Callable[
#           [bytes], client_datasets.Examples] = decompress_and_deserialize
#   ) -> 'SQLiteFederatedData':
#   def __init__(self,
#                connection: sqlite3.Connection,
#                parse_examples: Callable[[bytes], client_datasets.Examples],
#                start: Optional[federated_data.ClientId] = None,
#                stop: Optional[federated_data.ClientId] = None,
#                preprocess_client: federated_data
#                .ClientPreprocessor = federated_data.NoOpClientPreprocessor,
#                preprocess_batch: client_datasets
#                .BatchPreprocessor = client_datasets.NoOpBatchPreprocessor):
#   def slice(
#       self,
#       start: Optional[federated_data.ClientId] = None,
#       stop: Optional[federated_data.ClientId] = None) -> 'SQLiteFederatedData':
#   def preprocess_client(
#       self, fn: Callable[[federated_data.ClientId, client_datasets.Examples],
#                          client_datasets.Examples]
#   ) -> 'SQLiteFederatedData':
#   def preprocess_batch(
#       self, fn: Callable[[client_datasets.Examples], client_datasets.Examples]
#   ) -> 'SQLiteFederatedData':
#   def _range_where(self) -> str:
#   def num_clients(self) -> int:
#   def client_ids(self) -> Iterator[federated_data.ClientId]:
#   def client_sizes(self) -> Iterator[Tuple[federated_data.ClientId, int]]:
#   def client_size(self, client_id: federated_data.ClientId) -> int:
#   def clients(
#       self
#   ) -> Iterator[Tuple[federated_data.ClientId, client_datasets.ClientDataset]]:
#   def _read_clients(self):
#   def shuffled_clients(
#       self,
#       buffer_size: int,
#       seed: Optional[int] = None
#   ) -> Iterator[Tuple[federated_data.ClientId, client_datasets.ClientDataset]]:
#   def get_clients(
#       self, client_ids: Iterable[federated_data.ClientId]
#   ) -> Iterator[Tuple[federated_data.ClientId, client_datasets.ClientDataset]]:
#   def get_client(
#       self,
#       client_id: federated_data.ClientId) -> client_datasets.ClientDataset:
#   def _client_dataset(self, client_id: federated_data.ClientId,
#                       data: bytes) -> client_datasets.ClientDataset:
#   def __init__(self, path: str):
#   def __enter__(self):
#   def __exit__(self, exc_type, exc_value, exc_traceback):
#   def add_many(self,
#                client_ids_examples: Iterable[Tuple[bytes,
#                                                    client_datasets.Examples]]):
#     def prepare_parameters(ce):
#   def __init__(self, path: str,
#                parse_examples: Callable[[List[bytes]],
#                                         client_datasets.Examples],
#                split_name: str):
#   def __del__(self):
#   def __iter__(
#       self
#   ) -> Iterator[Tuple[federated_data.ClientId, client_datasets.ClientDataset]]:
#   def __next__(
#       self) -> Tuple[federated_data.ClientId, client_datasets.ClientDataset]:
# class SQLiteFederatedData(federated_data.FederatedData):
# class SQLiteFederatedDataBuilder(federated_data.FederatedDataBuilder):
# class TFFSQLiteClientsIterator:
# 
# Path: fedjax/core/util.py
# def safe_div(a, b):
#   def __getattr__(self, _):
# def import_tf():
# class FakeTensorFlowModule:


import itertools
import os.path
import zlib
import numpy as np
import numpy.testing as npt
import sqlite3
from absl import flags
from absl.testing import absltest
from fedjax.core import serialization
from fedjax.core import sqlite_federated_data
from fedjax.core import util
        ON client_metadata (client_id);""")

      for i in range(100):
        client_id = str(i)
        num_examples = i + 1
        for j in range(num_examples):
          tf_example = tf.train.Example(
              features=tf.train.Features(feature={
                  'x':
                      tf.train.Feature(
                          int64_list=tf.train.Int64List(value=[j]))
              }))
          serialized_example_proto = tf_example.SerializeToString()
          connection.execute('INSERT INTO examples VALUES (?, ?, ?);',
                             [split_name, client_id, serialized_example_proto])
        connection.execute('INSERT INTO client_metadata VALUES (?, ?, ?);',
                           [client_id, split_name, num_examples])
    connection.close()

    def parse_examples(vs):
      tf_examples = tf.io.parse_example(
          vs, features={'x': tf.io.FixedLenFeature(shape=(), dtype=tf.int64)})
      return tf.nest.map_structure(lambda t: t.numpy(), tf_examples)

    cls._path = path
    cls._parse_examples = parse_examples
    cls._split_name = split_name

  def test_iteration(self):
    client_ids = set()

# Path: fedjax/core/serialization.py
# def save_state(state, path):
# def load_state(path):
# def _ndarray_to_bytes(arr):
# def _dtype_from_name(name):
# def _ndarray_from_bytes(data):
# def _bytes_ndarray_to_bytes(x):
# def _object_ndarray_from_bytes(data):
# def _msgpack_ext_pack(x):
# def _msgpack_ext_unpack(code, data):
# def msgpack_serialize(pytree):
# def msgpack_deserialize(encoded_pytree):
# class _MsgpackExtType(enum.IntEnum):
# 
# Path: fedjax/core/sqlite_federated_data.py
# def decompress_and_deserialize(data: bytes):
#   def new(
#       path: str,
#       parse_examples: Callable[
#           [bytes], client_datasets.Examples] = decompress_and_deserialize
#   ) -> 'SQLiteFederatedData':
#   def __init__(self,
#                connection: sqlite3.Connection,
#                parse_examples: Callable[[bytes], client_datasets.Examples],
#                start: Optional[federated_data.ClientId] = None,
#                stop: Optional[federated_data.ClientId] = None,
#                preprocess_client: federated_data
#                .ClientPreprocessor = federated_data.NoOpClientPreprocessor,
#                preprocess_batch: client_datasets
#                .BatchPreprocessor = client_datasets.NoOpBatchPreprocessor):
#   def slice(
#       self,
#       start: Optional[federated_data.ClientId] = None,
#       stop: Optional[federated_data.ClientId] = None) -> 'SQLiteFederatedData':
#   def preprocess_client(
#       self, fn: Callable[[federated_data.ClientId, client_datasets.Examples],
#                          client_datasets.Examples]
#   ) -> 'SQLiteFederatedData':
#   def preprocess_batch(
#       self, fn: Callable[[client_datasets.Examples], client_datasets.Examples]
#   ) -> 'SQLiteFederatedData':
#   def _range_where(self) -> str:
#   def num_clients(self) -> int:
#   def client_ids(self) -> Iterator[federated_data.ClientId]:
#   def client_sizes(self) -> Iterator[Tuple[federated_data.ClientId, int]]:
#   def client_size(self, client_id: federated_data.ClientId) -> int:
#   def clients(
#       self
#   ) -> Iterator[Tuple[federated_data.ClientId, client_datasets.ClientDataset]]:
#   def _read_clients(self):
#   def shuffled_clients(
#       self,
#       buffer_size: int,
#       seed: Optional[int] = None
#   ) -> Iterator[Tuple[federated_data.ClientId, client_datasets.ClientDataset]]:
#   def get_clients(
#       self, client_ids: Iterable[federated_data.ClientId]
#   ) -> Iterator[Tuple[federated_data.ClientId, client_datasets.ClientDataset]]:
#   def get_client(
#       self,
#       client_id: federated_data.ClientId) -> client_datasets.ClientDataset:
#   def _client_dataset(self, client_id: federated_data.ClientId,
#                       data: bytes) -> client_datasets.ClientDataset:
#   def __init__(self, path: str):
#   def __enter__(self):
#   def __exit__(self, exc_type, exc_value, exc_traceback):
#   def add_many(self,
#                client_ids_examples: Iterable[Tuple[bytes,
#                                                    client_datasets.Examples]]):
#     def prepare_parameters(ce):
#   def __init__(self, path: str,
#                parse_examples: Callable[[List[bytes]],
#                                         client_datasets.Examples],
#                split_name: str):
#   def __del__(self):
#   def __iter__(
#       self
#   ) -> Iterator[Tuple[federated_data.ClientId, client_datasets.ClientDataset]]:
#   def __next__(
#       self) -> Tuple[federated_data.ClientId, client_datasets.ClientDataset]:
# class SQLiteFederatedData(federated_data.FederatedData):
# class SQLiteFederatedDataBuilder(federated_data.FederatedDataBuilder):
# class TFFSQLiteClientsIterator:
# 
# Path: fedjax/core/util.py
# def safe_div(a, b):
#   def __getattr__(self, _):
# def import_tf():
# class FakeTensorFlowModule:


# Path: fedjax/core/sqlite_federated_data_test.py
import itertools
import os.path
import zlib
import numpy as np
import numpy.testing as npt
import sqlite3
from absl import flags
from absl.testing import absltest
from fedjax.core import serialization
from fedjax.core import sqlite_federated_data
from fedjax.core import util

        ON client_metadata (client_id);""")

      for i in range(100):
        client_id = str(i)
        num_examples = i + 1
        for j in range(num_examples):
          tf_example = tf.train.Example(
              features=tf.train.Features(feature={
                  'x':
                      tf.train.Feature(
                          int64_list=tf.train.Int64List(value=[j]))
              }))
          serialized_example_proto = tf_example.SerializeToString()
          connection.execute('INSERT INTO examples VALUES (?, ?, ?);',
                             [split_name, client_id, serialized_example_proto])
        connection.execute('INSERT INTO client_metadata VALUES (?, ?, ?);',
                           [client_id, split_name, num_examples])
    connection.close()

    def parse_examples(vs):
      tf_examples = tf.io.parse_example(
          vs, features={'x': tf.io.FixedLenFeature(shape=(), dtype=tf.int64)})
      return tf.nest.map_structure(lambda t: t.numpy(), tf_examples)

    cls._path = path
    cls._parse_examples = parse_examples
    cls._split_name = split_name

  def test_iteration(self):
    client_ids = set()

    for client_id, client_dataset in sqlite_federated_data.TFFSQLiteClientsIterator(