




from typing import List, Optional, TypeVar

from .arrow_dataset import Dataset, _concatenate_map_style_datasets, _interleave_map_style_datasets
from .dataset_dict import DatasetDict, IterableDatasetDict
from .info import DatasetInfo
from .iterable_dataset import IterableDataset, _concatenate_iterable_datasets, _interleave_iterable_datasets
from .splits import NamedSplit
from .utils import logging
from .utils.py_utils import Literal


_UpperCAmelCase	:						Optional[Any]			=     logging.get_logger(__name__)


_UpperCAmelCase	:						Any			=     TypeVar("DatasetType", Dataset, IterableDataset)



def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase = None,			lowerCamelCase = None,			lowerCamelCase = None,			lowerCamelCase = None,			lowerCamelCase = "first_exhausted",			):
  from .arrow_dataset import Dataset
  from .iterable_dataset import IterableDataset

  if not datasets:
    raise ValueError("Unable to interleave an empty list of datasets."   )
  for i, dataset in enumerate(lowerCamelCase   ):
    if not isinstance(lowerCamelCase,			(Dataset, IterableDataset)   ):
      if isinstance(lowerCamelCase,			(DatasetDict, IterableDatasetDict)   ):
        if not dataset:
          raise ValueError(
              F"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} "
              "is an empty dataset dictionary."   )
        raise ValueError(
            F"Dataset at position {i} has at least one split: {list(lowerCamelCase   )}\n"
            F"Please pick one to interleave with the other datasets, for example: dataset['{next(iter(lowerCamelCase   )   )}']"   )
      raise ValueError(
          F"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(lowerCamelCase   ).__name__}."   )
    if i == 0:
      lowercase	,    lowercase    :Any      	=						(
          (Dataset, IterableDataset) if isinstance(lowerCamelCase,			lowerCamelCase   ) else (IterableDataset, Dataset)
      )
    elif not isinstance(lowerCamelCase,			lowerCamelCase   ):
      raise ValueError(
          F"Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects."   )
  if stopping_strategy not in ["first_exhausted", "all_exhausted"]:
    raise ValueError(F"{stopping_strategy} is not supported. Please enter a valid stopping_strategy."   )
  if dataset_type is Dataset:
    return _interleave_map_style_datasets(
        lowerCamelCase,			lowerCamelCase,			lowerCamelCase,			info=lowerCamelCase,			split=lowerCamelCase,			stopping_strategy=lowerCamelCase   )
  else:
    return _interleave_iterable_datasets(
        lowerCamelCase,			lowerCamelCase,			lowerCamelCase,			info=lowerCamelCase,			split=lowerCamelCase,			stopping_strategy=lowerCamelCase   )







def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase = None,			lowerCamelCase = None,			lowerCamelCase = 0,			):

  if not dsets:
    raise ValueError("Unable to concatenate an empty list of datasets."   )
  for i, dataset in enumerate(lowerCamelCase   ):
    if not isinstance(lowerCamelCase,			(Dataset, IterableDataset)   ):
      if isinstance(lowerCamelCase,			(DatasetDict, IterableDatasetDict)   ):
        if not dataset:
          raise ValueError(
              F"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} "
              "is an empty dataset dictionary."   )
        raise ValueError(
            F"Dataset at position {i} has at least one split: {list(lowerCamelCase   )}\n"
            F"Please pick one to interleave with the other datasets, for example: dataset['{next(iter(lowerCamelCase   )   )}']"   )
      raise ValueError(
          F"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(lowerCamelCase   ).__name__}."   )
    if i == 0:
      lowercase	,    lowercase    :Dict      	=						(
          (Dataset, IterableDataset) if isinstance(lowerCamelCase,			lowerCamelCase   ) else (IterableDataset, Dataset)
      )
    elif not isinstance(lowerCamelCase,			lowerCamelCase   ):
      raise ValueError(
          F"Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects."   )
  if dataset_type is Dataset:
    return _concatenate_map_style_datasets(lowerCamelCase,			info=lowerCamelCase,			split=lowerCamelCase,			axis=lowerCamelCase   )
  else:
    return _concatenate_iterable_datasets(lowerCamelCase,			info=lowerCamelCase,			split=lowerCamelCase,			axis=lowerCamelCase   )









