# Path: reservoir_nn/keras/rewiring.py
# class RoundRobinConnection():
# class CooMatrixConnection():
# class SparseMutableInterface(metaclass=abc.ABCMeta):
# class AdaptiveSparseReservoir(SparseMutableInterface, tf.keras.layers.Layer):
# class GlobalPolicy:
# class MutationPolicy:
# class MutationCallback(tf.keras.callbacks.Callback):
#   def __init__(self, num_connections: int):
#   def __call__(self, shape: Tuple[int, int]) -> Tuple[np.ndarray, np.ndarray]:
#   def __init__(self, weight_matrix: types.WeightMatrix):
#   def __call__(self, shape: Tuple[int, int]) -> Tuple[np.ndarray, np.ndarray]:
# def _get_reservoir_initializer(
#     reservoir_initializer: Any) -> types.ReservoirInitializer:
#   def get_reservoir_shape(self):
#   def get_sparse_tensors(self) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
#   def assign_sparse_tensors(self, sparse_indices: tf.Tensor,
#                             sparse_values: tf.Tensor, sparse_ages: tf.Tensor):
#   def __init__(
#       self,
#       units: int,
#       reservoir_initializer: Union[int, types.WeightMatrix,
#                                    types.ReservoirInitializer],
#       activation: Optional[types.Activation] = None,
#       use_bias: bool = False,
#       kernel_regularizer: Optional[types.Regularizer] = None,
#       bias_initializer: types.Initializer = 'zeros',
#       bias_regularizer: Optional[types.Regularizer] = None,
#       activity_regularizer: Optional[types.Regularizer] = None,
#       bias_constraint: Optional[types.Constraint] = None,
#       name: Optional[str] = None,
#       **kwargs: Any,
#   ):
#   def build(self, input_shape: Tuple[int, ...]):
#   def call(self, inputs: tf.Tensor) -> tf.Tensor:
#   def get_coo_weight_matrix(self) -> sparse.spmatrix:
#   def get_coo_age_matrix(self) -> sparse.spmatrix:
#   def get_reservoir_shape(self):
#   def get_sparse_tensors(self):
#   def assign_sparse_tensors(self, sparse_indices, sparse_values, sparse_ages):
#   def get_config(self):
#   def apply_global_policy(self, gpolicy: GlobalPolicy) -> 'MutationPolicy':
#   def mutation_step(self,
#                     layer: SparseMutableInterface,
#                     rng: Optional[np.random.RandomState] = None):
#   def compute_mutation_probability(self, sparse_values: tf.Tensor,
#                                    sparse_ages: tf.Tensor) -> tf.Tensor:
# def _mutate_sparse_reservoir(
#     sparse_indices: np.ndarray,
#     sparse_values: np.ndarray,
#     sparse_ages: np.ndarray,
#     mutation_prob: np.ndarray,
#     mutation_rate: float,
#     rng: np.random.RandomState,
#     shape: Tuple[int, int],
# ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
#   def __init__(self,
#                policy: Union[MutationPolicy, Dict[str, MutationPolicy]],
#                mutation_schedule: Optional[Callable[[int],
#                                                     GlobalPolicy]] = None,
#                rng: Optional[np.random.RandomState] = None,
#                verbose: int = 0):
#   def _get_policy(self, layer):
#   def on_epoch_begin(self, epoch, logs):
#   def on_train_batch_begin(self, batch, logs):
# 
# Path: reservoir_nn/local_learning/learning.py
# class LearningRuleName(enum.Enum):
# class LocalLearning(tf.keras.layers.Layer):
#   NONE = 'none'
#   HEBBIAN = 'hebbian'
#   OJA = 'oja'
#   CONTRASTIVEHEBBIAN = 'contrastive_hebbian'
#   def __init__(self, learning_rule: str = 'none', **kwargs):
#   def call(
#       self,
#       learning_params: Dict[str, Any],
#       activation: types.TensorLike,
#       weight: types.TensorLike,
#       expected: Optional[types.TensorLike] = None
#   ) -> Union[types.TensorLike, None]:
#   def hebbian_learning(self,
#                        activation: types.TensorLike,
#                        weight: types.TensorLike,
#                        eta: float = 0.1) -> types.TensorLike:
#   def oja_learning(self,
#                    activation: types.TensorLike,
#                    weight: types.TensorLike,
#                    eta: float = 0.1) -> types.TensorLike:
#   def contrastive_hebbian_learning(self,
#                                    activation: types.TensorLike,
#                                    weight: types.TensorLike,
#                                    expected: types.TensorLike,
#                                    eta: float = 0.1,
#                                    gamma: float = 0.1,
#                                    max_iter: int = 10) -> types.TensorLike:
# 
# Path: reservoir_nn/typing/types.py



