# Path: reservoir_nn/keras/reservoir_registry.py
# _LAYER_DICT = {
#     'DenseReservoir': layers.DenseReservoir,
#     'RecurrentDenseReservoir': layers.RecurrentDenseReservoir,
#     'LSTMDenseReservoir': layers.LSTMDenseReservoir,
#     'SparseReservoir': layers.SparseReservoir,
#     'RecurrentSparseReservoir': layers.RecurrentSparseReservoir,
#     'LSTMSparseReservoir': layers.LSTMSparseReservoir,
# }
# def get_reservoir(reservoir_name: str) -> Type[tf.keras.layers.Layer]:


from typing import Any, Dict, Optional, Tuple
from reservoir_nn.keras import reservoir_registry
import numpy as np
import tensorflow as tf
        embed_dim, 11, strides=2, padding='same', activation='relu')

  def call(self, inputs: tf.Tensor) -> tf.Tensor:
    x = self.conv1(inputs)
    x = self.conv2(x)
    return self.conv3(x)


class ReservoirTransformerEncoder(tf.keras.layers.Layer):
  """Encoder model for Seq2Seq Reservoir-based Transformer."""

  def __init__(self,
               embed_dim: int,
               num_heads: int,
               reservoir_base: str,
               reservoir_params: Dict[str, Any],
               dense_dim: int = 2048,
               build_reservoir: bool = False,
               **kwargs):
    super(ReservoirTransformerEncoder, self).__init__(**kwargs)
    self.embed_dim = embed_dim
    self.num_heads = num_heads
    self.dense_dim = dense_dim
    self.build_reservoir = build_reservoir
    self.attention = tf.keras.layers.MultiHeadAttention(
        num_heads=num_heads, key_dim=embed_dim)

    if self.build_reservoir:
      self.reservoir_top = tf.keras.layers.Dense(
          reservoir_params['weight'].shape[0])

# Path: reservoir_nn/keras/reservoir_registry.py
# _LAYER_DICT = {
#     'DenseReservoir': layers.DenseReservoir,
#     'RecurrentDenseReservoir': layers.RecurrentDenseReservoir,
#     'LSTMDenseReservoir': layers.LSTMDenseReservoir,
#     'SparseReservoir': layers.SparseReservoir,
#     'RecurrentSparseReservoir': layers.RecurrentSparseReservoir,
#     'LSTMSparseReservoir': layers.LSTMSparseReservoir,
# }
# def get_reservoir(reservoir_name: str) -> Type[tf.keras.layers.Layer]:


# Path: reservoir_nn/models/transformer.py
from typing import Any, Dict, Optional, Tuple
from reservoir_nn.keras import reservoir_registry
import numpy as np
import tensorflow as tf

        embed_dim, 11, strides=2, padding='same', activation='relu')

  def call(self, inputs: tf.Tensor) -> tf.Tensor:
    x = self.conv1(inputs)
    x = self.conv2(x)
    return self.conv3(x)


class ReservoirTransformerEncoder(tf.keras.layers.Layer):
  """Encoder model for Seq2Seq Reservoir-based Transformer."""

  def __init__(self,
               embed_dim: int,
               num_heads: int,
               reservoir_base: str,
               reservoir_params: Dict[str, Any],
               dense_dim: int = 2048,
               build_reservoir: bool = False,
               **kwargs):
    super(ReservoirTransformerEncoder, self).__init__(**kwargs)
    self.embed_dim = embed_dim
    self.num_heads = num_heads
    self.dense_dim = dense_dim
    self.build_reservoir = build_reservoir
    self.attention = tf.keras.layers.MultiHeadAttention(
        num_heads=num_heads, key_dim=embed_dim)

    if self.build_reservoir:
      self.reservoir_top = tf.keras.layers.Dense(
          reservoir_params['weight'].shape[0])

      self.reservoir_layer = reservoir_registry.get_reservoir(reservoir_base)(