# Path: reservoir_nn/typing/types.py



import enum
import tensorflow as tf
from typing import Any, Dict, Optional, Union
from reservoir_nn.typing import types
# 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.

"""Functions for Valkyrie learning rules."""




@enum.unique
class LearningRuleName(enum.Enum):
  NONE = 'none'
  HEBBIAN = 'hebbian'
  OJA = 'oja'
  CONTRASTIVEHEBBIAN = 'contrastive_hebbian'


class LocalLearning(tf.keras.layers.Layer):
  """Layer to perform local learning."""

  def __init__(self, learning_rule: str = 'none', **kwargs):
    super().__init__(**kwargs)

    self.learning_rule = LearningRuleName(learning_rule).value

  def call(
      self,
      learning_params: Dict[str, Any],

# Path: reservoir_nn/typing/types.py



# Path: reservoir_nn/local_learning/learning.py
import enum
import tensorflow as tf
from typing import Any, Dict, Optional, Union
from reservoir_nn.typing import types

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

"""Functions for Valkyrie learning rules."""




@enum.unique
class LearningRuleName(enum.Enum):
  NONE = 'none'
  HEBBIAN = 'hebbian'
  OJA = 'oja'
  CONTRASTIVEHEBBIAN = 'contrastive_hebbian'


class LocalLearning(tf.keras.layers.Layer):
  """Layer to perform local learning."""

  def __init__(self, learning_rule: str = 'none', **kwargs):
    super().__init__(**kwargs)

    self.learning_rule = LearningRuleName(learning_rule).value

  def call(
      self,
      learning_params: Dict[str, Any],

      activation: types.TensorLike,