# Path: clrs/_src/processors.py
# class Processor(hk.Module):
# class GAT(Processor):
# class GATFull(GAT):
# class GATv2(Processor):
# class GATv2Full(GATv2):
# class PGN(Processor):
# class DeepSets(PGN):
# class MPNN(PGN):
# class PGNMask(PGN):
# class MemNetMasked(Processor):
# class MemNetFull(MemNetMasked):
#   def __call__(
#       self,
#       node_fts: _Array,
#       edge_fts: _Array,
#       graph_fts: _Array,
#       adj_mat: _Array,
#       hidden: _Array,
#       **kwargs,
#   ) -> _Array:
#   def inf_bias(self):
#   def inf_bias_edge(self):
#   def __init__(
#       self,
#       out_size: int,
#       nb_heads: int,
#       activation: Optional[_Fn] = jax.nn.relu,
#       residual: bool = True,
#       name: str = 'gat_aggr',
#   ):
#   def __call__(
#       self,
#       node_fts: _Array,
#       edge_fts: _Array,
#       graph_fts: _Array,
#       adj_mat: _Array,
#       hidden: _Array,
#       **unused_kwargs,
#   ) -> _Array:
#   def __call__(self, node_fts: _Array, edge_fts: _Array, graph_fts: _Array,
#                adj_mat: _Array, hidden: _Array, **unused_kwargs) -> _Array:
#   def __init__(
#       self,
#       out_size: int,
#       nb_heads: int,
#       mid_size: Optional[int] = None,
#       activation: Optional[_Fn] = jax.nn.relu,
#       residual: bool = True,
#       name: str = 'gatv2_aggr',
#   ):
#   def __call__(
#       self,
#       node_fts: _Array,
#       edge_fts: _Array,
#       graph_fts: _Array,
#       adj_mat: _Array,
#       hidden: _Array,
#       **unused_kwargs,
#   ) -> _Array:
#   def __call__(self, node_fts: _Array, edge_fts: _Array, graph_fts: _Array,
#                adj_mat: _Array, hidden: _Array, **unused_kwargs) -> _Array:
#   def __init__(
#       self,
#       out_size: int,
#       mid_size: Optional[int] = None,
#       mid_act: Optional[_Fn] = None,
#       activation: Optional[_Fn] = jax.nn.relu,
#       reduction: _Fn = jnp.max,
#       msgs_mlp_sizes: Optional[List[int]] = None,
#       name: str = 'mpnn_aggr',
#   ):
#   def __call__(
#       self,
#       node_fts: _Array,
#       edge_fts: _Array,
#       graph_fts: _Array,
#       adj_mat: _Array,
#       hidden: _Array,
#       **unused_kwargs,
#   ) -> _Array:
#   def __call__(self, node_fts: _Array, edge_fts: _Array, graph_fts: _Array,
#                adj_mat: _Array, hidden: _Array, nb_nodes: int,
#                batch_size: int) -> _Array:
#   def __call__(self, node_fts: _Array, edge_fts: _Array, graph_fts: _Array,
#                adj_mat: _Array, hidden: _Array, **unused_kwargs) -> _Array:
#   def inf_bias(self):
#   def inf_bias_edge(self):
#   def __init__(
#       self,
#       vocab_size: int,
#       sentence_size: int,
#       linear_output_size: int,
#       embedding_size: int = 16,
#       memory_size: Optional[int] = 128,
#       num_hops: int = 1,
#       nonlin: Callable[[Any], Any] = jax.nn.relu,
#       apply_embeddings: bool = True,
#       init_func: hk.initializers.Initializer = jnp.zeros,
#       name: str = 'memnet') -> None:
#   def __call__(
#       self,
#       node_fts: _Array,
#       edge_fts: _Array,
#       graph_fts: _Array,
#       adj_mat: _Array,
#       hidden: _Array,
#       **unused_kwargs,
#   ) -> _Array:
#   def _apply(self, queries: _Array, stories: _Array) -> _Array:
#   def __call__(self, node_fts: _Array, edge_fts: _Array, graph_fts: _Array,
#                adj_mat: _Array, hidden: _Array, **unused_kwargs) -> _Array:
# def construct_processor(kind: str, hidden_dim: int, nb_heads: int) -> Processor:
# def _position_encoding(sentence_size: int, embedding_size: int) -> np.ndarray:


from absl.testing import absltest
from clrs._src import processors
import chex
import haiku as hk
import jax.numpy as jnp
# 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 processors.py."""



class MemnetTest(absltest.TestCase):

  def test_simple_run_and_check_shapes(self):

    batch_size = 64
    vocab_size = 177
    embedding_size = 64
    sentence_size = 11
    memory_size = 320
    linear_output_size = 128
    num_hops = 2

    def forward_fn(queries, stories):

