# Path: rigl/experimental/jax/datasets/cifar10.py
# class CIFAR10Dataset(dataset_base.ImageDataset):
#   NAME: str = 'cifar10'
#   MEAN_RGB: Sequence[float] = [0.4914 * 255, 0.4822 * 255, 0.4465 * 255]
#   STDDEV_RGB: Sequence[float] = [0.2470 * 255, 0.2435 * 255, 0.2616 * 255]
#   def __init__(self,
#                batch_size,
#                batch_size_test,
#                shuffle_buffer_size = 1024,
#                seed = 42):
#   def preprocess(
#       self, data):
# 
# Path: rigl/experimental/jax/datasets/dataset_base.py
# class Dataset(metaclass=abc.ABCMeta):
# class ImageDataset(Dataset):
#   DATAKEY: Optional[str] = None
#   LABELKEY: str = 'label'
#   DATAKEY = 'image'
#   def __init__(self,
#                name,
#                batch_size,
#                batch_size_test,
#                shuffle_buffer_size,
#                prefetch_size = 1,
#                seed = None):  # pytype: disable=annotation-type-mismatch
#   def _dataset_dir(self):
#   def get_train(self):
#   def get_train_len(self):
#   def get_test(self):
#   def get_test_len(self):
#   def preprocess(
#       self, data):
#   def augment(
#       self, data):
#   def preprocess(
#       self, data):
# 
# Path: rigl/experimental/jax/datasets/mnist.py
# class MNISTDataset(dataset_base.ImageDataset):
#   NAME: str = 'mnist'
#   def __init__(self,
#                batch_size,
#                batch_size_test,
#                shuffle_buffer_size = 1024,
#                seed = 42):
#   def preprocess(
#       self, data):


from typing import Any, Mapping, Type
from rigl.experimental.jax.datasets import cifar10
from rigl.experimental.jax.datasets import dataset_base
from rigl.experimental.jax.datasets import mnist
import tensorflow.compat.v2 as tf
# coding=utf-8
# Copyright 2022 RigL Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# 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.

# Lint as: python3
"""Dataset Factory.

Dataset factory to allow us to easily use tensorflow datasets (TFDS)
with JAX/FLAX, by defining a bunch of wrappers, including preprocessing.

Attributes:
  DATASETS: A list of the datasets that can be created.
"""




# Path: rigl/experimental/jax/datasets/cifar10.py
# class CIFAR10Dataset(dataset_base.ImageDataset):
#   NAME: str = 'cifar10'
#   MEAN_RGB: Sequence[float] = [0.4914 * 255, 0.4822 * 255, 0.4465 * 255]
#   STDDEV_RGB: Sequence[float] = [0.2470 * 255, 0.2435 * 255, 0.2616 * 255]
#   def __init__(self,
#                batch_size,
#                batch_size_test,
#                shuffle_buffer_size = 1024,
#                seed = 42):
#   def preprocess(
#       self, data):
# 
# Path: rigl/experimental/jax/datasets/dataset_base.py
# class Dataset(metaclass=abc.ABCMeta):
# class ImageDataset(Dataset):
#   DATAKEY: Optional[str] = None
#   LABELKEY: str = 'label'
#   DATAKEY = 'image'
#   def __init__(self,
#                name,
#                batch_size,
#                batch_size_test,
#                shuffle_buffer_size,
#                prefetch_size = 1,
#                seed = None):  # pytype: disable=annotation-type-mismatch
#   def _dataset_dir(self):
#   def get_train(self):
#   def get_train_len(self):
#   def get_test(self):
#   def get_test_len(self):
#   def preprocess(
#       self, data):
#   def augment(
#       self, data):
#   def preprocess(
#       self, data):
# 
# Path: rigl/experimental/jax/datasets/mnist.py
# class MNISTDataset(dataset_base.ImageDataset):
#   NAME: str = 'mnist'
#   def __init__(self,
#                batch_size,
#                batch_size_test,
#                shuffle_buffer_size = 1024,
#                seed = 42):
#   def preprocess(
#       self, data):


# Path: rigl/experimental/jax/datasets/dataset_factory.py
from typing import Any, Mapping, Type
from rigl.experimental.jax.datasets import cifar10
from rigl.experimental.jax.datasets import dataset_base
from rigl.experimental.jax.datasets import mnist
import tensorflow.compat.v2 as tf

# coding=utf-8
# Copyright 2022 RigL Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# 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.

# Lint as: python3
"""Dataset Factory.

Dataset factory to allow us to easily use tensorflow datasets (TFDS)
with JAX/FLAX, by defining a bunch of wrappers, including preprocessing.

Attributes:
  DATASETS: A list of the datasets that can be created.
"""




DATASETS: Mapping[str, Type[dataset_base.Dataset]] = {