from typing import Any, Dict, List, Optional, Tuple, Type, Union
from reservoir_nn.keras import rewiring
from reservoir_nn.local_learning import learning
from reservoir_nn.typing import types
from scipy import sparse
import numpy as np
import tensorflow as tf
              ) != 2 or self.reservoir_shape[0] != self.reservoir_shape[1]:
          raise ValueError(
              f'In a recurrent layer the `weight` must be a square matrix. '
              f'Got shape {self.reservoir_shape}.')

  @property
  def units(self):
    return self.reservoir_shape[-1]

  def call(self, inputs: tf.Tensor) -> tf.Tensor:
    if self.kernel is None:
      raise RuntimeError('The kernel of the DenseReservoir layer is not '
                         'initialized. Call DenseReservoir with a weight '
                         'matrix to initialize it.')
    matrices = [self.kernel] * (self.recurrence_degree + 1)

    rank = inputs.shape.rank
    if rank < 2:
      raise ValueError(f'Input tensor must be at least 2D, but rank={rank}.')

    outputs = inputs
    for i, kernel in enumerate(matrices):
      outputs = tf.tensordot(outputs, kernel, [[rank - 1], [0]])

      if self.activation_within_recurrence or i == len(matrices) - 1:
        if self.bias is not None:
          outputs = tf.nn.bias_add(outputs, self.bias)
        if self.activation is not None:
          outputs = self.activation(outputs)


