# 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):


from absl.testing import absltest
from absl.testing import parameterized
from reservoir_nn.keras import rewiring
import numpy as np
import tensorflow as tf
# Copyright 2021 Google LLC
#
# 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.

"""Tests for reservoir_nn.keras.rewiring."""



class AdaptiveSparseReservoirTest(parameterized.TestCase):

  def test_layer_with_num_connections_works(self):
    x = tf.constant([1.])

# 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/keras/rewiring_test.py
from absl.testing import absltest
from absl.testing import parameterized
from reservoir_nn.keras import rewiring
import numpy as np
import tensorflow as tf

# Copyright 2021 Google LLC
#
# 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.

"""Tests for reservoir_nn.keras.rewiring."""



class AdaptiveSparseReservoirTest(parameterized.TestCase):

  def test_layer_with_num_connections_works(self):
    x = tf.constant([1.])

    layer = rewiring.AdaptiveSparseReservoir(