# Path: reservoir_nn/typing/types.py



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

"""Keras-like blocks of layers."""




class DownSampleConvolution(tf.keras.layers.Layer):
  """A downsample convolution head block."""

  def __init__(
      self,
      conv_filters: Tuple[int, ...],
      conv_kernel_sizes: Tuple[int, ...],
      max_poolings: Tuple[bool, ...],
      conv_strides: Tuple[int, ...],

# Path: reservoir_nn/typing/types.py



# Path: reservoir_nn/keras/blocks.py
import math
import tensorflow as tf
from typing import Any, Optional, Tuple
from reservoir_nn.typing import types

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

"""Keras-like blocks of layers."""




class DownSampleConvolution(tf.keras.layers.Layer):
  """A downsample convolution head block."""

  def __init__(
      self,
      conv_filters: Tuple[int, ...],
      conv_kernel_sizes: Tuple[int, ...],
      max_poolings: Tuple[bool, ...],
      conv_strides: Tuple[int, ...],

      activation: Optional[types.Activation] = None,