# Path: clrs/_src/processors.py
# class Processor(hk.Module):
# class GAT(Processor):
# class GATFull(GAT):
# class GATv2(Processor):
# class GATv2Full(GATv2):
# class PGN(Processor):
# class DeepSets(PGN):
# class MPNN(PGN):
# class PGNMask(PGN):
# class MemNetMasked(Processor):
# class MemNetFull(MemNetMasked):
#   def __call__(
#       self,
#       node_fts: _Array,
#       edge_fts: _Array,
#       graph_fts: _Array,
#       adj_mat: _Array,
#       hidden: _Array,
#       **kwargs,
#   ) -> _Array:
#   def inf_bias(self):
#   def inf_bias_edge(self):
#   def __init__(
#       self,
#       out_size: int,
#       nb_heads: int,
#       activation: Optional[_Fn] = jax.nn.relu,
#       residual: bool = True,
#       name: str = 'gat_aggr',
#   ):
#   def __call__(
#       self,
#       node_fts: _Array,
#       edge_fts: _Array,
#       graph_fts: _Array,
#       adj_mat: _Array,
#       hidden: _Array,
#       **unused_kwargs,
#   ) -> _Array:
#   def __call__(self, node_fts: _Array, edge_fts: _Array, graph_fts: _Array,
#                adj_mat: _Array, hidden: _Array, **unused_kwargs) -> _Array:
#   def __init__(
#       self,
#       out_size: int,
#       nb_heads: int,
#       mid_size: Optional[int] = None,
#       activation: Optional[_Fn] = jax.nn.relu,
#       residual: bool = True,
#       name: str = 'gatv2_aggr',
#   ):
#   def __call__(
#       self,
#       node_fts: _Array,
#       edge_fts: _Array,
#       graph_fts: _Array,
#       adj_mat: _Array,
#       hidden: _Array,
#       **unused_kwargs,
#   ) -> _Array:
#   def __call__(self, node_fts: _Array, edge_fts: _Array, graph_fts: _Array,
#                adj_mat: _Array, hidden: _Array, **unused_kwargs) -> _Array:
#   def __init__(
#       self,
#       out_size: int,
#       mid_size: Optional[int] = None,
#       mid_act: Optional[_Fn] = None,
#       activation: Optional[_Fn] = jax.nn.relu,
#       reduction: _Fn = jnp.max,
#       msgs_mlp_sizes: Optional[List[int]] = None,
#       name: str = 'mpnn_aggr',
#   ):
#   def __call__(
#       self,
#       node_fts: _Array,
#       edge_fts: _Array,
#       graph_fts: _Array,
#       adj_mat: _Array,
#       hidden: _Array,
#       **unused_kwargs,
#   ) -> _Array:
#   def __call__(self, node_fts: _Array, edge_fts: _Array, graph_fts: _Array,
#                adj_mat: _Array, hidden: _Array, nb_nodes: int,
#                batch_size: int) -> _Array:
#   def __call__(self, node_fts: _Array, edge_fts: _Array, graph_fts: _Array,
#                adj_mat: _Array, hidden: _Array, **unused_kwargs) -> _Array:
#   def inf_bias(self):
#   def inf_bias_edge(self):
#   def __init__(
#       self,
#       vocab_size: int,
#       sentence_size: int,
#       linear_output_size: int,
#       embedding_size: int = 16,
#       memory_size: Optional[int] = 128,
#       num_hops: int = 1,
#       nonlin: Callable[[Any], Any] = jax.nn.relu,
#       apply_embeddings: bool = True,
#       init_func: hk.initializers.Initializer = jnp.zeros,
#       name: str = 'memnet') -> None:
#   def __call__(
#       self,
#       node_fts: _Array,
#       edge_fts: _Array,
#       graph_fts: _Array,
#       adj_mat: _Array,
#       hidden: _Array,
#       **unused_kwargs,
#   ) -> _Array:
#   def _apply(self, queries: _Array, stories: _Array) -> _Array:
#   def __call__(self, node_fts: _Array, edge_fts: _Array, graph_fts: _Array,
#                adj_mat: _Array, hidden: _Array, **unused_kwargs) -> _Array:
# def construct_processor(kind: str, hidden_dim: int, nb_heads: int) -> Processor:
# def _position_encoding(sentence_size: int, embedding_size: int) -> np.ndarray:


# Path: clrs/_src/processors_test.py
from absl.testing import absltest
from clrs._src import processors
import chex
import haiku as hk
import jax.numpy as jnp

# 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 processors.py."""



class MemnetTest(absltest.TestCase):

  def test_simple_run_and_check_shapes(self):

    batch_size = 64
    vocab_size = 177
    embedding_size = 64
    sentence_size = 11
    memory_size = 320
    linear_output_size = 128
    num_hops = 2

    def forward_fn(queries, stories):

      model = processors.MemNetFull(