text: from .fnn import FNN
  from .nn import NN
  from .. import activations
  from ...backend import tf
  
  
  class DeepONetCartesianProd(NN):
      """Deep operator network for dataset in the format of Cartesian product.
  
      Args:
          layer_size_branch: A list of integers as the width of a fully connected network,
              or `(dim, f)` where `dim` is the input dimension and `f` is a network
              function. The width of the last layer in the branch and trunk net should be
              equal.
          layer_size_trunk (list): A list of integers as the width of a fully connected
              network.
          activation: If `activation` is a ``string``, then the same activation is used in
              both trunk and branch nets. If `activation` is a ``dict``, then the trunk
              net uses the activation `activation["trunk"]`, and the branch net uses
              `activation["branch"]`.
      """
  
      def __init__(
          self,
          layer_sizes_branch,
          layer_sizes_trunk,
          activation,
          kernel_initializer,
      ):
          super(DeepONetCartesianProd, self).__init__()
          if isinstance(activation, dict):
              activation_branch = activation["branch"]
              self.activation_trunk = activations.get(activation["trunk"])
          else:
              activation_branch = self.activation_trunk = activations.get(activation)
  
          if callable(layer_sizes_branch[1]):
              # User-defined network
              self.branch = layer_sizes_branch[1]
          else:
              # Fully connected network
              self.branch = FNN(layer_sizes_branch, activation_branch, kernel_initializer)
          self.trunk = FNN(layer_sizes_trunk, self.activation_trunk, kernel_initializer)
          self.b = tf.Variable(tf.zeros(1))
  
      def call(self, inputs, training=False):
          x_func = inputs[0]
          x_loc = inputs[1]
  
          # Branch net to encode the input function
          x_func = self.branch(x_func)
          # Trunk net to encode the domain of the output function
          if self._input_transform is not None:
              x_loc = self._input_transform(x_loc)
          x_loc = self.activation_trunk(self.trunk(x_loc))
          # Dot product
          if x_func.shape[-1] != x_loc.shape[-1]:
              raise AssertionError(
                  "Output sizes of branch net and trunk net do not match."
              )
          x = tf.einsum("bi,ni->bn", x_func, x_loc)
          # Add bias
          x += self.b
  
          if self._output_transform is not None:
              x = self._output_transform(inputs, x)
  
          return x
  