import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


_UpperCAmelCase	:						List[Any]			=     logging.get_logger(__name__)



def       UpperCAmelCase__  (     lowerCamelCase   ):
  lowercase    :int      	=						r"\w+[.]\d+"
  lowercase    :Tuple      	=						re.findall(lowerCamelCase,			lowerCamelCase   )
  for pat in pats:
    lowercase    :List[str]      	=						key.replace(lowerCamelCase,			"_".join(pat.split("."   )   )   )
  return key



def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase,			lowerCamelCase   ):
  lowercase    :List[str]      	=						pt_tuple_key[:-1] + ("scale",)
  if (
      any("norm" in str_ for str_ in pt_tuple_key   )
      and (pt_tuple_key[-1] == "bias")
      and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
      and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
  ):
    lowercase    :Optional[int]      	=						pt_tuple_key[:-1] + ("scale",)
    return renamed_pt_tuple_key, pt_tensor
  elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
    lowercase    :int      	=						pt_tuple_key[:-1] + ("scale",)
    return renamed_pt_tuple_key, pt_tensor

  # embedding
  if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
    lowercase    :List[str]      	=						pt_tuple_key[:-1] + ("embedding",)
    return renamed_pt_tuple_key, pt_tensor

  # conv layer
  lowercase    :Union[str, Any]      	=						pt_tuple_key[:-1] + ("kernel",)
  if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
    lowercase    :List[Any]      	=						pt_tensor.transpose(2,			3,			1,			0   )
    return renamed_pt_tuple_key, pt_tensor

  # linear layer
  lowercase    :Optional[Any]      	=						pt_tuple_key[:-1] + ("kernel",)
  if pt_tuple_key[-1] == "weight":
    lowercase    :Tuple      	=						pt_tensor.T
    return renamed_pt_tuple_key, pt_tensor

  # old PyTorch layer norm weight
  lowercase    :Tuple      	=						pt_tuple_key[:-1] + ("weight",)
  if pt_tuple_key[-1] == "gamma":
    return renamed_pt_tuple_key, pt_tensor

  # old PyTorch layer norm bias
  lowercase    :int      	=						pt_tuple_key[:-1] + ("bias",)
  if pt_tuple_key[-1] == "beta":
    return renamed_pt_tuple_key, pt_tensor

  return pt_tuple_key, pt_tensor







def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase,			lowerCamelCase=42   ):
  # Step 1: Convert pytorch tensor to numpy
  lowercase    :Optional[int]      	=						{k: v.numpy() for k, v in pt_state_dict.items()}

  # Step 2: Since the model is stateless, get random Flax params
  lowercase    :str      	=						flax_model.init_weights(PRNGKey(lowerCamelCase   )   )

  lowercase    :Tuple      	=						flatten_dict(lowerCamelCase   )
  lowercase    :Optional[Any]      	=						{}

  # Need to change some parameters name to match Flax names
  for pt_key, pt_tensor in pt_state_dict.items():
    lowercase    :List[Any]      	=						rename_key(lowerCamelCase   )
    lowercase    :List[Any]      	=						tuple(renamed_pt_key.split("."   )   )

    # Correctly rename weight parameters
    lowercase	,    lowercase    :List[str]      	=						rename_key_and_reshape_tensor(lowerCamelCase,			lowerCamelCase,			lowerCamelCase   )

    if flax_key in random_flax_state_dict:
      if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
        raise ValueError(
            F"PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape "
            F"{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."   )

        # also add unexpected weight so that warning is thrown
    lowercase    :List[str]      	=						jnp.asarray(lowerCamelCase   )

  return unflatten_dict(lowerCamelCase   )



