class_docstr: The equilibrium aggregation layer from the `"Equilibrium Aggregation:
  Encoding Sets via Optimization" <https://arxiv.org/abs/2202.12795>`_ paper.
  
  The output of this layer :math:`\mathbf{y}` is defined implicitly via a
  potential function :math:`F(\mathbf{x}, \mathbf{y})`, a regularization term
  :math:`R(\mathbf{y})`, and the condition
  
  .. math::
      \mathbf{y} = \min_\mathbf{y} R(\mathbf{y}) + \sum_{i}
      F(\mathbf{x}_i, \mathbf{y}).
  
  The given implementation uses a ResNet-like model for the potential
  function and a simple :math:`L_2` norm :math:`R(\mathbf{y}) =
  \textrm{softplus}(\lambda) \cdot {\| \mathbf{y} \|}^2_2` for the
  regularizer with learnable weight :math:`\lambda`.
  
  Args:
      in_channels (int): Size of each input sample.
      out_channels (int): Size of each output sample.
      num_layers (List[int): List of hidden channels in the potential
          function.
      grad_iter (int): The number of steps to take in the internal gradient
          descent. (default: :obj:`5`)
      lamb (float): The initial regularization constant.
          (default: :obj:`0.1`)
class_docstr_tok: ["The","equilibrium","aggregation","layer","from","the","`","``","Equilibrium","Aggregation",":","Encoding","Sets","via","Optimization","''","<","https",":","\/\/arxiv.org\/abs\/2202.12795",">","`","_","paper",".","The","output","of","this","layer",":","math",":","`","\\mathbf","{","y","}","`","is","defined","implicitly","via","a","potential","function",":","math",":","`","F","(","\\mathbf","{","x","}",",","\\mathbf","{","y","}",")","`",",","a","regularization","term",":","math",":","`","R","(","\\mathbf","{","y","}",")","`",",","and","the","condition","..","math",":",":","\\mathbf","{","y","}","=","\\min_\\mathbf","{","y","}","R","(","\\mathbf","{","y","}",")","+","\\sum_","{","i","}","F","(","\\mathbf","{","x","}","_i",",","\\mathbf","{","y","}",")",".","The","given","implementation","uses","a","ResNet-like","model","for","the","potential","function","and","a","simple",":","math",":","`","L_2","`","norm",":","math",":","`","R","(","\\mathbf","{","y","}",")","=","\\textrm","{","softplus","}","(","\\lambda",")","\\cdot","{","\\|","\\mathbf","{","y","}","\\|","}","^2_2","`","for","the","regularizer","with","learnable","weight",":","math",":","`","\\lambda","`",".","Args",":","in_channels","(","int",")",":","Size","of","each","input","sample",".","out_channels","(","int",")",":","Size","of","each","output","sample",".","num_layers","(","List","[","int",")",":","List","of","hidden","channels","in","the","potential","function",".","grad_iter","(","int",")",":","The","number","of","steps","to","take","in","the","internal","gradient","descent",".","(","default",":",":","obj",":","`","5","`",")","lamb","(","float",")",":","The","initial","regularization","constant",".","(","default",":",":","obj",":","`","0.1","`",")"]
code: def energy(
      self,
      x: Tensor,
      y: Tensor,
      index: Optional[Tensor],
      dim_size: Optional[int] = None,
  ):
      return self.potential(x, y, index, dim_size) + self.reg(y)
  
imports: from typing import Callable, List, Optional, Tuple
  import torch
  from torch import Tensor
  from torch_geometric.nn.aggr import Aggregation
  from torch_geometric.nn.inits import reset
  from torch_geometric.utils import scatter
label_desc: Use image node_id 5 for calling the EquilibriumAggregation obj's underlying member method code with example usage: obj.energy(x, y, index, dim_size) and returns: unknown