# Path: t5x/checkpoint_importer.py
# class LazyArray(abc.ABC):
# class LazyThreadPoolArray(LazyArray):
# class LazyAwaitableArray(LazyArray):
# class CheckpointTranslator:
#   def __init__(self, shape: Sequence[int], dtype: jnp.dtype,
#                get_fn: Callable[[], np.ndarray]):
#   def shape(self) -> Tuple[int, ...]:
#   def dtype(self) -> jnp.dtype:
#   def nbytes(self) -> int:
#   def astype(self, dtype: np.dtype) -> 'LazyArray':
#   def get_async(self) -> asyncio.Future:
#   def get(self) -> np.ndarray:
#   def __repr__(self):
#   def get_async(self) -> asyncio.Future:
#   def get(self) -> np.ndarray:
#   def get_async(self) -> asyncio.Future:
#     async def _get_and_cast():
#   def get(self) -> np.ndarray:
#   def from_tensor_store_spec(
#       cls, ts_spec: ts.Spec,
#       get_fn: Callable[[], np.ndarray]) -> 'LazyAwaitableArray':
#   def from_array(cls, array: np.ndarray,
#                  get_fn: Callable[[], np.ndarray]) -> 'LazyAwaitableArray':
#   def from_tensor_store_spec_or_array(
#       cls, maybe_ts_spec: Union[ts.Spec, np.ndarray],
#       get_fn: Callable[[], np.ndarray]) -> 'LazyAwaitableArray':
#   def __init__(self):
#   def add(self, pattern):
#     def register_translation_fn_decorator(fn):
#   def apply(self, flatdict, **opts):
# def global_step(opts, key, val):
# def shared_embeddings(opts, key, val, slot):
# def separate_embeddings(opts, key, val, encdec, slot):
# def rel_embeddings(opts, key, val, encdec, blocknum, slot):
# def attention_layers(opts, key, val, encdec, blocknum, attntype, qkvo, slot):
# def mlpblock(opts, key, val, encdec, blocknum, io_name, io_num, slot):
# def layernorms(opts, key, val, encdec, blocknum, lyrnum, slot):
# def final_layernorms(opts, key, val, encdec, slot):
# def final_logits(opts, key, val, slot):
# def _add_missing_param_states(t5_data):
# def _maybe_correct_relpos_bias(t5_data):
# def load_tf_ckpt(path):
# def _update_state_dict(state_dict: Mapping[str, Any],
#                        t5_data: MutableMapping[str, LazyArray],
#                        strict: bool = True) -> Mapping[str, Any]:
# def restore_from_t5_checkpoint(
#     state_dict: Mapping[str, Any],
#     path: str,
#     lazy_parameters: bool = False,
#     strict: bool = True,
#     translator: Optional[CheckpointTranslator] = None) -> Mapping[str, Any]:
# SLOT_MAP = {'_slot_vc': 'v_col', '_slot_vr': 'v_row', '_slot_v': 'v'}
# TOWER_MAP = {'transformer': 'decoder'}


import json
import os
import jax
import numpy as np
import tensorflow as tf
from absl import flags
from absl.testing import absltest
from t5x import checkpoint_importer
#
#     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.

"""Tests for t5x.checkpoint_importer."""




class CheckpointImporterTest(absltest.TestCase):

  def test_rel_embeddings_shared_layers(self):
    # This represents a ckpt where the Mesh TensorFlow's
    # transformer_layers.SelfAttention.relative_attention_type = "bias_shared",
    # i.e., the same relative attention parameters are shared by all layers
    # within the (en|de)coder.
    ckpt_data = {
        'encoder/block_000/layer_000/SelfAttention/relative_attention_bias':
            1,
        'decoder/block_000/layer_000/SelfAttention/relative_attention_bias':
            2,
        'decoder/block_000/layer_000/SelfAttention/relative_attention_bias_slot_v':
            3,
    }
    t5_data = checkpoint_importer.t5_importer.apply(ckpt_data)

