iganet.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <any>
 
#include <boundary.hpp>
#include <functionspace.hpp>
#include <igabase.hpp>
#include <layer.hpp>
#include <optimizer.hpp>
#include <utils/container.hpp>
#include <utils/fqn.hpp>
#include <utils/tuple.hpp>
#include <utils/zip.hpp>
 
namespace iganet {
 
struct IgANetOptions {
  TORCH_ARG(int64_t, max_epoch) = 100;
  TORCH_ARG(int64_t, batch_size) = 1000;
  TORCH_ARG(double, min_loss) = 1e-4;
};
 
template <typename real_t>
class IgANetGeneratorImpl : public torch::nn::Module {
public:
  IgANetGeneratorImpl() = default;
 
  explicit IgANetGeneratorImpl(
      const std::vector<int64_t> &layers,
      const std::vector<std::vector<std::any>> &activations,
      Options<real_t> options = Options<real_t>{}) {
    assert(layers.size() == activations.size() + 1);
 
    // Generate vector of linear layers and register them as layer[i]
    for (auto i = 0; i < layers.size() - 1; ++i) {
      layers_.emplace_back(
          register_module("layer[" + std::to_string(i) + "]",
                          torch::nn::Linear(layers[i], layers[i + 1])));
      layers_.back()->to(options.device(), options.dtype(), true);
 
      torch::nn::init::xavier_uniform_(layers_.back()->weight);
      torch::nn::init::constant_(layers_.back()->bias, 0.0);
    }
 
    // Generate vector of activation functions
    for (const auto &a : activations)
      switch (std::any_cast<activation>(a[0])) {
        // No activation function
      case activation::none:
        switch (a.size()) {
        case 1:
          activations_.emplace_back(new None{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Batch Normalization
      case activation::batch_norm:
        switch (a.size()) {
        case 8:
          activations_.emplace_back(new BatchNorm{
              std::any_cast<torch::Tensor>(a[1]),
              std::any_cast<torch::Tensor>(a[2]),
              std::any_cast<torch::Tensor>(a[3]),
              std::any_cast<torch::Tensor>(a[4]), std::any_cast<double>(a[5]),
              std::any_cast<double>(a[6]), std::any_cast<bool>(a[7])});
          break;
        case 7:
          activations_.emplace_back(new BatchNorm{
              std::any_cast<torch::Tensor>(a[1]),
              std::any_cast<torch::Tensor>(a[2]),
              std::any_cast<torch::Tensor>(a[3]),
              std::any_cast<torch::Tensor>(a[4]), std::any_cast<double>(a[5]),
              std::any_cast<double>(a[6])});
          break;
        case 4:
          activations_.emplace_back(new BatchNorm{
              std::any_cast<torch::Tensor>(a[1]),
              std::any_cast<torch::Tensor>(a[2]),
              std::any_cast<torch::nn::functional::BatchNormFuncOptions>(
                  a[3])});
          break;
        case 3:
          activations_.emplace_back(
              new BatchNorm{std::any_cast<torch::Tensor>(a[1]),
                            std::any_cast<torch::Tensor>(a[2])});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // CELU
      case activation::celu:
        switch (a.size()) {
        case 3:
          activations_.emplace_back(
              new CELU{std::any_cast<double>(a[1]), std::any_cast<bool>(a[2])});
          break;
        case 2:
          try {
            activations_.emplace_back(new CELU{
                std::any_cast<torch::nn::functional::CELUFuncOptions>(a[1])});
          } catch (...) {
            activations_.emplace_back(new CELU{std::any_cast<double>(a[1])});
          }
          break;
        case 1:
          activations_.emplace_back(new CELU{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // ELU
      case activation::elu:
        switch (a.size()) {
        case 3:
          activations_.emplace_back(
              new ELU{std::any_cast<double>(a[1]), std::any_cast<bool>(a[2])});
          break;
        case 2:
          try {
            activations_.emplace_back(new ELU{
                std::any_cast<torch::nn::functional::ELUFuncOptions>(a[1])});
          } catch (...) {
            activations_.emplace_back(new ELU{std::any_cast<double>(a[1])});
          }
          break;
        case 1:
          activations_.emplace_back(new ELU{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // GELU
      case activation::gelu:
        switch (a.size()) {
        case 1:
          activations_.emplace_back(new GELU{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // GLU
      case activation::glu:
        switch (a.size()) {
        case 2:
          try {
            activations_.emplace_back(new GLU{
                std::any_cast<torch::nn::functional::GLUFuncOptions>(a[1])});
          } catch (...) {
            activations_.emplace_back(new GLU{std::any_cast<int64_t>(a[1])});
          }
          break;
        case 1:
          activations_.emplace_back(new GLU{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Group Normalization
      case activation::group_norm:
        switch (a.size()) {
        case 5:
          activations_.emplace_back(new GroupNorm{
              std::any_cast<int64_t>(a[1]), std::any_cast<torch::Tensor>(a[2]),
              std::any_cast<torch::Tensor>(a[3]), std::any_cast<double>(a[4])});
          break;
        case 2:
          try {
            activations_.emplace_back(new GroupNorm{
                std::any_cast<torch::nn::functional::GroupNormFuncOptions>(
                    a[1])});
          } catch (...) {
            activations_.emplace_back(
                new GroupNorm{std::any_cast<int64_t>(a[1])});
          }
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Gumbel-Softmax
      case activation::gumbel_softmax:
        switch (a.size()) {
        case 4:
          activations_.emplace_back(new GumbelSoftmax{
              std::any_cast<double>(a[1]), std::any_cast<int>(a[2]),
              std::any_cast<bool>(a[3])});
          break;
        case 2:
          activations_.emplace_back(new GumbelSoftmax{
              std::any_cast<torch::nn::functional::GumbelSoftmaxFuncOptions>(
                  a[1])});
          break;
        case 1:
          activations_.emplace_back(new GumbelSoftmax{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Hard shrinkish
      case activation::hardshrink:
        switch (a.size()) {
        case 2:
          try {
            activations_.emplace_back(new Hardshrink{
                std::any_cast<torch::nn::functional::HardshrinkFuncOptions>(
                    a[1])});
          } catch (...) {
            activations_.emplace_back(
                new Hardshrink{std::any_cast<double>(a[1])});
          }
          break;
        case 1:
          activations_.emplace_back(new Hardshrink{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Hardsigmoid
      case activation::hardsigmoid:
        switch (a.size()) {
        case 1:
          activations_.emplace_back(new Hardsigmoid{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Hardswish
      case activation::hardswish:
        switch (a.size()) {
        case 1:
          activations_.emplace_back(new Hardswish{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Hardtanh
      case activation::hardtanh:
        switch (a.size()) {
        case 4:
          activations_.emplace_back(new Hardtanh{std::any_cast<double>(a[1]),
                                                 std::any_cast<double>(a[2]),
                                                 std::any_cast<bool>(a[3])});
          break;
        case 3:
          activations_.emplace_back(new Hardtanh{std::any_cast<double>(a[1]),
                                                 std::any_cast<double>(a[2])});
          break;
        case 2:
          activations_.emplace_back(new Hardtanh{
              std::any_cast<torch::nn::functional::HardtanhFuncOptions>(a[1])});
          break;
        case 1:
          activations_.emplace_back(new Hardtanh{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Instance Normalization
      case activation::instance_norm:
        switch (a.size()) {
        case 8:
          activations_.emplace_back(new InstanceNorm{
              std::any_cast<torch::Tensor>(a[1]),
              std::any_cast<torch::Tensor>(a[2]),
              std::any_cast<torch::Tensor>(a[3]),
              std::any_cast<torch::Tensor>(a[4]), std::any_cast<double>(a[5]),
              std::any_cast<double>(a[6]), std::any_cast<bool>(a[7])});
          break;
        case 7:
          activations_.emplace_back(new InstanceNorm{
              std::any_cast<torch::Tensor>(a[1]),
              std::any_cast<torch::Tensor>(a[2]),
              std::any_cast<torch::Tensor>(a[3]),
              std::any_cast<torch::Tensor>(a[4]), std::any_cast<double>(a[5]),
              std::any_cast<double>(a[6])});
          break;
        case 2:
          activations_.emplace_back(new InstanceNorm{
              std::any_cast<torch::nn::functional::InstanceNormFuncOptions>(
                  a[1])});
          break;
        case 1:
          activations_.emplace_back(new InstanceNorm{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Layer Normalization
      case activation::layer_norm:
        switch (a.size()) {
        case 5:
          activations_.emplace_back(new LayerNorm{
              std::any_cast<std::vector<int64_t>>(a[1]),
              std::any_cast<torch::Tensor>(a[2]),
              std::any_cast<torch::Tensor>(a[3]), std::any_cast<double>(a[4])});
          break;
        case 2:
          try {
            activations_.emplace_back(new LayerNorm{
                std::any_cast<torch::nn::functional::LayerNormFuncOptions>(
                    a[1])});
          } catch (...) {
            activations_.emplace_back(
                new LayerNorm{std::any_cast<std::vector<int64_t>>(a[1])});
          }
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Leaky ReLU
      case activation::leaky_relu:
        switch (a.size()) {
        case 3:
          activations_.emplace_back(new LeakyReLU{std::any_cast<double>(a[1]),
                                                  std::any_cast<bool>(a[2])});
          break;
        case 2:
          try {
            activations_.emplace_back(new LeakyReLU{
                std::any_cast<torch::nn::functional::LeakyReLUFuncOptions>(
                    a[1])});
          } catch (...) {
            activations_.emplace_back(
                new LeakyReLU{std::any_cast<double>(a[1])});
          }
          break;
        case 1:
          activations_.emplace_back(new LeakyReLU{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Local response Normalization
      case activation::local_response_norm:
        switch (a.size()) {
        case 5:
          activations_.emplace_back(new LocalResponseNorm{
              std::any_cast<int64_t>(a[1]), std::any_cast<double>(a[2]),
              std::any_cast<double>(a[3]), std::any_cast<double>(a[4])});
          break;
        case 2:
          try {
            activations_.emplace_back(new LocalResponseNorm{std::any_cast<
                torch::nn::functional::LocalResponseNormFuncOptions>(a[1])});
          } catch (...) {
            activations_.emplace_back(
                new LocalResponseNorm{std::any_cast<int64_t>(a[1])});
          }
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // LogSigmoid
      case activation::logsigmoid:
        switch (a.size()) {
        case 1:
          activations_.emplace_back(new LogSigmoid{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // LogSoftmax
      case activation::logsoftmax:
        switch (a.size()) {
        case 2:
          try {
            activations_.emplace_back(new LogSoftmax{
                std::any_cast<torch::nn::functional::LogSoftmaxFuncOptions>(
                    a[1])});
          } catch (...) {
            activations_.emplace_back(
                new LogSoftmax{std::any_cast<int64_t>(a[1])});
          }
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Mish
      case activation::mish:
        switch (a.size()) {
        case 1:
          activations_.emplace_back(new Mish{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Lp Normalization
      case activation::normalize:
        switch (a.size()) {
        case 4:
          activations_.emplace_back(new Normalize{
              std::any_cast<double>(a[1]), std::any_cast<double>(a[2]),
              std::any_cast<int64_t>(a[3])});
          break;
        case 2:
          activations_.emplace_back(new Normalize{
              std::any_cast<torch::nn::functional::NormalizeFuncOptions>(
                  a[1])});
          break;
        case 1:
          activations_.emplace_back(new Normalize{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // PReLU
      case activation::prelu:
        switch (a.size()) {
        case 2:
          activations_.emplace_back(
              new PReLU{std::any_cast<torch::Tensor>(a[1])});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // ReLU
      case activation::relu:
        switch (a.size()) {
        case 2:
          try {
            activations_.emplace_back(new ReLU{
                std::any_cast<torch::nn::functional::ReLUFuncOptions>(a[1])});
          } catch (...) {
            activations_.emplace_back(new ReLU{std::any_cast<bool>(a[1])});
          }
          break;
        case 1:
          activations_.emplace_back(new ReLU{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Relu6
      case activation::relu6:
        switch (a.size()) {
        case 2:
          try {
            activations_.emplace_back(new ReLU6{
                std::any_cast<torch::nn::functional::ReLU6FuncOptions>(a[1])});
          } catch (...) {
            activations_.emplace_back(new ReLU6{std::any_cast<bool>(a[1])});
          }
          break;
        case 1:
          activations_.emplace_back(new ReLU6{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Randomized ReLU
      case activation::rrelu:
        switch (a.size()) {
        case 4:
          activations_.emplace_back(new RReLU{std::any_cast<double>(a[1]),
                                              std::any_cast<double>(a[2]),
                                              std::any_cast<bool>(a[3])});
          break;
        case 3:
          activations_.emplace_back(new RReLU{std::any_cast<double>(a[1]),
                                              std::any_cast<double>(a[2])});
          break;
        case 2:
          activations_.emplace_back(new RReLU{
              std::any_cast<torch::nn::functional::RReLUFuncOptions>(a[1])});
          break;
        case 1:
          activations_.emplace_back(new RReLU{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // SELU
      case activation::selu:
        switch (a.size()) {
        case 2:
          try {
            activations_.emplace_back(new SELU{
                std::any_cast<torch::nn::functional::SELUFuncOptions>(a[1])});
          } catch (...) {
            activations_.emplace_back(new SELU{std::any_cast<bool>(a[1])});
          }
          break;
        case 1:
          activations_.emplace_back(new SELU{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Sigmoid
      case activation::sigmoid:
        switch (a.size()) {
        case 1:
          activations_.emplace_back(new Sigmoid{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // SiLU
      case activation::silu:
        switch (a.size()) {
        case 1:
          activations_.emplace_back(new SiLU{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Softmax
      case activation::softmax:
        switch (a.size()) {
        case 2:
          try {
            activations_.emplace_back(new Softmax{
                std::any_cast<torch::nn::functional::SoftmaxFuncOptions>(
                    a[1])});
          } catch (...) {
            activations_.emplace_back(
                new Softmax{std::any_cast<int64_t>(a[1])});
          }
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Softmin
      case activation::softmin:
        switch (a.size()) {
        case 2:
          try {
            activations_.emplace_back(new Softmin{
                std::any_cast<torch::nn::functional::SoftminFuncOptions>(
                    a[1])});
          } catch (...) {
            activations_.emplace_back(
                new Softmin{std::any_cast<int64_t>(a[1])});
          }
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Softplus
      case activation::softplus:
        switch (a.size()) {
        case 3:
          activations_.emplace_back(new Softplus{std::any_cast<double>(a[1]),
                                                 std::any_cast<double>(a[2])});
          break;
        case 2:
          activations_.emplace_back(new Softplus{
              std::any_cast<torch::nn::functional::SoftplusFuncOptions>(a[1])});
          break;
        case 1:
          activations_.emplace_back(new Softplus{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Softshrink
      case activation::softshrink:
        switch (a.size()) {
        case 2:
          try {
            activations_.emplace_back(new Softshrink{
                std::any_cast<torch::nn::functional::SoftshrinkFuncOptions>(
                    a[1])});
          } catch (...) {
            activations_.emplace_back(
                new Softshrink{std::any_cast<double>(a[1])});
          }
          break;
        case 1:
          activations_.emplace_back(new Softshrink{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Softsign
      case activation::softsign:
        switch (a.size()) {
        case 1:
          activations_.emplace_back(new Softsign{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Tanh
      case activation::tanh:
        switch (a.size()) {
        case 1:
          activations_.emplace_back(new Tanh{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Tanhshrink
      case activation::tanhshrink:
        switch (a.size()) {
        case 1:
          activations_.emplace_back(new Tanhshrink{});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
        // Threshold
      case activation::threshold:
        switch (a.size()) {
        case 4:
          activations_.emplace_back(new Threshold{std::any_cast<double>(a[1]),
                                                  std::any_cast<double>(a[2]),
                                                  std::any_cast<bool>(a[3])});
          break;
        case 3:
          activations_.emplace_back(new Threshold{std::any_cast<double>(a[1]),
                                                  std::any_cast<double>(a[2])});
          break;
        case 2:
          activations_.emplace_back(new Threshold{
              std::any_cast<torch::nn::functional::ThresholdFuncOptions>(
                  a[1])});
          break;
        default:
          throw std::runtime_error("Invalid number of parameters");
        }
        break;
 
      default:
        throw std::runtime_error("Invalid activation function");
      }
  }
 
  torch::Tensor forward(torch::Tensor x) {
    torch::Tensor x_in = x.clone();
 
    // Standard feed-forward neural network
    for (auto [layer, activation] : utils::zip(layers_, activations_))
      x = activation->apply(layer->forward(x));
 
    return x;
  }
 
  inline torch::serialize::OutputArchive &
  write(torch::serialize::OutputArchive &archive,
        const std::string &key = "iganet") const {
    assert(layers_.size() == activations_.size());
 
    archive.write(key + ".layers", torch::full({1}, (int64_t)layers_.size()));
    for (std::size_t i = 0; i < layers_.size(); ++i) {
      archive.write(
          key + ".layer[" + std::to_string(i) + "].in_features",
          torch::full({1}, (int64_t)layers_[i]->options.in_features()));
      archive.write(
          key + ".layer[" + std::to_string(i) + "].outputs_features",
          torch::full({1}, (int64_t)layers_[i]->options.out_features()));
      archive.write(key + ".layer[" + std::to_string(i) + "].bias",
                    torch::full({1}, (int64_t)layers_[i]->options.bias()));
 
      activations_[i]->write(archive, key + ".layer[" + std::to_string(i) +
                                          "].activation");
    }
 
    return archive;
  }
 
  inline torch::serialize::InputArchive &
  read(torch::serialize::InputArchive &archive,
       const std::string &key = "iganet") {
    torch::Tensor layers, in_features, outputs_features, bias, activation;
 
    archive.read(key + ".layers", layers);
    for (int64_t i = 0; i < layers.item<int64_t>(); ++i) {
      archive.read(key + ".layer[" + std::to_string(i) + "].in_features",
                   in_features);
      archive.read(key + ".layer[" + std::to_string(i) + "].outputs_features",
                   outputs_features);
      archive.read(key + ".layer[" + std::to_string(i) + "].bias", bias);
      layers_.emplace_back(register_module(
          "layer[" + std::to_string(i) + "]",
          torch::nn::Linear(
              torch::nn::LinearOptions(in_features.item<int64_t>(),
                                       outputs_features.item<int64_t>())
                  .bias(bias.item<bool>()))));
 
      archive.read(key + ".layer[" + std::to_string(i) + "].activation.type",
                   activation);
      switch (static_cast<enum activation>(activation.item<int64_t>())) {
      case activation::none:
        activations_.emplace_back(new None{});
        break;
      case activation::batch_norm:
        activations_.emplace_back(
            new BatchNorm{torch::Tensor{}, torch::Tensor{}});
        break;
      case activation::celu:
        activations_.emplace_back(new CELU{});
        break;
      case activation::elu:
        activations_.emplace_back(new ELU{});
        break;
      case activation::gelu:
        activations_.emplace_back(new GELU{});
        break;
      case activation::glu:
        activations_.emplace_back(new GLU{});
        break;
      case activation::group_norm:
        activations_.emplace_back(new GroupNorm{0});
        break;
      case activation::gumbel_softmax:
        activations_.emplace_back(new GumbelSoftmax{});
        break;
      case activation::hardshrink:
        activations_.emplace_back(new Hardshrink{});
        break;
      case activation::hardsigmoid:
        activations_.emplace_back(new Hardsigmoid{});
        break;
      case activation::hardswish:
        activations_.emplace_back(new Hardswish{});
        break;
      case activation::hardtanh:
        activations_.emplace_back(new Hardtanh{});
        break;
      case activation::instance_norm:
        activations_.emplace_back(new InstanceNorm{});
        break;
      case activation::layer_norm:
        activations_.emplace_back(new LayerNorm{{}});
        break;
      case activation::leaky_relu:
        activations_.emplace_back(new LeakyReLU{});
        break;
      case activation::local_response_norm:
        activations_.emplace_back(new LocalResponseNorm{0});
        break;
      case activation::logsigmoid:
        activations_.emplace_back(new LogSigmoid{});
        break;
      case activation::logsoftmax:
        activations_.emplace_back(new LogSoftmax{0});
        break;
      case activation::mish:
        activations_.emplace_back(new Mish{});
        break;
      case activation::normalize:
        activations_.emplace_back(new Normalize{0, 0, 0});
        break;
      case activation::prelu:
        activations_.emplace_back(new PReLU{torch::Tensor{}});
        break;
      case activation::relu:
        activations_.emplace_back(new ReLU{});
        break;
      case activation::relu6:
        activations_.emplace_back(new ReLU6{});
        break;
      case activation::rrelu:
        activations_.emplace_back(new RReLU{});
        break;
      case activation::selu:
        activations_.emplace_back(new SELU{});
        break;
      case activation::sigmoid:
        activations_.emplace_back(new Sigmoid{});
        break;
      case activation::silu:
        activations_.emplace_back(new SiLU{});
        break;
      case activation::softmax:
        activations_.emplace_back(new Softmax{0});
        break;
      case activation::softmin:
        activations_.emplace_back(new Softmin{0});
        break;
      case activation::softplus:
        activations_.emplace_back(new Softplus{});
        break;
      case activation::softshrink:
        activations_.emplace_back(new Softshrink{});
        break;
      case activation::softsign:
        activations_.emplace_back(new Softsign{});
        break;
      case activation::tanh:
        activations_.emplace_back(new Tanh{});
        break;
      case activation::tanhshrink:
        activations_.emplace_back(new Tanhshrink{});
        break;
      case activation::threshold:
        activations_.emplace_back(new Threshold{0, 0});
        break;
      default:
        throw std::runtime_error("Invalid activation function");
      }
      activations_.back()->read(archive, key + ".layer[" + std::to_string(i) +
                                             "].activation");
    }
    return archive;
  }
 
  inline virtual void
  pretty_print(std::ostream &os = Log(log::info)) const noexcept override {
    os << "(\n";
 
    int i = 0;
    for (const auto &activation : activations_)
      os << "activation[" << i++ << "] = " << *activation << "\n";
    os << ")\n";
  }
 
private:
  std::vector<torch::nn::Linear> layers_;
 
  std::vector<std::unique_ptr<iganet::ActivationFunction>> activations_;
};
 
template <typename real_t>
class IgANetGenerator
    : public torch::nn::ModuleHolder<IgANetGeneratorImpl<real_t>> {
 
public:
  using torch::nn::ModuleHolder<IgANetGeneratorImpl<real_t>>::ModuleHolder;
  using Impl = IgANetGeneratorImpl<real_t>;
};
 
template <typename Optimizer, typename GeometryMap, typename Variable,
          template <typename, typename> typename IgABase = ::iganet::IgABase>
requires OptimizerType<Optimizer> && FunctionSpaceType<GeometryMap> && FunctionSpaceType<Variable>
class IgANet : public IgABase<GeometryMap, Variable>,
               utils::Serializable,
               private utils::FullQualifiedName {
public:
  using Base = IgABase<GeometryMap, Variable>;
 
  using optimizer_type = Optimizer;
 
  using optimizer_options_type = typename optimizer_options_type<Optimizer>::type;
  
protected:
  IgANetGenerator<typename Base::value_type> net_;
 
  std::unique_ptr<optimizer_type> opt_;
 
  IgANetOptions options_;
 
public:
  explicit IgANet(IgANetOptions defaults = {},
                  iganet::Options<typename Base::value_type> options =
                      iganet::Options<typename Base::value_type>{})
    : // Construct the base class
      Base(),
      // Construct the optimizer
      opt_(std::make_unique<optimizer_type>(net_->parameters())),
      // Set options
      options_(defaults) {}
 
  template <std::size_t NumCoeffs>
  IgANet(const std::vector<int64_t> &layers,
         const std::vector<std::vector<std::any>> &activations,
         std::array<int64_t, NumCoeffs> numCoeffs, IgANetOptions defaults = {},
         iganet::Options<typename Base::value_type> options =
             iganet::Options<typename Base::value_type>{})
      : IgANet(layers, activations, std::tuple{numCoeffs}, std::tuple{numCoeffs},
               defaults, options) {}
 
  template <std::size_t... NumCoeffs>
  IgANet(const std::vector<int64_t> &layers,
         const std::vector<std::vector<std::any>> &activations,
         std::tuple<std::array<int64_t, NumCoeffs>...> numCoeffs,
         IgANetOptions defaults = {},
         iganet::Options<typename Base::value_type> options =
             iganet::Options<typename Base::value_type>{})
      : IgANet(layers, activations, numCoeffs, numCoeffs, defaults, options) {}
 
  template <std::size_t GeometryMapNumCoeffs, std::size_t VariableNumCoeffs>
  IgANet(const std::vector<int64_t> &layers,
         const std::vector<std::vector<std::any>> &activations,
         std::array<int64_t, GeometryMapNumCoeffs> geometryMapNumCoeffs,
         std::array<int64_t, VariableNumCoeffs> variableNumCoeffs,
         IgANetOptions defaults = {},
         iganet::Options<typename Base::value_type> options =
             iganet::Options<typename Base::value_type>{})
      : IgANet(layers, activations, std::tuple{geometryMapNumCoeffs},
               std::tuple{variableNumCoeffs}, defaults, options) {}
 
  template <std::size_t... GeometryMapNumCoeffs,
            std::size_t... VariableNumCoeffs>
  IgANet(
      const std::vector<int64_t> &layers,
      const std::vector<std::vector<std::any>> &activations,
      std::tuple<std::array<int64_t, GeometryMapNumCoeffs>...>
          geometryMapNumCoeffs,
      std::tuple<std::array<int64_t, VariableNumCoeffs>...> variableNumCoeffs,
      IgANetOptions defaults = {},
      iganet::Options<typename Base::value_type> options =
          iganet::Options<typename Base::value_type>{})
      : // Construct the base class
        Base(geometryMapNumCoeffs, variableNumCoeffs, options),
        // Construct the deep neural network
        net_(utils::concat(std::vector<int64_t>{inputs(/* epoch */ 0).size(0)},
                           layers,
                           std::vector<int64_t>{Base::u_.as_tensor_size()}),
             activations, options),
 
        // Construct the optimizer
        opt_(std::make_unique<optimizer_type>(net_->parameters())),
 
        // Set options
        options_(defaults) {}
 
  inline const IgANetGenerator<typename Base::value_type> &net() const {
    return net_;
  }
 
  inline IgANetGenerator<typename Base::value_type> &net() { return net_; }
 
  inline const optimizer_type &optimizer() const { return *opt_; }
 
  inline optimizer_type &optimizer() { return *opt_; }
  
  inline void optimizerReset(bool resetOptions = true) {
    if (resetOptions)
      opt_ = std::make_unique<optimizer_type>(net_->parameters());
    else {
      std::vector<optimizer_options_type> options;
      for (auto & group : opt_->param_groups())
        options.push_back(static_cast<optimizer_options_type&>(group.options()));
      opt_ = std::make_unique<optimizer_type>(net_->parameters());
      for (auto [group, options] : utils::zip(opt_->param_groups(), options))
        static_cast<optimizer_options_type&>(group.options()) = options;
    }
  }
 
  inline void optimizerReset(const optimizer_options_type& optimizerOptions) {
    opt_ = std::make_unique<optimizer_type>(net_->parameters(), optimizerOptions);
  }
 
  inline optimizer_options_type &optimizerOptions(std::size_t param_group = 0) {
    if (param_group < opt_->param_groups().size())
      return static_cast<optimizer_options_type&>(opt_->param_groups()[param_group].options());
    else
      throw std::runtime_error("Index exceeds number of parameter groups");      
  }
 
  inline const optimizer_options_type &optimizerOptions(std::size_t param_group = 0) const {
    if (param_group < opt_->param_groups().size())
      return static_cast<optimizer_options_type&>(opt_->param_groups()[param_group].options());
    else
      throw std::runtime_error("Index exceeds number of parameter groups");      
  }
 
  inline void optimizerOptionsReset(const optimizer_options_type& options) {
    for (auto &group : opt_->param_groups())
      static_cast<optimizer_options_type&>(group.options()) = options;
  }
 
  inline void optimizerOptionsReset(optimizer_options_type&& options) {
    for (auto &group : opt_->param_groups())
      static_cast<optimizer_options_type&>(group.options()) = options;
  }
 
  inline void optimizerOptionsReset(const optimizer_options_type& options, std::size_t param_group) {
    if (param_group < opt_->param_groups().size())
      static_cast<optimizer_options_type&>(opt_->param_group().options()) = options;
    else
      throw std::runtime_error("Index exceeds number of parameter groups");      
  }
 
  inline void optimizerOptionsReset(optimizer_options_type&& options, std::size_t param_group) {
    if (param_group < opt_->param_groups().size())
      static_cast<optimizer_options_type&>(opt_->param_group().options()) = options;
    else
      throw std::runtime_error("Index exceeds number of parameter groups");      
  }
  
  inline const auto &options() const { return options_; }
 
  inline auto &options() { return options_; }
 
  virtual torch::Tensor inputs(int64_t epoch) const {
    if constexpr (Base::has_GeometryMap && Base::has_RefData)
      return torch::cat({Base::G_.as_tensor(), Base::f_.as_tensor()});
    else if constexpr (Base::has_GeometryMap && !Base::has_RefData)
      return Base::G_.as_tensor();
    else if constexpr (!Base::has_GeometryMap && Base::has_RefData)
      return Base::f_.as_tensor();
    else
      return torch::empty({0});
  }
 
  virtual bool epoch(int64_t) = 0;
 
  virtual torch::Tensor loss(const torch::Tensor &, int64_t) = 0;
 
  virtual void train(
#ifdef IGANET_WITH_MPI
      c10::intrusive_ptr<c10d::ProcessGroupMPI> pg =
          c10d::ProcessGroupMPI::createProcessGroupMPI()
#endif
  ) {
    torch::Tensor inputs, outputs, loss;
    typename Base::value_type previous_loss(-1.0);
 
    // Loop over epochs
    for (int64_t epoch = 0; epoch != options_.max_epoch(); ++epoch) {
 
      // Update epoch and inputs
      if (this->epoch(epoch))
        inputs = this->inputs(epoch);
 
      auto closure = [&]() {
        // Reset gradients
        net_->zero_grad();
 
        // Execute the model on the inputs
        outputs = net_->forward(inputs);
 
        // Compute the loss value
        loss = this->loss(outputs, epoch);
 
        // Compute gradients of the loss w.r.t. the model parameters
        loss.backward({}, true, false);
 
        return loss;
      };
 
#ifdef IGANET_WITH_MPI
      // Averaging the gradients of the parameters in all the processors
      // Note: This may lag behind DistributedDataParallel (DDP) in performance
      // since this synchronizes parameters after backward pass while DDP
      // overlaps synchronizing parameters and computing gradients in backward
      // pass
      std::vector<c10::intrusive_ptr<::c10d::Work>> works;
      for (auto &param : net_->named_parameters()) {
        std::vector<torch::Tensor> tmp = {param.value().grad()};
        works.emplace_back(pg->allreduce(tmp));
      }
 
      waitWork(pg, works);
 
      for (auto &param : net_->named_parameters()) {
        param.value().grad().data() =
            param.value().grad().data() / pg->getSize();
      }
#endif
 
      // Update the parameters based on the calculated gradients
      opt_->step(closure);
 
      typename Base::value_type current_loss = loss.template item<typename Base::value_type>();
      Log(log::verbose) << "Epoch " << std::to_string(epoch) << ": "
                        << current_loss
                        << std::endl;
 
      if (current_loss <
          options_.min_loss()) {
        Log(log::info) << "Total epochs: " << epoch << ", loss: "
                       << current_loss
                       << std::endl;
        break;
      }
 
      if (current_loss == previous_loss || std::abs(current_loss-previous_loss) < previous_loss/10) {
        Log(log::info) << "Total epochs: " << epoch << ", loss: "
                       << current_loss
                       << std::endl;
        break;
      }
 
      if (loss.isnan().template item<bool>()) {
        Log(log::info) << "Total epochs: " << epoch << ", loss: "
        << current_loss
        << std::endl;
        break;
      }
      previous_loss = current_loss;
    }
  }
 
  template <typename DataLoader>
  void train(DataLoader &loader
#ifdef IGANET_WITH_MPI
             ,
             c10::intrusive_ptr<c10d::ProcessGroupMPI> pg =
                 c10d::ProcessGroupMPI::createProcessGroupMPI()
#endif
  ) {
    torch::Tensor inputs, outputs, loss;
    typename Base::value_type previous_loss(-1.0);
 
    // Loop over epochs
    for (int64_t epoch = 0; epoch != options_.max_epoch(); ++epoch) {
 
      typename Base::value_type Loss(0);
 
      for (auto &batch : loader) {
        inputs = batch.data;
 
        if (inputs.dim() > 0) {
          if constexpr (Base::has_GeometryMap && Base::has_RefData) {
            Base::G_.from_tensor(
                inputs.slice(1, 0, Base::G_.as_tensor_size()).t());
            Base::f_.from_tensor(inputs
                                     .slice(1, Base::G_.as_tensor_size(),
                                            Base::G_.as_tensor_size() +
                                                Base::f_.as_tensor_size())
                                     .t());
          } else if constexpr (Base::has_GeometryMap && !Base::has_RefData)
            Base::G_.from_tensor(
                inputs.slice(1, 0, Base::G_.as_tensor_size()).t());
          else if constexpr (!Base::has_GeometryMap && Base::has_RefData)
            Base::f_.from_tensor(
                inputs.slice(1, 0, Base::f_.as_tensor_size()).t());
 
        } else {
          if constexpr (Base::has_GeometryMap && Base::has_RefData) {
            Base::G_.from_tensor(
                inputs.slice(1, 0, Base::G_.as_tensor_size()).flatten());
            Base::f_.from_tensor(inputs
                                     .slice(1, Base::G_.as_tensor_size(),
                                            Base::G_.as_tensor_size() +
                                                Base::f_.as_tensor_size())
                                     .flatten());
          } else if constexpr (Base::has_GeometryMap && !Base::has_RefData)
            Base::G_.from_tensor(
                inputs.slice(1, 0, Base::G_.as_tensor_size()).flatten());
          else if constexpr (!Base::has_GeometryMap && Base::has_RefData)
            Base::f_.from_tensor(
                inputs.slice(1, 0, Base::f_.as_tensor_size()).flatten());
        }
 
        this->epoch(epoch);
 
        auto closure = [&]() {
          // Reset gradients
          net_->zero_grad();
 
          // Execute the model on the inputs
          outputs = net_->forward(inputs);
 
          // Compute the loss value
          loss = this->loss(outputs, epoch);
 
          // Compute gradients of the loss w.r.t. the model parameters
          loss.backward({}, true, false);
 
          return loss;
        };
 
        // Update the parameters based on the calculated gradients
        opt_->step(closure);
 
        Loss += loss.template item<typename Base::value_type>();
      }
 
      Log(log::verbose) << "Epoch " << std::to_string(epoch) << ": " << Loss
                        << std::endl;
 
      if (Loss < options_.min_loss()) {
        Log(log::info) << "Total epochs: " << epoch << ", loss: " << Loss
                       << std::endl;
        break;
      }
 
      if (Loss == previous_loss) {
        Log(log::info) << "Total epochs: " << epoch << ", loss: " << Loss
                       << std::endl;
        break;
      }
      previous_loss = Loss;
 
      if (epoch == options_.max_epoch() - 1)
        Log(log::warning) << "Total epochs: " << epoch << ", loss: " << Loss
                          << std::endl;
    }
  }
 
  void eval() {
    torch::Tensor inputs = this->inputs(0);
    torch::Tensor outputs = net_->forward(inputs);
    Base::u_.from_tensor(outputs);
  }
 
  inline virtual nlohmann::json to_json() const override {
    return "Not implemented yet";
  }
 
  inline std::vector<torch::Tensor> parameters() const noexcept {
    return net_->parameters();
  }
 
  inline torch::OrderedDict<std::string, torch::Tensor>
  named_parameters() const noexcept {
    return net_->named_parameters();
  }
 
  inline std::size_t nparameters() const noexcept {
    std::size_t result = 0;
    for (const auto &param : this->parameters()) {
      result += param.numel();
    }
    return result;
  }
 
  inline virtual void
  pretty_print(std::ostream &os = Log(log::info)) const noexcept override {
    os << name() << "(\n"
       << "net = " << net_ << "\n";
    if constexpr (Base::has_GeometryMap)
      os << "G = " << Base::G_ << "\n";
    if constexpr (Base::has_RefData)
      os << "f = " << Base::f_ << "\n";
    if constexpr (Base::has_Solution)
      os << "u = " << Base::u_ << "\n)";
  }
 
  inline void save(const std::string &filename,
                   const std::string &key = "iganet") const {
    torch::serialize::OutputArchive archive;
    write(archive, key).save_to(filename);
  }
 
  inline void load(const std::string &filename,
                   const std::string &key = "iganet") {
    torch::serialize::InputArchive archive;
    archive.load_from(filename);
    read(archive, key);
  }
 
  inline torch::serialize::OutputArchive &
  write(torch::serialize::OutputArchive &archive,
        const std::string &key = "iganet") const {
    if constexpr (Base::has_GeometryMap)
      Base::G_.write(archive, key + ".geo");
    if constexpr (Base::has_RefData)
      Base::f_.write(archive, key + ".ref");
    if constexpr (Base::has_Solution)
      Base::u_.write(archive, key + ".out");
 
    net_->write(archive, key + ".net");
    torch::serialize::OutputArchive archive_net;
    net_->save(archive_net);
    archive.write(key + ".net.data", archive_net);
 
    torch::serialize::OutputArchive archive_opt;
    opt_->save(archive_opt);
    archive.write(key + ".opt", archive_opt);
 
    return archive;
  }
 
  inline torch::serialize::InputArchive &
  read(torch::serialize::InputArchive &archive,
       const std::string &key = "iganet") {
    if constexpr (Base::has_GeometryMap)
      Base::G_.read(archive, key + ".geo");
    if constexpr (Base::has_RefData)
      Base::f_.read(archive, key + ".ref");
    if constexpr (Base::has_Solution)
      Base::u_.read(archive, key + ".out");
 
    net_->read(archive, key + ".net");
    torch::serialize::InputArchive archive_net;
    archive.read(key + ".net.data", archive_net);
    net_->load(archive_net);
 
    opt_->add_parameters(net_->parameters());
    torch::serialize::InputArchive archive_opt;
    archive.read(key + ".opt", archive_opt);
    opt_->load(archive_opt);
 
    return archive;
  }
 
  bool operator==(const IgANet &other) const {
    bool result(true);
 
    if constexpr (Base::has_GeometryMap)
      result *= (Base::G_ == other.G());
    if constexpr (Base::has_RefData)
      result *= (Base::f_ == other.f());
    if constexpr (Base::has_Solution)
      result *= (Base::u_ == other.u());
 
    return result;
  }
 
  bool operator!=(const IgANet &other) const { return *this != other; }
 
#ifdef IGANET_WITH_MPI
private:
  static void waitWork(c10::intrusive_ptr<c10d::ProcessGroupMPI> pg,
                       std::vector<c10::intrusive_ptr<c10d::Work>> works) {
    for (auto &work : works) {
      try {
        work->wait();
      } catch (const std::exception &ex) {
        Log(log::error) << "Exception received during waitWork: " << ex.what()
                        << std::endl;
        pg->abort();
      }
    }
  }
#endif
};
 
  template <typename Optimizer, typename GeometryMap, typename Variable>
  requires OptimizerType<Optimizer> && FunctionSpaceType<GeometryMap> && FunctionSpaceType<Variable>
inline std::ostream &
operator<<(std::ostream &os,
           const IgANet<Optimizer, GeometryMap, Variable> &obj) {
  obj.pretty_print(os);
  return os;
}
 
  template <typename GeometryMap, typename Variable>
  requires FunctionSpaceType<GeometryMap> && FunctionSpaceType<Variable>
  class IgANetCustomizable {
public:
  using geometryMap_interior_knot_indices_type =
      decltype(std::declval<GeometryMap>()
                   .template find_knot_indices<functionspace::interior>(
                       std::declval<typename GeometryMap::eval_type>()));
 
  using geometryMap_boundary_knot_indices_type =
      decltype(std::declval<GeometryMap>()
                   .template find_knot_indices<functionspace::boundary>(
                       std::declval<
                           typename GeometryMap::boundary_eval_type>()));
 
  using variable_interior_knot_indices_type =
      decltype(std::declval<Variable>()
                   .template find_knot_indices<functionspace::interior>(
                       std::declval<typename Variable::eval_type>()));
 
  using variable_boundary_knot_indices_type =
      decltype(std::declval<Variable>()
                   .template find_knot_indices<functionspace::boundary>(
                       std::declval<typename Variable::boundary_eval_type>()));
 
  using geometryMap_interior_coeff_indices_type =
      decltype(std::declval<GeometryMap>()
                   .template find_coeff_indices<functionspace::interior>(
                       std::declval<typename GeometryMap::eval_type>()));
 
  using geometryMap_boundary_coeff_indices_type =
      decltype(std::declval<GeometryMap>()
                   .template find_coeff_indices<functionspace::boundary>(
                       std::declval<
                           typename GeometryMap::boundary_eval_type>()));
 
  using variable_interior_coeff_indices_type =
      decltype(std::declval<Variable>()
                   .template find_coeff_indices<functionspace::interior>(
                       std::declval<typename Variable::eval_type>()));
 
  using variable_boundary_coeff_indices_type =
      decltype(std::declval<Variable>()
                   .template find_coeff_indices<functionspace::boundary>(
                       std::declval<typename Variable::boundary_eval_type>()));
};
 
template <typename Optimizer, typename Inputs, typename Outputs, typename CollPts = void>
requires OptimizerType<Optimizer>
class IgANet2 : public IgABase2<Inputs, Outputs, CollPts>,
                utils::Serializable,
                private utils::FullQualifiedName {
public:
    using Base = IgABase2<Inputs, Outputs, CollPts>;
 
  using optimizer_type = Optimizer;
 
  using optimizer_options_type = typename optimizer_options_type<Optimizer>::type;
  
protected:
  IgANetGenerator<typename Base::value_type> net_;
 
  std::unique_ptr<optimizer_type> opt_;
 
  IgANetOptions options_;
 
public:
  explicit IgANet2(IgANetOptions defaults = {},
                  iganet::Options<typename Base::value_type> options =
                      iganet::Options<typename Base::value_type>{})
    : // Construct the base class
      Base(),
      // Construct the optimizer
      opt_(std::make_unique<optimizer_type>(net_->parameters())),
      // Set options
      options_(defaults) {}
 
  template <typename NumCoeffs>
  IgANet2(const std::vector<int64_t> &layers,
          const std::vector<std::vector<std::any>> &activations,
          const NumCoeffs &numCoeffs,
          enum init init = init::greville,
          IgANetOptions defaults = {},
          iganet::Options<typename Base::value_type> options =
          iganet::Options<typename Base::value_type>{})
    : IgANet2(layers, activations, numCoeffs, numCoeffs, init, defaults, options)
  {}
 
  template <typename NumCoeffsInputs, typename NumCoeffsOutputs>
  IgANet2(const std::vector<int64_t> &layers,
          const std::vector<std::vector<std::any>> &activations,
          const NumCoeffsInputs &numCoeffsInputs,
          const NumCoeffsOutputs &numCoeffsOutputs,
          enum init init = init::greville,
          IgANetOptions defaults = {},
          iganet::Options<typename Base::value_type> options =
          iganet::Options<typename Base::value_type>{})
    : // Construct the base class
    Base(numCoeffsInputs, numCoeffsOutputs, init, options),
    // Construct the deep neural network
    net_(utils::concat(std::vector<int64_t>{inputs(/* epoch */ 0).size(0)},
                       layers,
                       std::vector<int64_t>{outputs(/* epoch */ 0).size(0)}),
         activations, options),
    
    // Construct the optimizer
    opt_(std::make_unique<optimizer_type>(net_->parameters())),
    
    // Set options
    options_(defaults) {} 
  
  inline const IgANetGenerator<typename Base::value_type> &net() const {
    return net_;
  }
 
  inline IgANetGenerator<typename Base::value_type> &net() { return net_; }
 
  inline const optimizer_type &optimizer() const { return *opt_; }
 
  inline optimizer_type &optimizer() { return *opt_; }
  
  inline void optimizerReset(bool resetOptions = true) {
    if (resetOptions)
      opt_ = std::make_unique<optimizer_type>(net_->parameters());
    else {
      std::vector<optimizer_options_type> options;
      for (auto & group : opt_->param_groups())
        options.push_back(static_cast<optimizer_options_type&>(group.options()));
      opt_ = std::make_unique<optimizer_type>(net_->parameters());
      for (auto [group, options] : utils::zip(opt_->param_groups(), options))
        static_cast<optimizer_options_type&>(group.options()) = options;
    }
  }
 
  inline void optimizerReset(const optimizer_options_type& optimizerOptions) {
    opt_ = std::make_unique<optimizer_type>(net_->parameters(), optimizerOptions);
  }
 
  inline optimizer_options_type &optimizerOptions(std::size_t param_group = 0) {
    if (param_group < opt_->param_groups().size())
      return static_cast<optimizer_options_type&>(opt_->param_groups()[param_group].options());
    else
      throw std::runtime_error("Index exceeds number of parameter groups");      
  }
 
  inline const optimizer_options_type &optimizerOptions(std::size_t param_group = 0) const {
    if (param_group < opt_->param_groups().size())
      return static_cast<optimizer_options_type&>(opt_->param_groups()[param_group].options());
    else
      throw std::runtime_error("Index exceeds number of parameter groups");      
  }
 
  inline void optimizerOptionsReset(const optimizer_options_type& options) {
    for (auto &group : opt_->param_groups())
      static_cast<optimizer_options_type&>(group.options()) = options;
  }
 
  inline void optimizerOptionsReset(optimizer_options_type&& options) {
    for (auto &group : opt_->param_groups())
      static_cast<optimizer_options_type&>(group.options()) = options;
  }
 
  inline void optimizerOptionsReset(const optimizer_options_type& options, std::size_t param_group) {
    if (param_group < opt_->param_groups().size())
      static_cast<optimizer_options_type&>(opt_->param_group().options()) = options;
    else
      throw std::runtime_error("Index exceeds number of parameter groups");      
  }
 
  inline void optimizerOptionsReset(optimizer_options_type&& options, std::size_t param_group) {
    if (param_group < opt_->param_groups().size())
      static_cast<optimizer_options_type&>(opt_->param_group().options()) = options;
    else
      throw std::runtime_error("Index exceeds number of parameter groups");      
  }
  
  inline const auto &options() const { return options_; }
 
  inline auto &options() { return options_; }
 
  virtual torch::Tensor inputs(int64_t epoch) const {
    return utils::cat_tuple_into_tensor(Base::inputs_, [](const auto& obj){ return obj.as_tensor(); });
  }
 
  virtual torch::Tensor outputs(int64_t epoch) const {
    return utils::cat_tuple_into_tensor(Base::outputs_, [](const auto& obj){ return obj.as_tensor(); });
  }
 
  virtual void inputs(const torch::Tensor& tensor) {
    utils::slice_tensor_into_tuple(Base::inputs_, tensor,
                                   [](const auto& obj){ return obj.as_tensor_size(); },
                                   [](auto& obj, const auto& tensor){ return obj.from_tensor(tensor); });
  }
 
  virtual void outputs(const torch::Tensor& tensor) {
    utils::slice_tensor_into_tuple(Base::outputs_, tensor,
                                   [](const auto& obj){ return obj.as_tensor_size(); },
                                   [](auto& obj, const auto& tensor){ return obj.from_tensor(tensor); });
  }
  
  virtual bool epoch(int64_t) = 0;
 
  virtual torch::Tensor loss(const torch::Tensor &, int64_t) = 0;
 
  virtual void train(
#ifdef IGANET_WITH_MPI
      c10::intrusive_ptr<c10d::ProcessGroupMPI> pg =
          c10d::ProcessGroupMPI::createProcessGroupMPI()
#endif
  ) {
    torch::Tensor inputs, outputs, loss;
    typename Base::value_type previous_loss(-1.0);
 
    // Loop over epochs
    for (int64_t epoch = 0; epoch != options_.max_epoch(); ++epoch) {
 
      // Update epoch and inputs
      if (this->epoch(epoch))
        inputs = this->inputs(epoch);
 
      auto closure = [&]() {
        // Reset gradients
        net_->zero_grad();
 
        // Execute the model on the inputs
        outputs = net_->forward(inputs);
 
        // Compute the loss value
        loss = this->loss(outputs, epoch);
 
        // Compute gradients of the loss w.r.t. the model parameters
        loss.backward({}, true, false);
 
        return loss;
      };
 
#ifdef IGANET_WITH_MPI
      // Averaging the gradients of the parameters in all the processors
      // Note: This may lag behind DistributedDataParallel (DDP) in performance
      // since this synchronizes parameters after backward pass while DDP
      // overlaps synchronizing parameters and computing gradients in backward
      // pass
      std::vector<c10::intrusive_ptr<::c10d::Work>> works;
      for (auto &param : net_->named_parameters()) {
        std::vector<torch::Tensor> tmp = {param.value().grad()};
        works.emplace_back(pg->allreduce(tmp));
      }
 
      waitWork(pg, works);
 
      for (auto &param : net_->named_parameters()) {
        param.value().grad().data() =
            param.value().grad().data() / pg->getSize();
      }
#endif
 
      // Update the parameters based on the calculated gradients
      opt_->step(closure);
 
      typename Base::value_type current_loss = loss.template item<typename Base::value_type>();
      Log(log::verbose) << "Epoch " << std::to_string(epoch) << ": "
                        << current_loss
                        << std::endl;
 
      if (current_loss <
          options_.min_loss()) {
        Log(log::info) << "Total epochs: " << epoch << ", loss: "
                       << current_loss
                       << std::endl;
        break;
      }
 
      if (current_loss == previous_loss || std::abs(current_loss-previous_loss) < previous_loss/10) {
        Log(log::info) << "Total epochs: " << epoch << ", loss: "
                       << current_loss
                       << std::endl;
        break;
      }
 
      if (loss.isnan().template item<bool>()) {
        Log(log::info) << "Total epochs: " << epoch << ", loss: "
        << current_loss
        << std::endl;
        break;
      }
      previous_loss = current_loss;
    }
  }
 
  template <typename DataLoader>
  void train(DataLoader &loader
#ifdef IGANET_WITH_MPI
             ,
             c10::intrusive_ptr<c10d::ProcessGroupMPI> pg =
                 c10d::ProcessGroupMPI::createProcessGroupMPI()
#endif
  ) {
    torch::Tensor inputs, outputs, loss;
    typename Base::value_type previous_loss(-1.0);
 
    // Loop over epochs
    for (int64_t epoch = 0; epoch != options_.max_epoch(); ++epoch) {
 
      typename Base::value_type Loss(0);
 
      for (auto &batch : loader) {
        inputs = batch.data;
 
        if (inputs.dim() > 0) {
          // if constexpr (Base::has_GeometryMap && Base::has_RefData) {
          //   Base::G_.from_tensor(
          //       inputs.slice(1, 0, Base::G_.as_tensor_size()).t());
          //   Base::f_.from_tensor(inputs
          //                            .slice(1, Base::G_.as_tensor_size(),
          //                                   Base::G_.as_tensor_size() +
          //                                       Base::f_.as_tensor_size())
          //                            .t());
          // } else if constexpr (Base::has_GeometryMap && !Base::has_RefData)
          //   Base::G_.from_tensor(
          //       inputs.slice(1, 0, Base::G_.as_tensor_size()).t());
          // else if constexpr (!Base::has_GeometryMap && Base::has_RefData)
          //   Base::f_.from_tensor(
          //       inputs.slice(1, 0, Base::f_.as_tensor_size()).t());
 
        } else {
          // if constexpr (Base::has_GeometryMap && Base::has_RefData) {
          //   Base::G_.from_tensor(
          //       inputs.slice(1, 0, Base::G_.as_tensor_size()).flatten());
          //   Base::f_.from_tensor(inputs
          //                            .slice(1, Base::G_.as_tensor_size(),
          //                                   Base::G_.as_tensor_size() +
          //                                       Base::f_.as_tensor_size())
          //                            .flatten());
          // } else if constexpr (Base::has_GeometryMap && !Base::has_RefData)
          //   Base::G_.from_tensor(
          //       inputs.slice(1, 0, Base::G_.as_tensor_size()).flatten());
          // else if constexpr (!Base::has_GeometryMap && Base::has_RefData)
          //   Base::f_.from_tensor(
          //       inputs.slice(1, 0, Base::f_.as_tensor_size()).flatten());
        }
 
        this->epoch(epoch);
 
        auto closure = [&]() {
          // Reset gradients
          net_->zero_grad();
 
          // Execute the model on the inputs
          outputs = net_->forward(inputs);
 
          // Compute the loss value
          loss = this->loss(outputs, epoch);
 
          // Compute gradients of the loss w.r.t. the model parameters
          loss.backward({}, true, false);
 
          return loss;
        };
 
        // Update the parameters based on the calculated gradients
        opt_->step(closure);
 
        Loss += loss.template item<typename Base::value_type>();
      }
 
      Log(log::verbose) << "Epoch " << std::to_string(epoch) << ": " << Loss
                        << std::endl;
 
      if (Loss < options_.min_loss()) {
        Log(log::info) << "Total epochs: " << epoch << ", loss: " << Loss
                       << std::endl;
        break;
      }
 
      if (Loss == previous_loss) {
        Log(log::info) << "Total epochs: " << epoch << ", loss: " << Loss
                       << std::endl;
        break;
      }
      previous_loss = Loss;
 
      if (epoch == options_.max_epoch() - 1)
        Log(log::warning) << "Total epochs: " << epoch << ", loss: " << Loss
                          << std::endl;
    }
  }
 
  void eval() {
    torch::Tensor inputs = this->inputs(0);
    torch::Tensor outputs = net_->forward(inputs);
    this->outputs(outputs);
  }
 
  inline virtual nlohmann::json to_json() const override {
    return "Not implemented yet";
  }
 
  inline std::vector<torch::Tensor> parameters() const noexcept {
    return net_->parameters();
  }
 
  inline torch::OrderedDict<std::string, torch::Tensor>
  named_parameters() const noexcept {
    return net_->named_parameters();
  }
 
  inline std::size_t nparameters() const noexcept {
    std::size_t result = 0;
    for (const auto &param : this->parameters()) {
      result += param.numel();
    }
    return result;
  }
 
  inline virtual void
  pretty_print(std::ostream &os = Log(log::info)) const noexcept override {
    os << name() << "(\n"
       << "net = " << net_ << "\n";
    // if constexpr (Base::has_GeometryMap)
    //   os << "G = " << Base::G_ << "\n";
    // if constexpr (Base::has_RefData)
    //   os << "f = " << Base::f_ << "\n";
    // if constexpr (Base::has_Solution)
    //   os << "u = " << Base::u_ << "\n)";
  }
 
  inline void save(const std::string &filename,
                   const std::string &key = "iganet") const {
    torch::serialize::OutputArchive archive;
    write(archive, key).save_to(filename);
  }
 
  inline void load(const std::string &filename,
                   const std::string &key = "iganet") {
    torch::serialize::InputArchive archive;
    archive.load_from(filename);
    read(archive, key);
  }
 
  inline torch::serialize::OutputArchive &
  write(torch::serialize::OutputArchive &archive,
        const std::string &key = "iganet") const {
    // if constexpr (Base::has_GeometryMap)
    //   Base::G_.write(archive, key + ".geo");
    // if constexpr (Base::has_RefData)
    //   Base::f_.write(archive, key + ".ref");
    // if constexpr (Base::has_Solution)
    //   Base::u_.write(archive, key + ".out");
 
    net_->write(archive, key + ".net");
    torch::serialize::OutputArchive archive_net;
    net_->save(archive_net);
    archive.write(key + ".net.data", archive_net);
 
    torch::serialize::OutputArchive archive_opt;
    opt_->save(archive_opt);
    archive.write(key + ".opt", archive_opt);
 
    return archive;
  }
 
  inline torch::serialize::InputArchive &
  read(torch::serialize::InputArchive &archive,
       const std::string &key = "iganet") {
    // if constexpr (Base::has_GeometryMap)
    //   Base::G_.read(archive, key + ".geo");
    // if constexpr (Base::has_RefData)
    //   Base::f_.read(archive, key + ".ref");
    // if constexpr (Base::has_Solution)
    //   Base::u_.read(archive, key + ".out");
 
    net_->read(archive, key + ".net");
    torch::serialize::InputArchive archive_net;
    archive.read(key + ".net.data", archive_net);
    net_->load(archive_net);
 
    opt_->add_parameters(net_->parameters());
    torch::serialize::InputArchive archive_opt;
    archive.read(key + ".opt", archive_opt);
    opt_->load(archive_opt);
 
    return archive;
  }
 
  bool operator==(const IgANet2 &other) const {
    bool result(true);
 
    // if constexpr (Base::has_GeometryMap)
    //   result *= (Base::G_ == other.G());
    // if constexpr (Base::has_RefData)
    //   result *= (Base::f_ == other.f());
    // if constexpr (Base::has_Solution)
    //   result *= (Base::u_ == other.u());
 
    return result;
  }
 
  bool operator!=(const IgANet2 &other) const { return *this != other; }
 
#ifdef IGANET_WITH_MPI
private:
  static void waitWork(c10::intrusive_ptr<c10d::ProcessGroupMPI> pg,
                       std::vector<c10::intrusive_ptr<c10d::Work>> works) {
    for (auto &work : works) {
      try {
        work->wait();
      } catch (const std::exception &ex) {
        Log(log::error) << "Exception received during waitWork: " << ex.what()
                        << std::endl;
        pg->abort();
      }
    }
  }
#endif
};
 
template <typename Optimizer, typename Inputs, typename Outputs, typename CollPts>
requires OptimizerType<Optimizer>
inline std::ostream &
operator<<(std::ostream &os,
           const IgANet2<Optimizer, Inputs, Outputs, CollPts> &obj) {
  //  obj.pretty_print(os);
  return os;
}
 
template <typename, typename, typename = void>
class IgANetCustomizable2;
 
template <detail::HasAsTensor... Inputs,
          detail::HasAsTensor... Outputs>
class IgANetCustomizable2<std::tuple<Inputs...>,
                    std::tuple<Outputs...>, void> {
public:
  using inputs_interior_knot_indices_type =
    std::tuple<decltype(std::declval<Inputs>()
                   .template find_knot_indices<functionspace::interior>(
                                                                        std::declval<typename Inputs::eval_type>()))...>;
 
  template<std::size_t index>
  using input_interior_knot_indices_t = typename std::tuple_element_t<index, inputs_interior_knot_indices_type>;
  
  using inputs_boundary_knot_indices_type =
    std::tuple<decltype(std::declval<Inputs>()
                        .template find_knot_indices<functionspace::boundary>(
                                                                             std::declval<
                                                                             typename Inputs::boundary_eval_type>()))...>;
 
  template<std::size_t index>
  using input_boundary_knot_indices_t = typename std::tuple_element_t<index, inputs_boundary_knot_indices_type>;
  
  using outputs_interior_knot_indices_type =
    std::tuple<decltype(std::declval<Outputs>()
                   .template find_knot_indices<functionspace::interior>(
                                                                        std::declval<typename Outputs::eval_type>()))...>;
 
  template<std::size_t index>
  using output_interior_knot_indices_t = typename std::tuple_element_t<index, outputs_interior_knot_indices_type>;
  
  using outputs_boundary_knot_indices_type =
    std::tuple<decltype(std::declval<Outputs>()
                        .template find_knot_indices<functionspace::boundary>(
                                                                             std::declval<
                                                                             typename Outputs::boundary_eval_type>()))...>;
 
  template<std::size_t index>
  using output_boundary_knot_indices_t = typename std::tuple_element_t<index, outputs_boundary_knot_indices_type>;
  
  using inputs_interior_coeff_indices_type =
    std::tuple<decltype(std::declval<Inputs>()
                        .template find_coeff_indices<functionspace::interior>(
                                                                              std::declval<typename Inputs::eval_type>()))...>;
 
  template<std::size_t index>
  using input_interior_coeff_indices_t = typename std::tuple_element_t<index, inputs_interior_coeff_indices_type>;
  
  using inputs_boundary_coeff_indices_type =
    std::tuple<decltype(std::declval<Inputs>()
                        .template find_coeff_indices<functionspace::boundary>(
                                                                              std::declval<
                                                                              typename Inputs::boundary_eval_type>()))...>;
 
  template<std::size_t index>
  using input_boundary_coeff_indices_t = typename std::tuple_element_t<index, inputs_boundary_coeff_indices_type>;
  
  using outputs_interior_coeff_indices_type =
    std::tuple<decltype(std::declval<Outputs>()
                        .template find_coeff_indices<functionspace::interior>(
                                                                              std::declval<typename Outputs::eval_type>()))...>;
 
  template<std::size_t index>
  using output_interior_coeff_indices_t = typename std::tuple_element_t<index, outputs_interior_coeff_indices_type>;
  
  using outputs_boundary_coeff_indices_type =
    std::tuple<decltype(std::declval<Outputs>()
                        .template find_coeff_indices<functionspace::boundary>(
                                                                              std::declval<
                                                                              typename Outputs::boundary_eval_type>()))...>;
 
  template<std::size_t index>
  using output_boundary_coeff_indices_t = typename std::tuple_element_t<index, outputs_boundary_coeff_indices_type>;
};
 
template <detail::HasAsTensor... Inputs,
          detail::HasAsTensor... Outputs,
          detail::HasAsTensor... CollPts>
class IgANetCustomizable2<std::tuple<Inputs...>,
                    std::tuple<Outputs...>,
                          std::tuple<CollPts...>> : public IgANetCustomizable2<std::tuple<Inputs...>,
                                                                               std::tuple<Outputs...>, void> {
public:
  using collPts_interior_knot_indices_type =
    std::tuple<decltype(std::declval<CollPts>()
                   .template find_knot_indices<functionspace::interior>(
                                                                        std::declval<typename CollPts::eval_type>()))...>;
 
  using collPts_boundary_knot_indices_type =
    std::tuple<decltype(std::declval<CollPts>()
                        .template find_knot_indices<functionspace::boundary>(
                                                                             std::declval<
                                                                             typename CollPts::boundary_eval_type>()))...>;
 
  using collPts_interior_coeff_indices_type =
    std::tuple<decltype(std::declval<CollPts>()
                        .template find_coeff_indices<functionspace::interior>(
                                                                              std::declval<typename CollPts::eval_type>()))...>;
 
  using collPts_boundary_coeff_indices_type =
    std::tuple<decltype(std::declval<CollPts>()
                        .template find_coeff_indices<functionspace::boundary>(
                                                                              std::declval<
                                                                              typename CollPts::boundary_eval_type>()))...>; 
};  
 
  
} // namespace iganet