text: from functools import partial
  
  import torch
  
  from caldera.blocks import AggregatingEdgeBlock
  from caldera.blocks import AggregatingGlobalBlock
  from caldera.blocks import AggregatingNodeBlock
  from caldera.blocks import Aggregator
  from caldera.blocks import EdgeBlock
  from caldera.blocks import Flex
  from caldera.blocks import GlobalBlock
  from caldera.blocks import MLP
  from caldera.blocks import NodeBlock
  from caldera.data import GraphBatch
  from caldera.models.graph_core import GraphCore
  from caldera.models.graph_encoder import GraphEncoder
  
  
  class EncodeCoreDecode(torch.nn.Module):
      def __init__(
          self,
          latent_sizes=(128, 128, 1),
          output_sizes=(1, 1, 1),
          depths=(1, 1, 1),
          layer_norm: bool = True,
          dropout: float = None,
          pass_global_to_edge: bool = True,
          pass_global_to_node: bool = True,
      ):
          super().__init__()
          self.config = {
              "latent_size": {
                  "node": latent_sizes[1],
                  "edge": latent_sizes[0],
                  "global": latent_sizes[2],
                  "core_node_block_depth": depths[0],
                  "core_edge_block_depth": depths[1],
                  "core_global_block_depth": depths[2],
              },
              "output_size": {
                  "edge": output_sizes[0],
                  "node": output_sizes[1],
                  "global": output_sizes[2],
              },
              "node_block_aggregator": "add",
              "global_block_to_node_aggregator": "add",
              "global_block_to_edge_aggregator": "add",
              "pass_global_to_edge": pass_global_to_edge,
              "pass_global_to_node": pass_global_to_node,
          }
  
          def mlp(*layer_sizes):
              return Flex(MLP)(
                  Flex.d(), *layer_sizes, layer_norm=layer_norm, dropout=dropout
              )
  
          self.encoder = GraphEncoder(
              EdgeBlock(mlp(latent_sizes[0])),
              NodeBlock(mlp(latent_sizes[1])),
              GlobalBlock(mlp(latent_sizes[2])),
          )
  
          edge_layers = [self.config["latent_size"]["edge"]] * self.config["latent_size"][
              "core_edge_block_depth"
          ]
          node_layers = [self.config["latent_size"]["node"]] * self.config["latent_size"][
              "core_node_block_depth"
          ]
          global_layers = [self.config["latent_size"]["global"]] * self.config[
              "latent_size"
          ]["core_global_block_depth"]
  
          self.core = GraphCore(
              AggregatingEdgeBlock(mlp(*edge_layers)),
              AggregatingNodeBlock(
                  mlp(*node_layers), Aggregator(self.config["node_block_aggregator"])
              ),
              AggregatingGlobalBlock(
                  mlp(*global_layers),
                  edge_aggregator=Aggregator(
                      self.config["global_block_to_edge_aggregator"]
                  ),
                  node_aggregator=Aggregator(
                      self.config["global_block_to_node_aggregator"]
                  ),
              ),
              pass_global_to_edge=self.config["pass_global_to_edge"],
              pass_global_to_node=self.config["pass_global_to_node"],
          )
  
          self.decoder = GraphEncoder(
              EdgeBlock(mlp(latent_sizes[0])),
              NodeBlock(mlp(latent_sizes[1])),
              GlobalBlock(mlp(latent_sizes[2])),
          )
  
          self.output_transform = GraphEncoder(
              EdgeBlock(Flex(torch.nn.Linear)(Flex.d(), output_sizes[0])),
              NodeBlock(Flex(torch.nn.Linear)(Flex.d(), output_sizes[1])),
              GlobalBlock(Flex(torch.nn.Linear)(Flex.d(), output_sizes[2])),
          )
  
      def forward(self, data, steps):
          # encoded
          e, x, g = self.encoder(data)
          data = GraphBatch(x, e, g, data.edges, data.node_idx, data.edge_idx)
  
          # graph topography data
          edges = data.edges
          node_idx = data.node_idx
          edge_idx = data.edge_idx
          latent0 = data
  
          meta = (edges, node_idx, edge_idx)
  
          outputs = []
          for _ in range(steps):
              # core processing step
              e = torch.cat([latent0.e, e], dim=1)
              x = torch.cat([latent0.x, x], dim=1)
              g = torch.cat([latent0.g, g], dim=1)
              data = GraphBatch(x, e, g, *meta)
              e, x, g = self.core(data)
  
              # decode
              data = GraphBatch(x, e, g, *meta)
  
              _e, _x, _g = self.decoder(data)
              decoded = GraphBatch(_x, _e, _g, *meta)
  
              # transform
              _e, _x, _g = self.output_transform(decoded)
              outputs.append(GraphBatch(_x, _e, _g, edges, node_idx, edge_idx))
  
          # revise connectivity
  
          return outputs
  