# Path: reservoir_nn/keras/rewiring.py
# class RoundRobinConnection():
# class CooMatrixConnection():
# class SparseMutableInterface(metaclass=abc.ABCMeta):
# class AdaptiveSparseReservoir(SparseMutableInterface, tf.keras.layers.Layer):
# class GlobalPolicy:
# class MutationPolicy:
# class MutationCallback(tf.keras.callbacks.Callback):
#   def __init__(self, num_connections: int):
#   def __call__(self, shape: Tuple[int, int]) -> Tuple[np.ndarray, np.ndarray]:
#   def __init__(self, weight_matrix: types.WeightMatrix):
#   def __call__(self, shape: Tuple[int, int]) -> Tuple[np.ndarray, np.ndarray]:
# def _get_reservoir_initializer(
#     reservoir_initializer: Any) -> types.ReservoirInitializer:
#   def get_reservoir_shape(self):
#   def get_sparse_tensors(self) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
#   def assign_sparse_tensors(self, sparse_indices: tf.Tensor,
#                             sparse_values: tf.Tensor, sparse_ages: tf.Tensor):
#   def __init__(
#       self,
#       units: int,
#       reservoir_initializer: Union[int, types.WeightMatrix,
#                                    types.ReservoirInitializer],
#       activation: Optional[types.Activation] = None,
#       use_bias: bool = False,
#       kernel_regularizer: Optional[types.Regularizer] = None,
#       bias_initializer: types.Initializer = 'zeros',
#       bias_regularizer: Optional[types.Regularizer] = None,
#       activity_regularizer: Optional[types.Regularizer] = None,
#       bias_constraint: Optional[types.Constraint] = None,
#       name: Optional[str] = None,
#       **kwargs: Any,
#   ):
#   def build(self, input_shape: Tuple[int, ...]):
#   def call(self, inputs: tf.Tensor) -> tf.Tensor:
#   def get_coo_weight_matrix(self) -> sparse.spmatrix:
#   def get_coo_age_matrix(self) -> sparse.spmatrix:
#   def get_reservoir_shape(self):
#   def get_sparse_tensors(self):
#   def assign_sparse_tensors(self, sparse_indices, sparse_values, sparse_ages):
#   def get_config(self):
#   def apply_global_policy(self, gpolicy: GlobalPolicy) -> 'MutationPolicy':
#   def mutation_step(self,
#                     layer: SparseMutableInterface,
#                     rng: Optional[np.random.RandomState] = None):
#   def compute_mutation_probability(self, sparse_values: tf.Tensor,
#                                    sparse_ages: tf.Tensor) -> tf.Tensor:
# def _mutate_sparse_reservoir(
#     sparse_indices: np.ndarray,
#     sparse_values: np.ndarray,
#     sparse_ages: np.ndarray,
#     mutation_prob: np.ndarray,
#     mutation_rate: float,
#     rng: np.random.RandomState,
#     shape: Tuple[int, int],
# ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
#   def __init__(self,
#                policy: Union[MutationPolicy, Dict[str, MutationPolicy]],
#                mutation_schedule: Optional[Callable[[int],
#                                                     GlobalPolicy]] = None,
#                rng: Optional[np.random.RandomState] = None,
#                verbose: int = 0):
#   def _get_policy(self, layer):
#   def on_epoch_begin(self, epoch, logs):
#   def on_train_batch_begin(self, batch, logs):
# 
# Path: reservoir_nn/local_learning/learning.py
# class LearningRuleName(enum.Enum):
# class LocalLearning(tf.keras.layers.Layer):
#   NONE = 'none'
#   HEBBIAN = 'hebbian'
#   OJA = 'oja'
#   CONTRASTIVEHEBBIAN = 'contrastive_hebbian'
#   def __init__(self, learning_rule: str = 'none', **kwargs):
#   def call(
#       self,
#       learning_params: Dict[str, Any],
#       activation: types.TensorLike,
#       weight: types.TensorLike,
#       expected: Optional[types.TensorLike] = None
#   ) -> Union[types.TensorLike, None]:
#   def hebbian_learning(self,
#                        activation: types.TensorLike,
#                        weight: types.TensorLike,
#                        eta: float = 0.1) -> types.TensorLike:
#   def oja_learning(self,
#                    activation: types.TensorLike,
#                    weight: types.TensorLike,
#                    eta: float = 0.1) -> types.TensorLike:
#   def contrastive_hebbian_learning(self,
#                                    activation: types.TensorLike,
#                                    weight: types.TensorLike,
#                                    expected: types.TensorLike,
#                                    eta: float = 0.1,
#                                    gamma: float = 0.1,
#                                    max_iter: int = 10) -> types.TensorLike:
# 
# Path: reservoir_nn/typing/types.py



# Path: reservoir_nn/keras/layers.py
from typing import Any, Dict, List, Optional, Tuple, Type, Union
from reservoir_nn.keras import rewiring
from reservoir_nn.local_learning import learning
from reservoir_nn.typing import types
from scipy import sparse
import numpy as np
import tensorflow as tf

              ) != 2 or self.reservoir_shape[0] != self.reservoir_shape[1]:
          raise ValueError(
              f'In a recurrent layer the `weight` must be a square matrix. '
              f'Got shape {self.reservoir_shape}.')

  @property
  def units(self):
    return self.reservoir_shape[-1]

  def call(self, inputs: tf.Tensor) -> tf.Tensor:
    if self.kernel is None:
      raise RuntimeError('The kernel of the DenseReservoir layer is not '
                         'initialized. Call DenseReservoir with a weight '
                         'matrix to initialize it.')
    matrices = [self.kernel] * (self.recurrence_degree + 1)

    rank = inputs.shape.rank
    if rank < 2:
      raise ValueError(f'Input tensor must be at least 2D, but rank={rank}.')

    outputs = inputs
    for i, kernel in enumerate(matrices):
      outputs = tf.tensordot(outputs, kernel, [[rank - 1], [0]])

      if self.activation_within_recurrence or i == len(matrices) - 1:
        if self.bias is not None:
          outputs = tf.nn.bias_add(outputs, self.bias)
        if self.activation is not None:
          outputs = self.activation(outputs)


    self.kernel = learning.LocalLearning(