text: """
  Materials Graph Network
  """
  from copy import deepcopy
  from typing import List, Optional
  
  import tensorflow as tf
  
  from m3gnet.utils import register
  
  from ._atom import AtomNetwork
  from ._base import GraphUpdate, GraphUpdateFunc
  from ._bond import BondNetwork, PairRadialBasisExpansion
  from ._core import Embedding
  from ._state import StateNetwork
  
  
  @register
  class GraphNetworkLayer(GraphUpdate):
      """
      A graph network layer features bond/atom/state update in the sequence
      bond -> atom -> state. The input and output of each step are graphs/graph
      converted lists
      """
  
      def __init__(
          self,
          bond_network: Optional[GraphUpdate] = None,
          atom_network: Optional[GraphUpdate] = None,
          state_network: Optional[GraphUpdate] = None,
          **kwargs,
      ):
          """
  
          Args:
              bond_network (GraphUpdate): bond update network
              atom_network (GraphUpdate): atom update network
              state_network (GraphUpdate): state update network
              **kwargs:
          """
          self.bond_network = bond_network or GraphUpdate()
          self.atom_network = atom_network or GraphUpdate()
          self.state_network = state_network or GraphUpdate()
          super().__init__(**kwargs)
  
      def call(self, graph: List, **kwargs) -> List:
          """
          Args:
              graph (List): a graph in list representation
              **kwargs:
          Returns: tf.Tensor
          """
          out = self.state_network(self.atom_network(self.bond_network(graph)))
          return out
  
      def get_config(self) -> dict:
          """
          Get config dict for serialization
          Returns: dict
          """
          config = super().get_config()
          config.update(
              {
                  "bond_network": self.bond_network,
                  "atom_network": self.atom_network,
                  "state_network": self.state_network,
              }
          )
          return config
  
  
  def _get_bond_featurizer(nfeat_bond, n_bond_types, bond_embedding_dim, rbf_type, kwargs):
      # bond network settings for the featurization layer
      # no need for bond network
      if nfeat_bond is not None:
          return BondNetwork()
  
      # no feature length, meaning that the features are either
      # integer  types or just distances
      if bond_embedding_dim is not None:
          # bond attributes are integer types
          return GraphFieldEmbedding(
              nvocal=n_bond_types,
              embedding_dim=bond_embedding_dim,
              field="bonds",
              name="bond_embedding",
          )
  
      if rbf_type.lower() == "gaussian":
          centers = kwargs.pop("centers", None)
          width = kwargs.pop("width", None)
          if centers is None and width is None:
              raise ValueError(
                  "If the bond attributes are single float values, "
                  "we expect the value to be expanded before passing "
                  "to the pretrained. Therefore, `centers` and `width` for "
                  "Gaussian basis expansion are needed"
              )
          # bond attributes are distances
          bond_network = PairRadialBasisExpansion(rbf_type="Gaussian", centers=centers, width=width)
          return bond_network
  
      if rbf_type.lower() == "sphericalbessel":
          max_l = kwargs.pop("max_l")
          max_n = kwargs.pop("max_n")
          cutoff = kwargs.pop("cutoff")
          smooth = kwargs.pop("smooth", False)
          return PairRadialBasisExpansion(
              rbf_type="SphericalBessel",
              max_l=max_l,
              max_n=max_n,
              cutoff=cutoff,
              smooth=smooth,
          )
      raise ValueError("Cannot derive bond network type")
  
  
  @register
  class GraphFeaturizer(GraphNetworkLayer):
      """
      Graph featurizer that does several things to convert an initial graph
      with atomic number atom attributes and bond distance bond attributes to
      a graph with proper feature dimensions
      """
  
      def __init__(
          self,
          nfeat_bond: Optional[int] = None,
          nfeat_atom: Optional[int] = None,
          nfeat_state: Optional[int] = None,
          n_bond_types: Optional[int] = None,
          n_atom_types: Optional[int] = None,
          n_state_types: Optional[int] = None,
          bond_embedding_dim: Optional[int] = None,
          atom_embedding_dim: Optional[int] = None,
          state_embedding_dim: Optional[int] = None,
          rbf_type="SpherialBessel",
          **kwargs,
      ):
          """
  
          Args:
              nfeat_bond (int): bond feature dimension
              nfeat_atom (int): atom feature dimension
              nfeat_state (int): state feature dimension
              n_bond_types (int): number of bond types, used only when the
                  bond is categorical
              n_atom_types (int): number of atom types
              n_state_types (int): number of state types, used only when state is
                  of categorical type
              bond_embedding_dim (int): embedding dimension for bond
              atom_embedding_dim (int): embedding dimension for atom
              state_embedding_dim (int): embedding dimension for state
              rbf_type (str): radial basis function type, choose between
                  "SphericalBessel" and "Gaussian"
              **kwargs:
          """
          self.kwargs = deepcopy(kwargs)
  
          # bond network settings for the featurization layer
          bond_network = _get_bond_featurizer(nfeat_bond, n_bond_types, bond_embedding_dim, rbf_type, kwargs)
          # atom network for the featurization layer
          if nfeat_atom is None:
              if n_atom_types is None or atom_embedding_dim is None:
                  raise ValueError("Either specify nfeat_atom or " "n_atom_types and atom_embedding_dim")
              atom_network = GraphFieldEmbedding(
                  nvocal=n_atom_types + 1,
                  embedding_dim=atom_embedding_dim,
                  field="atoms",
                  name="atom_embedding",
              )
          else:
              atom_network = AtomNetwork()
  
          # state network for the featurization layer
          if nfeat_state is None and state_embedding_dim is not None:
              if n_state_types is None:
                  raise ValueError("Either specify nfeat_state or " "n_state_types and state_embedding_dim")
              state_network = GraphFieldEmbedding(
                  nvocal=n_state_types,
                  embedding_dim=state_embedding_dim,
                  field="states",
                  name="state_embedding",
              )
          else:
              state_network = StateNetwork()
  
          self.nfeat_bond = nfeat_bond
          self.nfeat_atom = nfeat_atom
          self.nfeat_state = nfeat_state
          self.n_bond_types = n_bond_types
          self.n_atom_types = n_atom_types
          self.n_state_types = n_state_types
          self.bond_embedding_dim = bond_embedding_dim
          self.atom_embedding_dim = atom_embedding_dim
          self.state_embedding_dim = state_embedding_dim
          self.rbf_type = rbf_type
          super().__init__(bond_network=bond_network, atom_network=atom_network, state_network=state_network, **kwargs)
  
      def get_config(self) -> dict:
          """
          Get config dict for serialization
          Returns:
          """
          config = tf.keras.layers.Layer.get_config(self)
          config.update(
              nfeat_bond=self.nfeat_bond,
              nfeat_atom=self.nfeat_atom,
              nfeat_state=self.nfeat_state,
              n_bond_types=self.n_bond_types,
              n_atom_types=self.n_atom_types,
              n_state_types=self.n_state_types,
              bond_embedding_dim=self.bond_embedding_dim,
              atom_embedding_dim=self.atom_embedding_dim,
              state_embedding_dim=self.state_embedding_dim,
              rbf_type=self.rbf_type,
          )
          config.update(**self.kwargs)
          return config
  
  
  @register
  class GraphFieldEmbedding(GraphUpdateFunc):
      """
      Embedding the categorical field of a graph to continuous space
      """
  
      def __init__(self, nvocal: int = 95, embedding_dim: int = 16, field: str = "atoms", **kwargs):
          """
  
          Args:
              nvocal (int): number of vocabulary
              embedding_dim (int): the embedding dimension
              field (str): graph field for embedding
              **kwargs:
          """
          self.nvocal = nvocal
          self.embedding_dim = embedding_dim
          self.field = field
          self.kwargs = kwargs
          self.embedding = Embedding(nvocal, embedding_dim, **kwargs)
          super().__init__(update_func=self.embedding, update_field=field, **kwargs)
  
      def get_config(self) -> dict:
          """
          Get config dict for serialization
          Returns:
          """
          return {
              "nvocal": self.nvocal,
              "embedding_dim": self.embedding_dim,
              "field": self.field,
              "name": self.name,
              "trainable": self.trainable,
              "dtype": self.dtype,
          }
  