# Path: reservoir_nn/typing/types.py



import abc
import dataclasses
import functools
import logging
import numpy as np
import tensorflow as tf
from typing import Any, Callable, Dict, Optional, Tuple, Union
from reservoir_nn.typing import types
from scipy import sparse




class RoundRobinConnection():
  """Creates a roundrobin sparse reservoir."""

  def __init__(self, num_connections: int):
    self.num_connections = num_connections

  def __call__(self, shape: Tuple[int, int]) -> Tuple[np.ndarray, np.ndarray]:
    if shape[0] * shape[-1] < self.num_connections:
      raise ValueError(f'Cannot build layer because the target'
                       f'({self.num_connections}) is larger than max number'
                       f'of possible connections between input units '
                       f'({shape[0]}) and output units ({shape[1]})')

    initial_indices = np.zeros(shape=(self.num_connections, 2), dtype='int64')
    initial_indices[...] = np.arange(self.num_connections)[:, None]
    initial_indices[:, 0] %= shape[0]
    initial_indices[:, 1] %= shape[1]

    initial_values = np.zeros(self.num_connections)

    return initial_indices, initial_values


class CooMatrixConnection():
  """Creates a sparse reservoir for the input."""


# Path: reservoir_nn/typing/types.py



# Path: reservoir_nn/keras/rewiring.py
import abc
import dataclasses
import functools
import logging
import numpy as np
import tensorflow as tf
from typing import Any, Callable, Dict, Optional, Tuple, Union
from reservoir_nn.typing import types
from scipy import sparse





class RoundRobinConnection():
  """Creates a roundrobin sparse reservoir."""

  def __init__(self, num_connections: int):
    self.num_connections = num_connections

  def __call__(self, shape: Tuple[int, int]) -> Tuple[np.ndarray, np.ndarray]:
    if shape[0] * shape[-1] < self.num_connections:
      raise ValueError(f'Cannot build layer because the target'
                       f'({self.num_connections}) is larger than max number'
                       f'of possible connections between input units '
                       f'({shape[0]}) and output units ({shape[1]})')

    initial_indices = np.zeros(shape=(self.num_connections, 2), dtype='int64')
    initial_indices[...] = np.arange(self.num_connections)[:, None]
    initial_indices[:, 0] %= shape[0]
    initial_indices[:, 1] %= shape[1]

    initial_values = np.zeros(self.num_connections)

    return initial_indices, initial_values


class CooMatrixConnection():
  """Creates a sparse reservoir for the input."""


  def __init__(self, weight_matrix: types.WeightMatrix):