# Path: t5x/checkpoint_importer.py
# class LazyArray(abc.ABC):
# class LazyThreadPoolArray(LazyArray):
# class LazyAwaitableArray(LazyArray):
# class CheckpointTranslator:
#   def __init__(self, shape: Sequence[int], dtype: jnp.dtype,
#                get_fn: Callable[[], np.ndarray]):
#   def shape(self) -> Tuple[int, ...]:
#   def dtype(self) -> jnp.dtype:
#   def nbytes(self) -> int:
#   def astype(self, dtype: np.dtype) -> 'LazyArray':
#   def get_async(self) -> asyncio.Future:
#   def get(self) -> np.ndarray:
#   def __repr__(self):
#   def get_async(self) -> asyncio.Future:
#   def get(self) -> np.ndarray:
#   def get_async(self) -> asyncio.Future:
#     async def _get_and_cast():
#   def get(self) -> np.ndarray:
#   def from_tensor_store_spec(
#       cls, ts_spec: ts.Spec,
#       get_fn: Callable[[], np.ndarray]) -> 'LazyAwaitableArray':
#   def from_array(cls, array: np.ndarray,
#                  get_fn: Callable[[], np.ndarray]) -> 'LazyAwaitableArray':
#   def from_tensor_store_spec_or_array(
#       cls, maybe_ts_spec: Union[ts.Spec, np.ndarray],
#       get_fn: Callable[[], np.ndarray]) -> 'LazyAwaitableArray':
#   def __init__(self):
#   def add(self, pattern):
#     def register_translation_fn_decorator(fn):
#   def apply(self, flatdict, **opts):
# def global_step(opts, key, val):
# def shared_embeddings(opts, key, val, slot):
# def separate_embeddings(opts, key, val, encdec, slot):
# def rel_embeddings(opts, key, val, encdec, blocknum, slot):
# def attention_layers(opts, key, val, encdec, blocknum, attntype, qkvo, slot):
# def mlpblock(opts, key, val, encdec, blocknum, io_name, io_num, slot):
# def layernorms(opts, key, val, encdec, blocknum, lyrnum, slot):
# def final_layernorms(opts, key, val, encdec, slot):
# def final_logits(opts, key, val, slot):
# def _add_missing_param_states(t5_data):
# def _maybe_correct_relpos_bias(t5_data):
# def load_tf_ckpt(path):
# def _update_state_dict(state_dict: Mapping[str, Any],
#                        t5_data: MutableMapping[str, LazyArray],
#                        strict: bool = True) -> Mapping[str, Any]:
# def restore_from_t5_checkpoint(
#     state_dict: Mapping[str, Any],
#     path: str,
#     lazy_parameters: bool = False,
#     strict: bool = True,
#     translator: Optional[CheckpointTranslator] = None) -> Mapping[str, Any]:
# SLOT_MAP = {'_slot_vc': 'v_col', '_slot_vr': 'v_row', '_slot_v': 'v'}
# TOWER_MAP = {'transformer': 'decoder'}


# Path: t5x/checkpoint_importer_test.py
import json
import os
import jax
import numpy as np
import tensorflow as tf
from absl import flags
from absl.testing import absltest
from t5x import checkpoint_importer

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

"""Tests for t5x.checkpoint_importer."""




class CheckpointImporterTest(absltest.TestCase):

  def test_rel_embeddings_shared_layers(self):
    # This represents a ckpt where the Mesh TensorFlow's
    # transformer_layers.SelfAttention.relative_attention_type = "bias_shared",
    # i.e., the same relative attention parameters are shared by all layers
    # within the (en|de)coder.
    ckpt_data = {
        'encoder/block_000/layer_000/SelfAttention/relative_attention_bias':
            1,
        'decoder/block_000/layer_000/SelfAttention/relative_attention_bias':
            2,
        'decoder/block_000/layer_000/SelfAttention/relative_attention_bias_slot_v':
            3,
    }
    t5_data = checkpoint_importer.t5_importer.apply(ckpt_data)

    t5_data = checkpoint_importer._maybe_correct_relpos_bias(t5_data)