yms-kan 0.0.1__py3-none-any.whl

yms_kan/MultKAN.py

@@ -0,0 +1,3085 @@
+import copy
+import math
+import os
+import random
+import sys
+
+import matplotlib.pyplot as plt
+import numpy as np
+# from .MultKANLayer import MultKANLayer
+import pandas as pd
+import sympy
+import torch
+import torch.nn as nn
+import yaml
+from PIL import Image
+from sklearn.metrics import accuracy_score, recall_score, precision_score, f1_score
+from sympy import *
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+from tqdm import tqdm
+
+from .KANLayer import KANLayer
+from .LBFGS import *
+# from .Symbolic_MultKANLayer import *
+from .Symbolic_KANLayer import Symbolic_KANLayer
+from .hypothesis import plot_tree
+from .spline import curve2coef
+from .utils import SYMBOLIC_LIB
+
+
+class MultKAN(nn.Module):
+    """
+    KAN class
+
+    Attributes:
+    -----------
+        grid : int
+            the number of grid intervals
+        k : int
+            spline order
+        act_fun : a list of KANLayers
+        symbolic_fun: a list of Symbolic_KANLayer
+        depth : int
+            depth of KAN
+        width : list
+            number of neurons in each layer.
+            Without multiplication nodes, [2,5,5,3] means 2D inputs, 3D outputs, with 2 layers of 5 hidden neurons.
+            With multiplication nodes, [2,[5,3],[5,1],3] means besides the [2,5,53] KAN, there are 3 (1) mul nodes in layer 1 (2).
+        mult_arity : int, or list of int lists
+            multiplication arity for each multiplication node (the number of numbers to be multiplied)
+        grid : int
+            the number of grid intervals
+        k : int
+            the order of piecewise polynomial
+        base_fun : fun
+            residual function b(x). an activation function phi(x) = sb_scale * b(x) + sp_scale * spline(x)
+        symbolic_fun : a list of Symbolic_KANLayer
+            Symbolic_KANLayers
+        symbolic_enabled : bool
+            If False, the symbolic front is not computed (to save time). Default: True.
+        width_in : list
+            The number of input neurons for each layer
+        width_out : list
+            The number of output neurons for each layer
+        base_fun_name : str
+            The base function b(x)
+        grip_eps : float
+            The parameter that interpolates between uniform grid and adaptive grid (based on sample quantile)
+        node_bias : a list of 1D torch.float
+        node_scale : a list of 1D torch.float
+        subnode_bias : a list of 1D torch.float
+        subnode_scale : a list of 1D torch.float
+        symbolic_enabled : bool
+            when symbolic_enabled = False, the symbolic branch (symbolic_fun) will be ignored in computation (set to zero)
+        affine_trainable : bool
+            indicate whether affine parameters are trainable (node_bias, node_scale, subnode_bias, subnode_scale)
+        sp_trainable : bool
+            indicate whether the overall magnitude of splines is trainable
+        sb_trainable : bool
+            indicate whether the overall magnitude of base function is trainable
+        save_act : bool
+            indicate whether intermediate activations are saved in forward pass
+        node_scores : None or list of 1D torch.float
+            node attribution score
+        edge_scores : None or list of 2D torch.float
+            edge attribution score
+        subnode_scores : None or list of 1D torch.float
+            subnode attribution score
+        cache_data : None or 2D torch.float
+            cached input data
+        acts : None or a list of 2D torch.float
+            activations on nodes
+        auto_save : bool
+            indicate whether to automatically save a checkpoint once the model is modified
+        state_id : int
+            the state of the model (used to save checkpoint)
+        ckpt_path : str
+            the folder to store checkpoints
+        round : int
+            the number of times rewind() has been called
+        device : str
+    """
+
+    def __init__(self, width=None, grid=3, k=3, mult_arity=2, noise_scale=0.3, scale_base_mu=0.0, scale_base_sigma=1.0,
+                 base_fun='silu', symbolic_enabled=True, affine_trainable=False, grid_eps=0.02, grid_range=[-1, 1],
+                 sp_trainable=True, sb_trainable=True, seed=1, save_act=True, sparse_init=False, auto_save=True,
+                 first_init=True, ckpt_path='./model', state_id=0, round=0, device='cpu'):
+        """
+        initalize a KAN model
+
+        Args:
+        -----
+            width : list of int
+                Without multiplication nodes: :math:`[n_0, n_1, .., n_{L-1}]` specify the number of neurons in each layer (including inputs/outputs)
+                With multiplication nodes: :math:`[[n_0,m_0=0], [n_1,m_1], .., [n_{L-1},m_{L-1}]]` specify the number of addition/multiplication nodes in each layer (including inputs/outputs)
+            grid : int
+                number of grid intervals. Default: 3.
+            k : int
+                order of piecewise polynomial. Default: 3.
+            mult_arity : int, or list of int lists
+                multiplication arity for each multiplication node (the number of numbers to be multiplied)
+            noise_scale : float
+                initial injected noise to spline.
+            base_fun : str
+                the residual function b(x). Default: 'silu'
+            symbolic_enabled : bool
+                compute (True) or skip (False) symbolic computations (for efficiency). By default: True.
+            affine_trainable : bool
+                affine parameters are updated or not. Affine parameters include node_scale, node_bias, subnode_scale, subnode_bias
+            grid_eps : float
+                When grid_eps = 1, the grid is uniform; when grid_eps = 0, the grid is partitioned using percentiles of samples. 0 < grid_eps < 1 interpolates between the two extremes.
+            grid_range : list/np.array of shape (2,))
+                setting the range of grids. Default: [-1,1]. This argument is not important if fit(update_grid=True) (by default updata_grid=True)
+            sp_trainable : bool
+                If true, scale_sp is trainable. Default: True.
+            sb_trainable : bool
+                If true, scale_base is trainable. Default: True.
+            device : str
+                device
+            seed : int
+                random seed
+            save_act : bool
+                indicate whether intermediate activations are saved in forward pass
+            sparse_init : bool
+                sparse initialization (True) or normal dense initialization. Default: False.
+            auto_save : bool
+                indicate whether to automatically save a checkpoint once the model is modified
+            state_id : int
+                the state of the model (used to save checkpoint)
+            ckpt_path : str
+                the folder to store checkpoints. Default: './model'
+            round : int
+                the number of times rewind() has been called
+            device : str
+
+        Returns:
+        --------
+            self
+
+        Example
+        -------
+        # >>> from yms_kan import *
+        # >>> model = KAN(width=[2,5,1], grid=5, k=3, seed=0)
+        checkpoint directory created: ./model
+        saving model version 0.0
+        """
+        super(MultKAN, self).__init__()
+
+        torch.manual_seed(seed)
+        np.random.seed(seed)
+        random.seed(seed)
+
+        ### initializeing the numerical front ###
+
+        self.act_fun = []
+        self.depth = len(width) - 1
+
+        #print('haha1', width)
+        for i in range(len(width)):
+            #print(type(width[i]), type(width[i]) == int)
+            if type(width[i]) == int or type(width[i]) == np.int64:
+                width[i] = [width[i], 0]
+
+        #print('haha2', width)
+
+        self.width = width
+
+        # if mult_arity is just a scalar, we extend it to a list of lists
+        # e.g, mult_arity = [[2,3],[4]] means that in the first hidden layer, 2 mult ops have arity 2 and 3, respectively;
+        # in the second hidden layer, 1 mult op has arity 4.
+        if isinstance(mult_arity, int):
+            self.mult_homo = True  # when homo is True, parallelization is possible
+        else:
+            self.mult_homo = False  # when home if False, for loop is required.
+        self.mult_arity = mult_arity
+
+        width_in = self.width_in
+        width_out = self.width_out
+
+        self.base_fun_name = base_fun
+        if base_fun == 'silu':
+            base_fun = torch.nn.SiLU()
+        elif base_fun == 'identity':
+            base_fun = torch.nn.Identity()
+        elif base_fun == 'zero':
+            base_fun = lambda x: x * 0.
+
+        self.grid_eps = grid_eps
+        self.grid_range = grid_range
+
+        for l in range(self.depth):
+            # splines
+            if isinstance(grid, list):
+                grid_l = grid[l]
+            else:
+                grid_l = grid
+
+            if isinstance(k, list):
+                k_l = k[l]
+            else:
+                k_l = k
+
+            sp_batch = KANLayer(in_dim=width_in[l], out_dim=width_out[l + 1], num=grid_l, k=k_l,
+                                noise_scale=noise_scale, scale_base_mu=scale_base_mu, scale_base_sigma=scale_base_sigma,
+                                scale_sp=1., base_fun=base_fun, grid_eps=grid_eps, grid_range=grid_range,
+                                sp_trainable=sp_trainable, sb_trainable=sb_trainable, sparse_init=sparse_init)
+            self.act_fun.append(sp_batch)
+
+        self.node_bias = []
+        self.node_scale = []
+        self.subnode_bias = []
+        self.subnode_scale = []
+
+        globals()['self.node_bias_0'] = torch.nn.Parameter(torch.zeros(3, 1)).requires_grad_(False)
+        exec('self.node_bias_0' + " = torch.nn.Parameter(torch.zeros(3,1)).requires_grad_(False)")
+
+        for l in range(self.depth):
+            exec(
+                f'self.node_bias_{l} = torch.nn.Parameter(torch.zeros(width_in[l+1])).requires_grad_(affine_trainable)')
+            exec(
+                f'self.node_scale_{l} = torch.nn.Parameter(torch.ones(width_in[l+1])).requires_grad_(affine_trainable)')
+            exec(
+                f'self.subnode_bias_{l} = torch.nn.Parameter(torch.zeros(width_out[l+1])).requires_grad_(affine_trainable)')
+            exec(
+                f'self.subnode_scale_{l} = torch.nn.Parameter(torch.ones(width_out[l+1])).requires_grad_(affine_trainable)')
+            exec(f'self.node_bias.append(self.node_bias_{l})')
+            exec(f'self.node_scale.append(self.node_scale_{l})')
+            exec(f'self.subnode_bias.append(self.subnode_bias_{l})')
+            exec(f'self.subnode_scale.append(self.subnode_scale_{l})')
+
+        self.act_fun = nn.ModuleList(self.act_fun)
+
+        self.grid = grid
+        self.k = k
+        self.base_fun = base_fun
+
+        ### initializing the symbolic front ###
+        self.symbolic_fun = []
+        for l in range(self.depth):
+            sb_batch = Symbolic_KANLayer(in_dim=width_in[l], out_dim=width_out[l + 1])
+            self.symbolic_fun.append(sb_batch)
+
+        self.symbolic_fun = nn.ModuleList(self.symbolic_fun)
+        self.symbolic_enabled = symbolic_enabled
+        self.affine_trainable = affine_trainable
+        self.sp_trainable = sp_trainable
+        self.sb_trainable = sb_trainable
+
+        self.save_act = save_act
+
+        self.node_scores = None
+        self.edge_scores = None
+        self.subnode_scores = None
+
+        self.cache_data = None
+        self.acts = None
+
+        self.auto_save = auto_save
+        self.state_id = 0
+        self.ckpt_path = ckpt_path
+        self.round = round
+
+        self.device = device
+        self.to(device)
+
+        if auto_save:
+            if first_init:
+                if not os.path.exists(ckpt_path):
+                    # Create the directory
+                    os.makedirs(ckpt_path)
+                print(f"checkpoint directory created: {ckpt_path}")
+                print('saving model version 0.0')
+
+                history_path = self.ckpt_path + '/history.txt'
+                with open(history_path, 'w') as file:
+                    file.write(f'### Round {self.round} ###' + '\n')
+                    file.write('init => 0.0' + '\n')
+                self.saveckpt(path=self.ckpt_path + '/' + '0.0')
+            else:
+                self.state_id = state_id
+
+        self.input_id = torch.arange(self.width_in[0], )
+
+    def to(self, device):
+        """
+        move the model to device
+
+        Args:
+        -----
+            device : str or device
+
+        Returns:
+        --------
+            self
+
+        Example
+        -------
+        # >>> from yms_kan import *
+        # >>> device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        # >>> model = KAN(width=[2,5,1], grid=5, k=3, seed=0)
+        # >>> model.to(device)
+        """
+        super(MultKAN, self).to(device)
+        self.device = device
+
+        for kanlayer in self.act_fun:
+            kanlayer.to(device)
+
+        for symbolic_kanlayer in self.symbolic_fun:
+            symbolic_kanlayer.to(device)
+
+        return self
+
+    @property
+    def width_in(self):
+        """
+        The number of input nodes for each layer
+        """
+        width = self.width
+        width_in = [width[l][0] + width[l][1] for l in range(len(width))]
+        return width_in
+
+    @property
+    def width_out(self):
+        """
+        The number of output subnodes for each layer
+        """
+        width = self.width
+        if self.mult_homo == True:
+            width_out = [width[l][0] + self.mult_arity * width[l][1] for l in range(len(width))]
+        else:
+            width_out = [width[l][0] + int(np.sum(self.mult_arity[l])) for l in range(len(width))]
+        return width_out
+
+    @property
+    def n_sum(self):
+        """
+        The number of addition nodes for each layer
+        """
+        width = self.width
+        n_sum = [width[l][0] for l in range(1, len(width) - 1)]
+        return n_sum
+
+    @property
+    def n_mult(self):
+        """
+        The number of multiplication nodes for each layer
+        """
+        width = self.width
+        n_mult = [width[l][1] for l in range(1, len(width) - 1)]
+        return n_mult
+
+    @property
+    def feature_score(self):
+        """
+        attribution scores for inputs
+        """
+        self.attribute()
+        if self.node_scores is None:
+            return None
+        else:
+            return self.node_scores[0]
+
+    def initialize_from_another_model(self, another_model, x):
+        """
+        initialize from another model of the same width, but their 'grid' parameter can be different.
+        Note this is equivalent to refine() when we don't want to keep another_model
+
+        Args:
+        -----
+            another_model : MultKAN
+            x : 2D torch.float
+
+        Returns:
+        --------
+            self
+
+        Example
+        -------
+        # >>> from yms_kan import *
+        # >>> model1 = KAN(width=[2,5,1], grid=3)
+        # >>> model2 = KAN(width=[2,5,1], grid=10)
+        # >>> x = torch.rand(100,2)
+        # >>> model2.initialize_from_another_model(model1, x)
+        """
+        another_model(x)  # get activations
+        batch = x.shape[0]
+
+        self.initialize_grid_from_another_model(another_model, x)
+
+        for l in range(self.depth):
+            spb = self.act_fun[l]
+            #spb_parent = another_model.act_fun[l]
+
+            # spb = spb_parent
+            preacts = another_model.spline_preacts[l]
+            postsplines = another_model.spline_postsplines[l]
+            self.act_fun[l].coef.data = curve2coef(preacts[:, 0, :], postsplines.permute(0, 2, 1), spb.grid, k=spb.k)
+            self.act_fun[l].scale_base.data = another_model.act_fun[l].scale_base.data
+            self.act_fun[l].scale_sp.data = another_model.act_fun[l].scale_sp.data
+            self.act_fun[l].mask.data = another_model.act_fun[l].mask.data
+
+        for l in range(self.depth):
+            self.node_bias[l].data = another_model.node_bias[l].data
+            self.node_scale[l].data = another_model.node_scale[l].data
+
+            self.subnode_bias[l].data = another_model.subnode_bias[l].data
+            self.subnode_scale[l].data = another_model.subnode_scale[l].data
+
+        for l in range(self.depth):
+            self.symbolic_fun[l] = another_model.symbolic_fun[l]
+
+        return self.to(self.device)
+
+    def log_history(self, method_name):
+
+        if self.auto_save:
+            # save to log file
+            # print(func.__name__)
+            with open(self.ckpt_path + '/history.txt', 'a') as file:
+                file.write(str(self.round) + '.' + str(self.state_id) + ' => ' + method_name + ' => ' + str(
+                    self.round) + '.' + str(self.state_id + 1) + '\n')
+
+            # update state_id
+            self.state_id += 1
+
+            # save to ckpt
+            self.saveckpt(path=self.ckpt_path + '/' + str(self.round) + '.' + str(self.state_id))
+            print('saving model version ' + str(self.round) + '.' + str(self.state_id))
+
+    def refine(self, new_grid):
+        """
+        grid refinement
+
+        Args:
+        -----
+            new_grid : init
+                the number of grid intervals after refinement
+
+        Returns:
+        --------
+            a refined model : MultKAN
+
+        Example
+        -------
+        # >>> from yms_kan import *
+        # >>> device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        # >>> model = KAN(width=[2,5,1], grid=5, k=3, seed=0)
+        # >>> print(model.grid)
+        # >>> x = torch.rand(100,2)
+        # >>> model.get_act(x)
+        # >>> model = model.refine(10)
+        # >>> print(model.grid)
+        checkpoint directory created: ./model
+        saving model version 0.0
+        5
+        saving model version 0.1
+        10
+        """
+
+        model_new = MultKAN(width=self.width,
+                            grid=new_grid,
+                            k=self.k,
+                            mult_arity=self.mult_arity,
+                            base_fun=self.base_fun_name,
+                            symbolic_enabled=self.symbolic_enabled,
+                            affine_trainable=self.affine_trainable,
+                            grid_eps=self.grid_eps,
+                            grid_range=self.grid_range,
+                            sp_trainable=self.sp_trainable,
+                            sb_trainable=self.sb_trainable,
+                            ckpt_path=self.ckpt_path,
+                            auto_save=True,
+                            first_init=False,
+                            state_id=self.state_id,
+                            round=self.round,
+                            device=self.device)
+
+        model_new.initialize_from_another_model(self, self.cache_data)
+        model_new.cache_data = self.cache_data
+        model_new.grid = new_grid
+
+        self.log_history('refine')
+        model_new.state_id += 1
+
+        return model_new.to(self.device)
+
+    def saveckpt(self, path='model'):
+        """
+        save the current model to files (configuration file and state file)
+
+        Args:
+        -----
+            path : str
+                the path where checkpoints are saved
+
+        Returns:
+        --------
+            None
+
+        Example
+        -------
+        # >>> from yms_kan import *
+        # >>> device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        # >>> model = KAN(width=[2,5,1], grid=5, k=3, seed=0)
+        # >>> model.saveckpt('./mark')
+        # There will be three files appearing in the current folder: mark_cache_data, mark_config.yml, mark_state
+        """
+
+        model = self
+
+        dic = dict(
+            width=model.width,
+            grid=model.grid,
+            k=model.k,
+            mult_arity=model.mult_arity,
+            base_fun_name=model.base_fun_name,
+            symbolic_enabled=model.symbolic_enabled,
+            affine_trainable=model.affine_trainable,
+            grid_eps=model.grid_eps,
+            grid_range=model.grid_range,
+            sp_trainable=model.sp_trainable,
+            sb_trainable=model.sb_trainable,
+            state_id=model.state_id,
+            auto_save=model.auto_save,
+            ckpt_path=model.ckpt_path,
+            round=model.round,
+            device=str(model.device)
+        )
+
+        if dic["device"].isdigit():
+            dic["device"] = int(model.device)
+
+        for i in range(model.depth):
+            dic[f'symbolic.funs_name.{i}'] = model.symbolic_fun[i].funs_name
+
+        with open(f'{path}_config.yml', 'w') as outfile:
+            yaml.dump(dic, outfile, default_flow_style=False)
+
+        torch.save(model.state_dict(), f'{path}_state')
+        torch.save(model.cache_data, f'{path}_cache_data')
+
+    @staticmethod
+    def loadckpt(path='model'):
+        """
+        load checkpoint from path
+
+        Args:
+        -----
+            path : str
+                the path where checkpoints are saved
+
+        Returns:
+        --------
+            MultKAN
+
+        Example
+        -------
+        # >>> from yms_kan import *
+        # >>> model = KAN(width=[2,5,1], grid=5, k=3, seed=0)
+        # >>> model.saveckpt('./mark')
+        # >>> KAN.loadckpt('./mark')
+        """
+        with open(f'{path}_config.yml', 'r') as stream:
+            config = yaml.safe_load(stream)
+
+        state = torch.load(f'{path}_state', map_location=torch.device('cpu'), weights_only=False)
+
+        model_load = MultKAN(width=config['width'],
+                             grid=config['grid'],
+                             k=config['k'],
+                             mult_arity=config['mult_arity'],
+                             base_fun=config['base_fun_name'],
+                             symbolic_enabled=config['symbolic_enabled'],
+                             affine_trainable=config['affine_trainable'],
+                             grid_eps=config['grid_eps'],
+                             grid_range=config['grid_range'],
+                             sp_trainable=config['sp_trainable'],
+                             sb_trainable=config['sb_trainable'],
+                             state_id=config['state_id'],
+                             auto_save=config['auto_save'],
+                             first_init=False,
+                             ckpt_path=config['ckpt_path'],
+                             round=config['round'] + 1,
+                             device='cpu')
+
+        model_load.load_state_dict(state)
+        model_load.cache_data = torch.load(f'{path}_cache_data', map_location=torch.device('cpu'),
+                                           weights_only=False)
+
+        depth = len(model_load.width) - 1
+        for l in range(depth):
+            out_dim = model_load.symbolic_fun[l].out_dim
+            in_dim = model_load.symbolic_fun[l].in_dim
+            funs_name = config[f'symbolic.funs_name.{l}']
+            for j in range(out_dim):
+                for i in range(in_dim):
+                    fun_name = funs_name[j][i]
+                    model_load.symbolic_fun[l].funs_name[j][i] = fun_name
+                    model_load.symbolic_fun[l].funs[j][i] = SYMBOLIC_LIB[fun_name][0]
+                    model_load.symbolic_fun[l].funs_sympy[j][i] = SYMBOLIC_LIB[fun_name][1]
+                    model_load.symbolic_fun[l].funs_avoid_singularity[j][i] = SYMBOLIC_LIB[fun_name][3]
+        model_load.device = config['device']
+        return model_load
+
+    def copy(self):
+        """
+        deepcopy
+
+        Args:
+        -----
+            path : str
+                the path where checkpoints are saved
+
+        Returns:
+        --------
+            MultKAN
+
+        Example
+        -------
+        # >>> from yms_kan import *
+        # >>> model = KAN(width=[1,1], grid=5, k=3, seed=0)
+        # >>> model2 = model.copy()
+        # >>> model2.act_fun[0].coef.data *= 2
+        # >>> print(model2.act_fun[0].coef.data)
+        # >>> print(model.act_fun[0].coef.data)
+        """
+        path = 'copy_temp'
+        self.saveckpt(path)
+        return KAN.loadckpt(path)
+
+    def rewind(self, model_id):
+        """
+        rewind to an old version
+
+        Args:
+        -----
+            model_id : str
+                in format '{a}.{b}' where a is the round number, b is the version number in that round
+
+        Returns:
+        --------
+            MultKAN
+
+        Example
+        -------
+        Please refer to tutorials. API 12: Checkpoint, save & load model
+        """
+        self.round += 1
+        self.state_id = model_id.split('.')[-1]
+
+        history_path = self.ckpt_path + '/history.txt'
+        with open(history_path, 'a') as file:
+            file.write(f'### Round {self.round} ###' + '\n')
+
+        self.saveckpt(path=self.ckpt_path + '/' + f'{self.round}.{self.state_id}')
+
+        print(
+            'rewind to model version ' + f'{self.round - 1}.{self.state_id}' + ', renamed as ' + f'{self.round}.{self.state_id}')
+
+        return MultKAN.loadckpt(path=self.ckpt_path + '/' + str(model_id))
+
+    def checkout(self, model_id):
+        """
+        check out an old version
+
+        Args:
+        -----
+            model_id : str
+                in format '{a}.{b}' where a is the round number, b is the version number in that round
+
+        Returns:
+        --------
+            MultKAN
+
+        Example
+        -------
+        Same use as rewind, although checkout doesn't change states
+        """
+        return MultKAN.loadckpt(path=self.ckpt_path + '/' + str(model_id))
+
+    def update_grid_from_samples(self, x):
+        """
+        update grid from samples
+
+        Args:
+        -----
+            x : 2D torch.tensor
+                inputs
+
+        Returns:
+        --------
+            None
+
+        Example
+        -------
+        # >>> from yms_kan import *
+        # >>> model = KAN(width=[1,1], grid=5, k=3, seed=0)
+        # >>> print(model.act_fun[0].grid)
+        # >>> x = torch.linspace(-10,10,steps=101)[:,None]
+        # >>> model.update_grid_from_samples(x)
+        # >>> print(model.act_fun[0].grid)
+        """
+        for l in range(self.depth):
+            self.get_act(x)
+            self.act_fun[l].update_grid_from_samples(self.acts[l])
+
+    def update_grid(self, x):
+        """
+        call update_grid_from_samples. This seems unnecessary but we retain it for the sake of classes that might inherit from MultKAN
+        """
+        self.update_grid_from_samples(x)
+
+    def initialize_grid_from_another_model(self, model, x):
+        """
+        initialize grid from another model
+
+        Args:
+        -----
+            model : MultKAN
+                parent model
+            x : 2D torch.tensor
+                inputs
+
+        Returns:
+        --------
+            None
+
+        Example
+        -------
+        # >>> from yms_kan import *
+        # >>> model = KAN(width=[1,1], grid=5, k=3, seed=0)
+        # >>> print(model.act_fun[0].grid)
+        # >>> x = torch.linspace(-10,10,steps=101)[:,None]
+        # >>> model2 = KAN(width=[1,1], grid=10, k=3, seed=0)
+        # >>> model2.initialize_grid_from_another_model(model, x)
+        # >>> print(model2.act_fun[0].grid)
+        """
+        model(x)
+        for l in range(self.depth):
+            self.act_fun[l].initialize_grid_from_parent(model.act_fun[l], model.acts[l])
+
+    def forward(self, x, singularity_avoiding=False, y_th=10.):
+        """
+        forward pass
+
+        Args:
+        -----
+            x : 2D torch.tensor
+                inputs
+            singularity_avoiding : bool
+                whether to avoid singularity for the symbolic branch
+            y_th : float
+                the threshold for singularity
+
+        Returns:
+        --------
+            None
+
+        Example1
+        --------
+        # >>> from yms_kan import *
+        # >>> model = KAN(width=[2,5,1], grid=5, k=3, seed=0)
+        # >>> x = torch.rand(100,2)
+        # >>> model(x).shape
+
+        Example2
+        --------
+        # >>> from yms_kan import *
+        # >>> model = KAN(width=[1,1], grid=5, k=3, seed=0)
+        # >>> x = torch.tensor([[1],[-0.01]])
+        # >>> model.fix_symbolic(0,0,0,'log',fit_params_bool=False)
+        # >>> print(model(x))
+        # >>> print(model(x, singularity_avoiding=True))
+        # >>> print(model(x, singularity_avoiding=True, y_th=1.))
+        """
+        x = x[:, self.input_id.long()]
+        assert x.shape[1] == self.width_in[0]
+
+        # cache data
+        self.cache_data = x
+
+        self.acts = []  # shape ([batch, n0], [batch, n1], ..., [batch, n_L])
+        self.acts_premult = []
+        self.spline_preacts = []
+        self.spline_postsplines = []
+        self.spline_postacts = []
+        self.acts_scale = []
+        self.acts_scale_spline = []
+        self.subnode_actscale = []
+        self.edge_actscale = []
+        # self.neurons_scale = []
+
+        self.acts.append(x)  # acts shape: (batch, width[l])
+
+        for l in range(self.depth):
+
+            x_numerical, preacts, postacts_numerical, postspline = self.act_fun[l](x)
+            #print(preacts, postacts_numerical, postspline)
+
+            if self.symbolic_enabled:
+                x_symbolic, postacts_symbolic = self.symbolic_fun[l](x, singularity_avoiding=singularity_avoiding,
+                                                                     y_th=y_th)
+            else:
+                x_symbolic = 0.
+                postacts_symbolic = 0.
+
+            x = x_numerical + x_symbolic
+
+            if self.save_act:
+                # save subnode_scale
+                self.subnode_actscale.append(torch.std(x, dim=0).detach())
+
+            # subnode affine transform
+            x = self.subnode_scale[l][None, :] * x + self.subnode_bias[l][None, :]
+
+            if self.save_act:
+                postacts = postacts_numerical + postacts_symbolic
+
+                # self.neurons_scale.append(torch.mean(torch.abs(x), dim=0))
+                #grid_reshape = self.act_fun[l].grid.reshape(self.width_out[l + 1], self.width_in[l], -1)
+                input_range = torch.std(preacts, dim=0) + 0.1
+                output_range_spline = torch.std(postacts_numerical,
+                                                dim=0)  # for training, only penalize the spline part
+                output_range = torch.std(postacts,
+                                         dim=0)  # for visualization, include the contribution from both spline + symbolic
+                # save edge_scale
+                self.edge_actscale.append(output_range)
+
+                self.acts_scale.append((output_range / input_range).detach())
+                self.acts_scale_spline.append(output_range_spline / input_range)
+                self.spline_preacts.append(preacts.detach())
+                self.spline_postacts.append(postacts.detach())
+                self.spline_postsplines.append(postspline.detach())
+
+                self.acts_premult.append(x.detach())
+
+            # multiplication
+            dim_sum = self.width[l + 1][0]
+            dim_mult = self.width[l + 1][1]
+
+            if self.mult_homo:
+                for i in range(self.mult_arity - 1):
+                    if i == 0:
+                        x_mult = x[:, dim_sum::self.mult_arity] * x[:, dim_sum + 1::self.mult_arity]
+                    else:
+                        x_mult = x_mult * x[:, dim_sum + i + 1::self.mult_arity]
+
+            else:
+                for j in range(dim_mult):
+                    acml_id = dim_sum + np.sum(self.mult_arity[l + 1][:j])
+                    for i in range(self.mult_arity[l + 1][j] - 1):
+                        if i == 0:
+                            x_mult_j = x[:, [acml_id]] * x[:, [acml_id + 1]]
+                        else:
+                            x_mult_j = x_mult_j * x[:, [acml_id + i + 1]]
+
+                    if j == 0:
+                        x_mult = x_mult_j
+                    else:
+                        x_mult = torch.cat([x_mult, x_mult_j], dim=1)
+
+            if self.width[l + 1][1] > 0:
+                x = torch.cat([x[:, :dim_sum], x_mult], dim=1)
+
+            # x = x + self.biases[l].weight
+            # node affine transform
+            x = self.node_scale[l][None, :] * x + self.node_bias[l][None, :]
+
+            self.acts.append(x.detach())
+
+        return x
+
+    def set_mode(self, l, i, j, mode, mask_n=None):
+        if mode == "s":
+            mask_n = 0.
+            mask_s = 1.
+        elif mode == "n":
+            mask_n = 1.
+            mask_s = 0.
+        elif mode == "sn" or mode == "ns":
+            if mask_n is None:
+                mask_n = 1.
+            else:
+                mask_n = mask_n
+            mask_s = 1.
+        else:
+            mask_n = 0.
+            mask_s = 0.
+
+        self.act_fun[l].mask.data[i][j] = mask_n
+        self.symbolic_fun[l].mask.data[j, i] = mask_s
+
+    def fix_symbolic(self, l, i, j, fun_name, fit_params_bool=True, a_range=(-10, 10), b_range=(-10, 10), verbose=True,
+                     random=False, log_history=True):
+        '''
+        set (l,i,j) activation to be symbolic (specified by fun_name)
+        
+        Args:
+        -----
+            l : int
+                layer index
+            i : int
+                input neuron index
+            j : int
+                output neuron index
+            fun_name : str
+                function name
+            fit_params_bool : bool
+                obtaining affine parameters through fitting (True) or setting default values (False)
+            a_range : tuple
+                sweeping range of a
+            b_range : tuple
+                sweeping range of b
+            verbose : bool
+                If True, more information is printed.
+            random : bool
+                initialize affine parameteres randomly or as [1,0,1,0]
+            log_history : bool
+                indicate whether to log history when the function is called
+        
+        Returns:
+        --------
+            None or r2 (coefficient of determination)
+            
+        Example 1 
+        ---------
+        >>> # when fit_params_bool = False
+        >>> model = KAN(width=[2,5,1], grid=5, k=3)
+        >>> model.fix_symbolic(0,1,3,'sin',fit_params_bool=False)
+        >>> print(model.act_fun[0].mask.reshape(2,5))
+        >>> print(model.symbolic_fun[0].mask.reshape(2,5))
+                    
+        Example 2
+        ---------
+        >>> # when fit_params_bool = True
+        >>> model = KAN(width=[2,5,1], grid=5, k=3, noise_scale=1.)
+        >>> x = torch.normal(0,1,size=(100,2))
+        >>> model(x) # obtain activations (otherwise model does not have attributes acts)
+        >>> model.fix_symbolic(0,1,3,'sin',fit_params_bool=True)
+        >>> print(model.act_fun[0].mask.reshape(2,5))
+        >>> print(model.symbolic_fun[0].mask.reshape(2,5))
+        '''
+        if not fit_params_bool:
+            self.symbolic_fun[l].fix_symbolic(i, j, fun_name, verbose=verbose, random=random)
+            r2 = None
+        else:
+            x = self.acts[l][:, i]
+            mask = self.act_fun[l].mask
+            y = self.spline_postacts[l][:, j, i]
+            # y = self.postacts[l][:, j, i]
+            r2 = self.symbolic_fun[l].fix_symbolic(i, j, fun_name, x, y, a_range=a_range, b_range=b_range,
+                                                   verbose=verbose)
+            if mask[i, j] == 0:
+                r2 = - 1e8
+        self.set_mode(l, i, j, mode="s")
+
+        if log_history:
+            self.log_history('fix_symbolic')
+        return r2
+
+    def unfix_symbolic(self, l, i, j, log_history=True):
+        """
+        unfix the (l,i,j) activation function.
+        """
+        self.set_mode(l, i, j, mode="n")
+        self.symbolic_fun[l].funs_name[j][i] = "0"
+        if log_history:
+            self.log_history('unfix_symbolic')
+
+    def unfix_symbolic_all(self, log_history=True):
+        '''
+        unfix all activation functions.
+        '''
+        for l in range(len(self.width) - 1):
+            for i in range(self.width_in[l]):
+                for j in range(self.width_out[l + 1]):
+                    self.unfix_symbolic(l, i, j, log_history)
+
+    def get_range(self, l, i, j, verbose=True):
+        '''
+        Get the input range and output range of the (l,i,j) activation
+        
+        Args:
+        -----
+            l : int
+                layer index
+            i : int
+                input neuron index
+            j : int
+                output neuron index
+        
+        Returns:
+        --------
+            x_min : float
+                minimum of input
+            x_max : float
+                maximum of input
+            y_min : float
+                minimum of output
+            y_max : float
+                maximum of output
+        
+        Example
+        -------
+        >>> model = KAN(width=[2,3,1], grid=5, k=3, noise_scale=1.)
+        >>> x = torch.normal(0,1,size=(100,2))
+        >>> model(x) # do a forward pass to obtain model.acts
+        >>> model.get_range(0,0,0)
+        '''
+        x = self.spline_preacts[l][:, j, i]
+        y = self.spline_postacts[l][:, j, i]
+        x_min = torch.min(x).cpu().detach().numpy()
+        x_max = torch.max(x).cpu().detach().numpy()
+        y_min = torch.min(y).cpu().detach().numpy()
+        y_max = torch.max(y).cpu().detach().numpy()
+        if verbose:
+            print('x range: [' + '%.2f' % x_min, ',', '%.2f' % x_max, ']')
+            print('y range: [' + '%.2f' % y_min, ',', '%.2f' % y_max, ']')
+        return x_min, x_max, y_min, y_max
+
+    def _draw_operator_symbol(self, ax, fig, x, y, symbol_type, scale, size):
+        """辅助函数绘制运算符符号"""
+        path = os.path.join(os.path.dirname(__file__), 'assets', 'img', f'{symbol_type}_symbol.png')
+        try:
+            with Image.open(path) as img:
+                trans = ax.transData.transform
+                inv_trans = fig.transFigure.inverted().transform
+                bbox = [x - size / 2, y - size / 2, size, size]
+                bbox = inv_trans.transform_bbox(trans(bbox))
+                ax.imshow(img, extent=bbox)
+        except Exception as e:
+            print(f"无法加载符号 {symbol_type}: {str(e)}")
+
+    def _draw_variable_labels(self, ax, variables, layer, scale, varscale, y_offset, y0, z0):
+        """辅助函数绘制变量标签"""
+        n = self.width_in[layer]
+        for i, var in enumerate(variables):
+            x = 1 / (2 * n) + i / n
+            y = layer * (y0 + z0) + y_offset
+            label = f'${latex(var)}$' if isinstance(var, sympy.Expr) else var
+            ax.text(x, y, label, fontsize=15 * scale * varscale,
+                    ha='center', va='center', transform=ax.transData)
+
+    def plot(self, folder="./figures", beta=3, metric='backward', scale=0.5, tick=False, sample=False, in_vars=None,
+             out_vars=None, title=None, varscale=1.0):
+        '''
+        plot KAN
+        
+        Args:
+        -----
+            folder : str
+                the folder to store pngs
+            beta : float
+                positive number. control the transparency of each activation. transparency = tanh(beta*l1).
+            mask : bool
+                If True, plot with mask (need to run prune() first to obtain mask). If False (by default), plot all activation functions.
+            mode : bool
+                "supervised" or "unsupervised". If "supervised", l1 is measured by absolution value (not subtracting mean); if "unsupervised", l1 is measured by standard deviation (subtracting mean).
+            scale : float
+                control the size of the diagram
+            in_vars: None or list of str
+                the name(s) of input variables
+            out_vars: None or list of str
+                the name(s) of output variables
+            title: None or str
+                title
+            varscale : float
+                the size of input variables
+            
+        Returns:
+        --------
+            Figure
+            
+        Example
+        -------
+        >>> # see more interactive examples in demos
+        >>> model = KAN(width=[2,3,1], grid=3, k=3, noise_scale=1.0)
+        >>> x = torch.normal(0,1,size=(100,2))
+        >>> model(x) # do a forward pass to obtain model.acts
+        >>> model.plot()
+        '''
+        global Symbol
+
+        if not self.save_act:
+            print('cannot plot since data are not saved. Set save_act=True first.')
+
+        if self.acts is None:
+            if self.cache_data is None:
+                raise Exception('model hasn\'t seen any data yet.')
+            self.forward(self.cache_data)
+
+        if metric == 'backward':
+            self.attribute()
+
+        if not os.path.exists(folder):
+            os.makedirs(folder)
+
+        depth = len(self.width) - 1
+        for l in range(depth):
+            w_large = 2.0
+            for i in range(self.width_in[l]):
+                for j in range(self.width_out[l + 1]):
+                    rank = torch.argsort(self.acts[l][:, i])
+                    fig, ax = plt.subplots(figsize=(w_large, w_large))
+
+                    num = rank.shape[0]
+                    symbolic_mask = self.symbolic_fun[l].mask[j][i]
+                    numeric_mask = self.act_fun[l].mask[i][j]
+                    if symbolic_mask > 0. and numeric_mask > 0.:
+                        color = 'purple'
+                        alpha_mask = 1
+                    if symbolic_mask > 0. and numeric_mask == 0.:
+                        color = "red"
+                        alpha_mask = 1
+                    if symbolic_mask == 0. and numeric_mask > 0.:
+                        color = "black"
+                        alpha_mask = 1
+                    if symbolic_mask == 0. and numeric_mask == 0.:
+                        color = "white"
+                        alpha_mask = 0
+
+                    if tick:
+                        ax.tick_params(axis="y", direction="in", pad=-22, labelsize=50)
+                        ax.tick_params(axis="x", direction="in", pad=-15, labelsize=50)
+                        x_min, x_max, y_min, y_max = self.get_range(l, i, j, verbose=False)
+                        plt.xticks([x_min, x_max], ['%2.f' % x_min, '%2.f' % x_max])
+                        plt.yticks([y_min, y_max], ['%2.f' % y_min, '%2.f' % y_max])
+                    else:
+                        plt.xticks([])
+                        plt.yticks([])
+                    if alpha_mask == 1:
+                        plt.gca().patch.set_edgecolor('black')
+                    else:
+                        plt.gca().patch.set_edgecolor('white')
+                    plt.gca().patch.set_linewidth(1.5)
+                    # plt.axis('off')
+
+                    plt.plot(self.acts[l][:, i][rank].cpu().detach().numpy(),
+                             self.spline_postacts[l][:, j, i][rank].cpu().detach().numpy(), color=color, lw=5)
+                    if sample:
+                        plt.scatter(self.acts[l][:, i][rank].cpu().detach().numpy(),
+                                    self.spline_postacts[l][:, j, i][rank].cpu().detach().numpy(), color=color,
+                                    s=400 * scale ** 2)
+                    plt.gca().spines[:].set_color(color)
+
+                    plt.savefig(f'{folder}/sp_{l}_{i}_{j}.png', bbox_inches="tight", dpi=400)
+                    plt.close()
+
+        def score2alpha(score):
+            return np.tanh(beta * score)
+
+        if metric == 'forward_n':
+            scores = self.acts_scale
+        elif metric == 'forward_u':
+            scores = self.edge_actscale
+        elif metric == 'backward':
+            scores = self.edge_scores
+        else:
+            raise Exception(f'metric = \'{metric}\' not recognized')
+
+        alpha = [score2alpha(score.cpu().detach().numpy()) for score in scores]
+
+        # draw skeleton
+        width = np.array(self.width)
+        width_in = np.array(self.width_in)
+        width_out = np.array(self.width_out)
+        A = 1
+        y0 = 0.3  # height: from input to pre-mult
+        z0 = 0.1  # height: from pre-mult to post-mult (input of next layer)
+
+        neuron_depth = len(width)
+        min_spacing = A / np.maximum(np.max(width_out), 5)
+
+        max_neuron = np.max(width_out)
+        max_num_weights = np.max(width_in[:-1] * width_out[1:])
+        y1 = 0.4 / np.maximum(max_num_weights, 5)  # size (height/width) of 1D function diagrams
+        y2 = 0.15 / np.maximum(max_neuron, 5)  # size (height/width) of operations (sum and mult)
+
+        fig, ax = plt.subplots(figsize=(10 * scale, 10 * scale * (neuron_depth - 1) * (y0 + z0)))
+        # fig, ax = plt.subplots(figsize=(5,5*(neuron_depth-1)*y0))
+
+        # -- Transformation functions
+        DC_to_FC = ax.transData.transform
+        FC_to_NFC = fig.transFigure.inverted().transform
+        # -- Take data coordinates and transform them to normalized figure coordinates
+        DC_to_NFC = lambda x: FC_to_NFC(DC_to_FC(x))
+
+        # plot scatters and lines
+        for l in range(neuron_depth):
+
+            n = width_in[l]
+
+            # scatters
+            for i in range(n):
+                plt.scatter(1 / (2 * n) + i / n, l * (y0 + z0), s=min_spacing ** 2 * 10000 * scale ** 2, color='black')
+
+            # plot connections (input to pre-mult)
+            for i in range(n):
+                if l < neuron_depth - 1:
+                    n_next = width_out[l + 1]
+                    N = n * n_next
+                    for j in range(n_next):
+                        id_ = i * n_next + j
+
+                        symbol_mask = self.symbolic_fun[l].mask[j][i]
+                        numerical_mask = self.act_fun[l].mask[i][j]
+                        if symbol_mask == 1. and numerical_mask > 0.:
+                            color = 'purple'
+                            alpha_mask = 1.
+                        if symbol_mask == 1. and numerical_mask == 0.:
+                            color = "red"
+                            alpha_mask = 1.
+                        if symbol_mask == 0. and numerical_mask == 1.:
+                            color = "black"
+                            alpha_mask = 1.
+                        if symbol_mask == 0. and numerical_mask == 0.:
+                            color = "white"
+                            alpha_mask = 0.
+
+                        plt.plot([1 / (2 * n) + i / n, 1 / (2 * N) + id_ / N],
+                                 [l * (y0 + z0), l * (y0 + z0) + y0 / 2 - y1], color=color, lw=2 * scale,
+                                 alpha=alpha[l][j][i] * alpha_mask)
+                        plt.plot([1 / (2 * N) + id_ / N, 1 / (2 * n_next) + j / n_next],
+                                 [l * (y0 + z0) + y0 / 2 + y1, l * (y0 + z0) + y0], color=color, lw=2 * scale,
+                                 alpha=alpha[l][j][i] * alpha_mask)
+
+            # plot connections (pre-mult to post-mult, post-mult = next-layer input)
+            if l < neuron_depth - 1:
+                n_in = width_out[l + 1]
+                n_out = width_in[l + 1]
+                mult_id = 0
+                for i in range(n_in):
+                    if i < width[l + 1][0]:
+                        j = i
+                    else:
+                        if i == width[l + 1][0]:
+                            if isinstance(self.mult_arity, int):
+                                ma = self.mult_arity
+                            else:
+                                ma = self.mult_arity[l + 1][mult_id]
+                            current_mult_arity = ma
+                        if current_mult_arity == 0:
+                            mult_id += 1
+                            if isinstance(self.mult_arity, int):
+                                ma = self.mult_arity
+                            else:
+                                ma = self.mult_arity[l + 1][mult_id]
+                            current_mult_arity = ma
+                        j = width[l + 1][0] + mult_id
+                        current_mult_arity -= 1
+                        #j = (i-width[l+1][0])//self.mult_arity + width[l+1][0]
+                    plt.plot([1 / (2 * n_in) + i / n_in, 1 / (2 * n_out) + j / n_out],
+                             [l * (y0 + z0) + y0, (l + 1) * (y0 + z0)], color='black', lw=2 * scale)
+
+            plt.xlim(0, 1)
+            plt.ylim(-0.1 * (y0 + z0), (neuron_depth - 1 + 0.1) * (y0 + z0))
+
+        plt.axis('off')
+
+        for l in range(neuron_depth - 1):
+            # plot splines
+            n = width_in[l]
+            for i in range(n):
+                n_next = width_out[l + 1]
+                N = n * n_next
+                for j in range(n_next):
+                    id_ = i * n_next + j
+                    im = plt.imread(f'{folder}/sp_{l}_{i}_{j}.png')
+                    left = DC_to_NFC([1 / (2 * N) + id_ / N - y1, 0])[0]
+                    right = DC_to_NFC([1 / (2 * N) + id_ / N + y1, 0])[0]
+                    bottom = DC_to_NFC([0, l * (y0 + z0) + y0 / 2 - y1])[1]
+                    up = DC_to_NFC([0, l * (y0 + z0) + y0 / 2 + y1])[1]
+                    newax = fig.add_axes([left, bottom, right - left, up - bottom])
+                    # newax = fig.add_axes([1/(2*N)+id_/N-y1, (l+1/2)*y0-y1, y1, y1], anchor='NE')
+                    newax.imshow(im, alpha=alpha[l][j][i])
+                    newax.axis('off')
+
+            # plot sum symbols
+            N = n = width_out[l + 1]
+            for j in range(n):
+                id_ = j
+                path = os.path.dirname(os.path.abspath(__file__)) + "/assets/img/sum_symbol.png"
+                im = plt.imread(path)
+                left = DC_to_NFC([1 / (2 * N) + id_ / N - y2, 0])[0]
+                right = DC_to_NFC([1 / (2 * N) + id_ / N + y2, 0])[0]
+                bottom = DC_to_NFC([0, l * (y0 + z0) + y0 - y2])[1]
+                up = DC_to_NFC([0, l * (y0 + z0) + y0 + y2])[1]
+                newax = fig.add_axes([left, bottom, right - left, up - bottom])
+                newax.imshow(im)
+                newax.axis('off')
+
+            # plot mult symbols
+            N = n = width_in[l + 1]
+            n_sum = width[l + 1][0]
+            n_mult = width[l + 1][1]
+            for j in range(n_mult):
+                id_ = j + n_sum
+                path = os.path.dirname(os.path.abspath(__file__)) + "/assets/img/mult_symbol.png"
+                im = plt.imread(path)
+                left = DC_to_NFC([1 / (2 * N) + id_ / N - y2, 0])[0]
+                right = DC_to_NFC([1 / (2 * N) + id_ / N + y2, 0])[0]
+                bottom = DC_to_NFC([0, (l + 1) * (y0 + z0) - y2])[1]
+                up = DC_to_NFC([0, (l + 1) * (y0 + z0) + y2])[1]
+                newax = fig.add_axes([left, bottom, right - left, up - bottom])
+                newax.imshow(im)
+                newax.axis('off')
+
+        if in_vars is not None:
+            n = self.width_in[0]
+            for i in range(n):
+                if isinstance(in_vars[i], sympy.Expr):
+                    plt.gcf().get_axes()[0].text(1 / (2 * (n)) + i / (n), -0.1, f'${latex(in_vars[i])}$',
+                                                 fontsize=40 * scale * varscale, horizontalalignment='center',
+                                                 verticalalignment='center')
+                else:
+                    plt.gcf().get_axes()[0].text(1 / (2 * (n)) + i / (n), -0.1, in_vars[i],
+                                                 fontsize=40 * scale * varscale, horizontalalignment='center',
+                                                 verticalalignment='center')
+
+        if out_vars is not None:
+            n = self.width_in[-1]
+            for i in range(n):
+                if isinstance(out_vars[i], sympy.Expr):
+                    plt.gcf().get_axes()[0].text(1 / (2 * (n)) + i / (n), (y0 + z0) * (len(self.width) - 1) + 0.15,
+                                                 f'${latex(out_vars[i])}$', fontsize=40 * scale * varscale,
+                                                 horizontalalignment='center', verticalalignment='center')
+                else:
+                    plt.gcf().get_axes()[0].text(1 / (2 * (n)) + i / (n), (y0 + z0) * (len(self.width) - 1) + 0.15,
+                                                 out_vars[i], fontsize=40 * scale * varscale,
+                                                 horizontalalignment='center', verticalalignment='center')
+
+        if title is not None:
+            plt.gcf().get_axes()[0].text(0.5, (y0 + z0) * (len(self.width) - 1) + 0.3, title, fontsize=40 * scale,
+                                         horizontalalignment='center', verticalalignment='center')
+
+    def reg(self, reg_metric, lamb_l1, lamb_entropy, lamb_coef, lamb_coefdiff):
+        """
+        Get regularization
+
+        Args:
+        -----
+            reg_metric : the regularization metric
+                'edge_forward_spline_n', 'edge_forward_spline_u', 'edge_forward_sum', 'edge_backward', 'node_backward'
+            lamb_l1 : float
+                l1 penalty strength
+            lamb_entropy : float
+                entropy penalty strength
+            lamb_coef : float
+                coefficient penalty strength
+            lamb_coefdiff : float
+                coefficient smoothness strength
+
+        Returns:
+        --------
+            reg_ : torch.float
+
+        Example
+        -------
+        >>> model = KAN(width=[2,3,1], grid=5, k=3, noise_scale=1.)
+        >>> x = torch.rand(100,2)
+        >>> model.get_act(x)
+        >>> model.reg('edge_forward_spline_n', 1.0, 2.0, 1.0, 1.0)
+        """
+        if reg_metric == 'edge_forward_spline_n':
+            acts_scale = self.acts_scale_spline
+
+        elif reg_metric == 'edge_forward_sum':
+            acts_scale = self.acts_scale
+
+        elif reg_metric == 'edge_forward_spline_u':
+            acts_scale = self.edge_actscale
+
+        elif reg_metric == 'edge_backward':
+            acts_scale = self.edge_scores
+
+        elif reg_metric == 'node_backward':
+            acts_scale = self.node_attribute_scores
+
+        else:
+            raise Exception(f'reg_metric = {reg_metric} not recognized!')
+
+        reg_ = 0.
+        for i in range(len(acts_scale)):
+            vec = acts_scale[i]
+
+            l1 = torch.sum(vec)
+            p_row = vec / (torch.sum(vec, dim=1, keepdim=True) + 1)
+            p_col = vec / (torch.sum(vec, dim=0, keepdim=True) + 1)
+            entropy_row = - torch.mean(torch.sum(p_row * torch.log2(p_row + 1e-4), dim=1))
+            entropy_col = - torch.mean(torch.sum(p_col * torch.log2(p_col + 1e-4), dim=0))
+            reg_ += lamb_l1 * l1 + lamb_entropy * (entropy_row + entropy_col)  # both l1 and entropy
+
+        # regularize coefficient to encourage spline to be zero
+        for i in range(len(self.act_fun)):
+            coeff_l1 = torch.sum(torch.mean(torch.abs(self.act_fun[i].coef), dim=1))
+            coeff_diff_l1 = torch.sum(torch.mean(torch.abs(torch.diff(self.act_fun[i].coef)), dim=1))
+            reg_ += lamb_coef * coeff_l1 + lamb_coefdiff * coeff_diff_l1
+
+        return reg_
+
+    def get_reg(self, reg_metric, lamb_l1, lamb_entropy, lamb_coef, lamb_coefdiff):
+        """
+        Get regularization. This seems unnecessary but in case a class wants to inherit this, it may want to rewrite get_reg, but not reg.
+        """
+        return self.reg(reg_metric, lamb_l1, lamb_entropy, lamb_coef, lamb_coefdiff)
+
+    def disable_symbolic_in_fit(self, lamb):
+        """
+        during fitting, disable symbolic if either is true (lamb = 0, none of symbolic functions is active)
+        """
+        old_save_act = self.save_act
+        if lamb == 0.:
+            self.save_act = False
+
+        # skip symbolic if no symbolic is turned on
+        depth = len(self.symbolic_fun)
+        no_symbolic = True
+        for l in range(depth):
+            no_symbolic *= torch.sum(torch.abs(self.symbolic_fun[l].mask)) == 0
+
+        old_symbolic_enabled = self.symbolic_enabled
+
+        if no_symbolic:
+            self.symbolic_enabled = False
+
+        return old_save_act, old_symbolic_enabled
+
+    def get_params(self):
+        """
+        Get parameters
+        """
+        return self.parameters()
+
+    def fit(self, dataset, opt="LBFGS", steps=100, log=1, lamb=0., lamb_l1=1., lamb_entropy=2., lamb_coef=0.,
+            lamb_coefdiff=0., update_grid=True, grid_update_num=10, loss_fn=None, lr=1., start_grid_update_step=-1,
+            stop_grid_update_step=50, batch=-1,
+            metrics=None, save_fig=False, in_vars=None, out_vars=None, beta=3, save_fig_freq=1, img_folder='./video',
+            singularity_avoiding=False, y_th=1000., reg_metric='edge_forward_spline_n', display_metrics=None):
+        '''
+        training
+
+        Args:
+        -----
+            dataset : dic
+                contains dataset['train_input'], dataset['train_label'], dataset['test_input'], dataset['test_label']
+            opt : str
+                "LBFGS" or "Adam"
+            steps : int
+                training steps
+            log : int
+                logging frequency
+            lamb : float
+                overall penalty strength
+            lamb_l1 : float
+                l1 penalty strength
+            lamb_entropy : float
+                entropy penalty strength
+            lamb_coef : float
+                coefficient magnitude penalty strength
+            lamb_coefdiff : float
+                difference of nearby coefficits (smoothness) penalty strength
+            update_grid : bool
+                If True, update grid regularly before stop_grid_update_step
+            grid_update_num : int
+                the number of grid updates before stop_grid_update_step
+            start_grid_update_step : int
+                no grid updates before this training step
+            stop_grid_update_step : int
+                no grid updates after this training step
+            loss_fn : function
+                loss function
+            lr : float
+                learning rate
+            batch : int
+                batch size, if -1 then full.
+            save_fig_freq : int
+                save figure every (save_fig_freq) steps
+            singularity_avoiding : bool
+                indicate whether to avoid singularity for the symbolic part
+            y_th : float
+                singularity threshold (anything above the threshold is considered singular and is softened in some ways)
+            reg_metric : str
+                regularization metric. Choose from {'edge_forward_spline_n', 'edge_forward_spline_u', 'edge_forward_sum', 'edge_backward', 'node_backward'}
+            metrics : a list of metrics (as functions)
+                the metrics to be computed in training
+            display_metrics : a list of functions
+                the metric to be displayed in tqdm progress bar
+
+        Returns:
+        --------
+            results : dic
+                results['train_loss'], 1D array of training losses (RMSE)
+                results['test_loss'], 1D array of test losses (RMSE)
+                results['reg'], 1D array of regularization
+                other metrics specified in metrics
+
+        Example
+        -------
+        # >>> from yms_kan import *
+        # >>> model = KAN(width=[2,5,1], grid=5, k=3, noise_scale=0.3, seed=2)
+        # >>> f = lambda x: torch.exp(torch.sin(torch.pi*x[:,[0]]) + x[:,[1]]**2)
+        # >>> dataset = create_dataset(f, n_var=2)
+        # >>> model.fit(dataset, opt='LBFGS', steps=20, lamb=0.001);
+        # >>> model.plot()
+        # Most examples in toturals involve the fit() method. Please check them for useness.
+        '''
+
+        if lamb > 0. and not self.save_act:
+            print('setting lamb=0. If you want to set lamb > 0, set self.save_act=True')
+
+        old_save_act, old_symbolic_enabled = self.disable_symbolic_in_fit(lamb)
+
+        pbar = tqdm(range(steps), desc='description', ncols=100)
+
+        if loss_fn is None:
+            loss_fn = loss_fn_eval = lambda x, y: torch.mean((x - y) ** 2)
+        else:
+            loss_fn = loss_fn_eval = loss_fn
+
+        grid_update_freq = int(stop_grid_update_step / grid_update_num)
+
+        if opt == "Adam":
+            optimizer = torch.optim.Adam(self.get_params(), lr=lr)
+        elif opt == "LBFGS":
+            optimizer = LBFGS(self.get_params(), lr=lr, history_size=10, line_search_fn="strong_wolfe",
+                              tolerance_grad=1e-32, tolerance_change=1e-32, tolerance_ys=1e-32)
+
+        results = {'train_loss': [], 'test_loss': [], 'reg': []}
+        if metrics is not None:
+            for i in range(len(metrics)):
+                results[metrics[i].__name__] = []
+
+        if batch == -1 or batch > dataset['train_input'].shape[0]:
+            batch_size = dataset['train_input'].shape[0]
+            batch_size_test = dataset['test_input'].shape[0]
+        else:
+            batch_size = batch
+            batch_size_test = batch
+
+        global train_loss, reg_
+
+        def closure():
+            global train_loss, reg_
+            optimizer.zero_grad()
+            pred = self.forward(dataset['train_input'][train_id], singularity_avoiding=singularity_avoiding, y_th=y_th)
+            train_loss = loss_fn(pred, dataset['train_label'][train_id])
+            if self.save_act:
+                if reg_metric == 'edge_backward':
+                    self.attribute()
+                if reg_metric == 'node_backward':
+                    self.node_attribute()
+                reg_ = self.get_reg(reg_metric, lamb_l1, lamb_entropy, lamb_coef, lamb_coefdiff)
+            else:
+                reg_ = torch.tensor(0.)
+            objective = train_loss + lamb * reg_
+            objective.backward()
+            return objective
+
+        if save_fig:
+            if not os.path.exists(img_folder):
+                os.makedirs(img_folder)
+
+        for _ in pbar:
+
+            if _ == steps - 1 and old_save_act:
+                self.save_act = True
+
+            if save_fig and _ % save_fig_freq == 0:
+                save_act = self.save_act
+                self.save_act = True
+
+            train_id = np.random.choice(dataset['train_input'].shape[0], batch_size, replace=False)
+            test_id = np.random.choice(dataset['test_input'].shape[0], batch_size_test, replace=False)
+
+            if _ % grid_update_freq == 0 and _ < stop_grid_update_step and update_grid and _ >= start_grid_update_step:
+                self.update_grid(dataset['train_input'][train_id])
+
+            if opt == "LBFGS":
+                optimizer.step(closure)
+
+            if opt == "Adam":
+                pred = self.forward(dataset['train_input'][train_id], singularity_avoiding=singularity_avoiding,
+                                    y_th=y_th)
+                train_loss = loss_fn(pred, dataset['train_label'][train_id])
+                if self.save_act:
+                    if reg_metric == 'edge_backward':
+                        self.attribute()
+                    if reg_metric == 'node_backward':
+                        self.node_attribute()
+                    reg_ = self.get_reg(reg_metric, lamb_l1, lamb_entropy, lamb_coef, lamb_coefdiff)
+                else:
+                    reg_ = torch.tensor(0.)
+                loss = train_loss + lamb * reg_
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+            test_loss = loss_fn_eval(self.forward(dataset['test_input'][test_id]), dataset['test_label'][test_id])
+
+            if metrics is not None:
+                for i in range(len(metrics)):
+                    results[metrics[i].__name__].append(metrics[i]().item())
+
+            results['train_loss'].append(torch.sqrt(train_loss).cpu().detach().numpy())
+            results['test_loss'].append(torch.sqrt(test_loss).cpu().detach().numpy())
+            results['reg'].append(reg_.cpu().detach().numpy())
+
+            if _ % log == 0:
+                if display_metrics is None:
+                    pbar.set_description("| train_loss: %.6f | test_loss: %.6f | reg: %.6f | " % (
+                        torch.sqrt(train_loss).cpu().detach().numpy(), torch.sqrt(test_loss).cpu().detach().numpy(),
+                        reg_.cpu().detach().numpy()))
+                else:
+                    string = ''
+                    data = ()
+                    for metric in display_metrics:
+                        string += f' {metric}: %.6f |'
+                        try:
+                            results[metric]
+                        except:
+                            raise Exception(f'{metric} not recognized')
+                        data += (results[metric][-1],)
+                    pbar.set_description(string % data)
+
+            if save_fig and _ % save_fig_freq == 0:
+                self.plot(folder=img_folder, in_vars=in_vars, out_vars=out_vars, title="Step {}".format(_), beta=beta)
+                plt.savefig(img_folder + '/' + str(_) + '.jpg', bbox_inches='tight', dpi=100)
+                plt.close()
+                self.save_act = save_act
+
+        self.log_history('fit')
+        # revert back to original state
+        self.symbolic_enabled = old_symbolic_enabled
+        return results
+
+    def fix(self, dataset: dict, batch_size, batch_size_test, opt="LBFGS", epochs=100, log=1, lamb=0.,
+            lamb_l1=1., label=None,
+            lamb_entropy=2., lamb_coef=0.,
+            lamb_coefdiff=0., update_grid=True, grid_update_num=10, loss_fn=None, lr=1., start_grid_update_step=-1,
+            stop_grid_update_step=100,
+            metrics=None, save_fig=False, in_vars=None, out_vars=None, beta=3, save_fig_freq=1, img_folder='./video',
+            singularity_avoiding=False, y_th=1000., reg_metric='edge_forward_spline_n', display_metrics=None):
+        """
+        Args:
+            label:
+            batch_size:
+            batch_size_test:
+            lamb_coefdiff:
+            metrics:
+            stop_grid_update_step:
+            save_fig:
+            singularity_avoiding:
+            y_th:
+            in_vars:
+            update_grid:
+            grid_update_num:
+            loss_fn:
+            lr:
+            beta:
+            out_vars:
+            save_fig_freq:
+            reg_metric:
+            display_metrics:
+            img_folder:
+            start_grid_update_step:
+            lamb_coef:
+            lamb_entropy:
+            lamb_l1:
+            lamb:
+            log:
+            epochs:
+            opt:
+            dataset (dict):
+        """
+
+        def calculate_results(all_labels, all_predictions, average='macro'):
+            result = {
+                'accuracy': accuracy_score(y_true=all_labels, y_pred=all_predictions),
+                'precision': precision_score(y_true=all_labels, y_pred=all_predictions, average=average),
+                'recall': recall_score(y_true=all_labels, y_pred=all_predictions, average=average),
+                'f1_score': f1_score(y_true=all_labels, y_pred=all_predictions, average=average),
+                # 'cm': confusion_matrix(y_true=all_labels, y_pred=all_predictions, labels=np.arange(len(classes)))
+            }
+            return result
+
+        all_predictions = []
+        all_labels = []
+        if lamb > 0. and not self.save_act:
+            print('setting lamb=0. If you want to set lamb > 0, set self.save_act=True')
+
+        old_save_act, old_symbolic_enabled = self.disable_symbolic_in_fit(lamb)
+        if label is not None:
+            label = label.to(self.device)
+        pbar = tqdm(range(epochs), desc='description', file=sys.stdout, colour='red')
+
+        if loss_fn is None:
+            loss_fn = lambda x, y: torch.mean((x - y) ** 2)
+        else:
+            loss_fn = loss_fn
+
+        grid_update_freq = int(stop_grid_update_step / grid_update_num)
+
+        if opt == "Adam":
+            optimizer = torch.optim.Adam(self.get_params(), lr=lr)
+        elif opt == "LBFGS":
+            optimizer = LBFGS(self.get_params(), lr=lr, history_size=10, line_search_fn="strong_wolfe",
+                              tolerance_grad=1e-32, tolerance_change=1e-32, tolerance_ys=1e-32)
+        else:
+            optimizer = torch.optim.SGD(self.get_params(), lr=lr)
+        lr_scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=5, min_lr=1e-9)
+        results = {'train_loss': [], 'test_loss': [], 'reg': [], 'accuracy': [],
+                   'precision': [], 'recall': [], 'f1_score': []}
+        if metrics is not None:
+            for i in range(len(metrics)):
+                results[metrics[i].__name__] = []
+        steps = math.ceil(dataset['train_input'].shape[0] / batch_size)
+
+        train_loss = torch.zeros(1).to(self.device)
+        reg_ = torch.zeros(1).to(self.device)
+
+        def closure():
+            nonlocal train_loss, reg_
+            optimizer.zero_grad()
+            pred = self.forward(batch_train_input, singularity_avoiding=singularity_avoiding, y_th=y_th)
+            loss = loss_fn(pred, batch_train_label)
+            if self.save_act:
+                if reg_metric == 'edge_backward':
+                    self.attribute()
+                if reg_metric == 'node_backward':
+                    self.node_attribute()
+                reg_ = self.get_reg(reg_metric, lamb_l1, lamb_entropy, lamb_coef, lamb_coefdiff)
+            else:
+                reg_ = torch.tensor(0.)
+            objective = loss + lamb * reg_
+            train_loss = (train_loss * batch_num + objective.detach()) / (batch_num + 1)
+            objective.backward()
+            return objective
+
+        if save_fig:
+            if not os.path.exists(img_folder):
+                os.makedirs(img_folder)
+
+        for epoch in pbar:
+
+            if epoch == epochs - 1 and old_save_act:
+                self.save_act = True
+
+            if save_fig and epoch % save_fig_freq == 0:
+                save_act = self.save_act
+                self.save_act = True
+
+            train_indices = np.arange(dataset['train_input'].shape[0])
+            np.random.shuffle(train_indices)
+            for batch_num, i in enumerate(range(0, len(train_indices), batch_size), start=0):
+                step = epoch * steps + batch_num + 1
+                batch_train_id = train_indices[i:i + batch_size]
+                batch_train_input = dataset['train_input'][batch_train_id].to(self.device)
+                batch_train_label = dataset['train_label'][batch_train_id].to(self.device)
+
+                if step % grid_update_freq == 0 and step < stop_grid_update_step and update_grid and step >= start_grid_update_step:
+                    self.update_grid(batch_train_input)
+
+                if opt == "LBFGS":
+                    optimizer.step(closure)
+
+                # if opt == "Adam":
+                else:
+                    optimizer.zero_grad()
+                    pred = self.forward(batch_train_input, singularity_avoiding=singularity_avoiding,
+                                        y_th=y_th)
+                    loss = loss_fn(pred, batch_train_label)
+                    if self.save_act:
+                        if reg_metric == 'edge_backward':
+                            self.attribute()
+                        if reg_metric == 'node_backward':
+                            self.node_attribute()
+                        reg_ = self.get_reg(reg_metric, lamb_l1, lamb_entropy, lamb_coef, lamb_coefdiff)
+                    else:
+                        reg_ = torch.tensor(0.)
+                    loss = loss + lamb * reg_
+                    train_loss = (train_loss * batch_num + loss.detach()) / (batch_num + 1)
+                    loss.backward()
+                    optimizer.step()
+
+            test_loss = torch.zeros(1).to(self.device)
+            with torch.no_grad():
+                test_indices = np.arange(dataset['test_input'].shape[0])
+                np.random.shuffle(test_indices)
+                for batch_num, i in enumerate(range(0, len(test_indices), batch_size_test)):
+                    batch_test_id = test_indices[i:i + batch_size_test]
+                    batch_test_input = dataset['test_input'][batch_test_id].to(self.device)
+                    batch_test_label = dataset['test_label'][batch_test_id].to(self.device)
+
+                    outputs = self.forward(batch_test_input)
+
+                    loss = loss_fn(outputs, batch_test_label)
+
+                    test_loss = (test_loss * batch_num + loss.detach()) / (batch_num + 1)
+                    if label is not None:
+                        diffs = torch.abs(outputs - label)
+                        closest_indices = torch.argmin(diffs, dim=1)
+                        closest_values = label[closest_indices]
+                        all_predictions.extend(closest_values.detach().cpu().numpy())
+                        all_labels.extend(batch_test_label.detach().cpu().numpy())
+
+            lr_scheduler.step(test_loss)
+            if metrics is not None:
+                for i in range(len(metrics)):
+                    results[metrics[i].__name__].append(metrics[i]().item())
+
+            results['train_loss'].append(torch.sqrt(train_loss).cpu().detach().numpy())
+            results['test_loss'].append(torch.sqrt(test_loss).cpu().detach().numpy())
+            results['reg'].append(reg_.cpu().detach().numpy())
+            if label is not None:
+                res = calculate_results(all_labels=all_labels, all_predictions=all_predictions)
+                results.update({key: results[key] + [value] for key, value in res.items() if key in results})
+
+            if epoch % log == 0:
+                if display_metrics is None:
+                    pbar.set_description(f"|train_loss:%.6f|test_loss:%.6f|reg:%.6f|" % (
+                        train_loss.item(), test_loss.item(), reg_.item()))
+                else:
+                    string = ''
+                    data = ()
+                    for metric in display_metrics:
+                        string += f' {metric}: %.6f |'
+                        try:
+                            results[metric]
+                        except:
+                            raise Exception(f'{metric} not recognized')
+                        data += (results[metric][-1],)
+                    pbar.set_description(string % data)
+
+            if save_fig and epoch % save_fig_freq == 0:
+                self.plot(folder=img_folder, in_vars=in_vars, out_vars=out_vars, title="Step {}".format(epoch),
+                          beta=beta)
+                plt.savefig(img_folder + '/' + str(epoch) + '.jpg', bbox_inches='tight', dpi=100)
+                plt.close()
+                self.save_act = save_act
+
+        self.log_history('fit')
+        self.symbolic_enabled = old_symbolic_enabled
+        return results
+
+    def prune_node(self, threshold=1e-2, mode="auto", active_neurons_id=None, log_history=True):
+        """
+        pruning nodes
+
+        Args:
+        -----
+            threshold : float
+                if the attribution score of a neuron is below the threshold, it is considered dead and will be removed
+            mode : str
+                'auto' or 'manual'. with 'auto', nodes are automatically pruned using threshold. with 'manual', active_neurons_id should be passed in.
+
+        Returns:
+        --------
+            pruned network : MultKAN
+
+        Example
+        -------
+        # >>> from yms_kan import *
+        # >>> model = KAN(width=[2,5,1], grid=5, k=3, noise_scale=0.3, seed=2)
+        # >>> f = lambda x: torch.exp(torch.sin(torch.pi*x[:,[0]]) + x[:,[1]]**2)
+        # >>> dataset = create_dataset(f, n_var=2)
+        # >>> model.fit(dataset, opt='LBFGS', steps=20, lamb=0.001);
+        # >>> model = model.prune_node()
+        # >>> model.plot()
+        """
+        if self.acts is None:
+            self.get_act()
+
+        mask_up = [torch.ones(self.width_in[0], device=self.device)]
+        mask_down = []
+        active_neurons_up = [list(range(self.width_in[0]))]
+        active_neurons_down = []
+        num_sums = []
+        num_mults = []
+        mult_arities = [[]]
+
+        if active_neurons_id != None:
+            mode = "manual"
+
+        for i in range(len(self.acts_scale) - 1):
+
+            mult_arity = []
+
+            if mode == "auto":
+                self.attribute()
+                overall_important_up = self.node_scores[i + 1] > threshold
+
+            elif mode == "manual":
+                overall_important_up = torch.zeros(self.width_in[i + 1], dtype=torch.bool, device=self.device)
+                overall_important_up[active_neurons_id[i]] = True
+
+            num_sum = torch.sum(overall_important_up[:self.width[i + 1][0]])
+            num_mult = torch.sum(overall_important_up[self.width[i + 1][0]:])
+            if self.mult_homo == True:
+                overall_important_down = torch.cat([overall_important_up[:self.width[i + 1][0]], (
+                    overall_important_up[self.width[i + 1][0]:][None, :].expand(self.mult_arity, -1)).T.reshape(-1, )],
+                                                   dim=0)
+            else:
+                overall_important_down = overall_important_up[:self.width[i + 1][0]]
+                for j in range(overall_important_up[self.width[i + 1][0]:].shape[0]):
+                    active_bool = overall_important_up[self.width[i + 1][0] + j]
+                    arity = self.mult_arity[i + 1][j]
+                    overall_important_down = torch.cat(
+                        [overall_important_down, torch.tensor([active_bool] * arity).to(self.device)])
+                    if active_bool:
+                        mult_arity.append(arity)
+
+            num_sums.append(num_sum.item())
+            num_mults.append(num_mult.item())
+
+            mask_up.append(overall_important_up.float())
+            mask_down.append(overall_important_down.float())
+
+            active_neurons_up.append(torch.where(overall_important_up == True)[0])
+            active_neurons_down.append(torch.where(overall_important_down == True)[0])
+
+            mult_arities.append(mult_arity)
+
+        active_neurons_down.append(list(range(self.width_out[-1])))
+        mask_down.append(torch.ones(self.width_out[-1], device=self.device))
+
+        if not self.mult_homo:
+            mult_arities.append(self.mult_arity[-1])
+
+        self.mask_up = mask_up
+        self.mask_down = mask_down
+
+        # update act_fun[l].mask up
+        for l in range(len(self.acts_scale) - 1):
+            for i in range(self.width_in[l + 1]):
+                if i not in active_neurons_up[l + 1]:
+                    self.remove_node(l + 1, i, mode='up', log_history=False)
+
+            for i in range(self.width_out[l + 1]):
+                if i not in active_neurons_down[l]:
+                    self.remove_node(l + 1, i, mode='down', log_history=False)
+
+        model2 = MultKAN(copy.deepcopy(self.width), grid=self.grid, k=self.k, base_fun=self.base_fun_name,
+                         mult_arity=self.mult_arity, ckpt_path=self.ckpt_path, auto_save=True, first_init=False,
+                         state_id=self.state_id, round=self.round).to(self.device)
+        model2.load_state_dict(self.state_dict())
+
+        width_new = [self.width[0]]
+
+        for i in range(len(self.acts_scale)):
+
+            if i < len(self.acts_scale) - 1:
+                num_sum = num_sums[i]
+                num_mult = num_mults[i]
+                model2.node_bias[i].data = model2.node_bias[i].data[active_neurons_up[i + 1]]
+                model2.node_scale[i].data = model2.node_scale[i].data[active_neurons_up[i + 1]]
+                model2.subnode_bias[i].data = model2.subnode_bias[i].data[active_neurons_down[i]]
+                model2.subnode_scale[i].data = model2.subnode_scale[i].data[active_neurons_down[i]]
+                model2.width[i + 1] = [num_sum, num_mult]
+
+                model2.act_fun[i].out_dim_sum = num_sum
+                model2.act_fun[i].out_dim_mult = num_mult
+
+                model2.symbolic_fun[i].out_dim_sum = num_sum
+                model2.symbolic_fun[i].out_dim_mult = num_mult
+
+                width_new.append([num_sum, num_mult])
+
+            model2.act_fun[i] = model2.act_fun[i].get_subset(active_neurons_up[i], active_neurons_down[i])
+            model2.symbolic_fun[i] = self.symbolic_fun[i].get_subset(active_neurons_up[i], active_neurons_down[i])
+
+        model2.cache_data = self.cache_data
+        model2.acts = None
+
+        width_new.append(self.width[-1])
+        model2.width = width_new
+
+        if not self.mult_homo:
+            model2.mult_arity = mult_arities
+
+        if log_history:
+            self.log_history('prune_node')
+            model2.state_id += 1
+
+        return model2
+
+    def prune_edge(self, threshold=3e-2, log_history=True):
+        """
+        pruning edges
+
+        Args:
+        -----
+            threshold : float
+                if the attribution score of an edge is below the threshold, it is considered dead and will be set to zero.
+
+        Returns:
+        --------
+            pruned network : MultKAN
+
+        Example
+        -------
+        # >>> from yms_kan import *
+        # >>> model = KAN(width=[2,5,1], grid=5, k=3, noise_scale=0.3, seed=2)
+        # >>> f = lambda x: torch.exp(torch.sin(torch.pi*x[:,[0]]) + x[:,[1]]**2)
+        # >>> dataset = create_dataset(f, n_var=2)
+        # >>> model.fit(dataset, opt='LBFGS', steps=20, lamb=0.001);
+        # >>> model = model.prune_edge()
+        # >>> model.plot()
+        """
+        if self.acts is None:
+            self.get_act()
+
+        for i in range(len(self.width) - 1):
+            # self.act_fun[i].mask.data = ((self.acts_scale[i] > threshold).permute(1,0)).float()
+            old_mask = self.act_fun[i].mask.data
+            self.act_fun[i].mask.data = ((self.edge_scores[i] > threshold).permute(1, 0) * old_mask).float()
+
+        if log_history:
+            self.log_history('fix_symbolic')
+
+    def prune(self, node_th=1e-2, edge_th=3e-2):
+        """
+        prune (both nodes and edges)
+
+        Args:
+        -----
+            node_th : float
+                if the attribution score of a node is below node_th, it is considered dead and will be set to zero.
+            edge_th : float
+                if the attribution score of an edge is below node_th, it is considered dead and will be set to zero.
+
+        Returns:
+        --------
+            pruned network : MultKAN
+
+        Example
+        -------
+        # >>> from yms_kan import *
+        # >>> model = KAN(width=[2,5,1], grid=5, k=3, noise_scale=0.3, seed=2)
+        # >>> f = lambda x: torch.exp(torch.sin(torch.pi*x[:,[0]]) + x[:,[1]]**2)
+        # >>> dataset = create_dataset(f, n_var=2)
+        # >>> model.fit(dataset, opt='LBFGS', steps=20, lamb=0.001);
+        # >>> model = model.prune()
+        # >>> model.plot()
+        """
+        if self.acts is None:
+            self.get_act()
+
+        self = self.prune_node(node_th, log_history=False)
+        # self.prune_node(node_th, log_history=False)
+        self.forward(self.cache_data)
+        self.attribute()
+        self.prune_edge(edge_th, log_history=False)
+        self.log_history('prune')
+        return self
+
+    def prune_input(self, threshold=1e-2, active_inputs=None, log_history=True):
+        """
+        prune inputs
+
+        Args:
+        -----
+            threshold : float
+                if the attribution score of the input feature is below threshold, it is considered irrelevant.
+            active_inputs : None or list
+                if a list is passed, the manual mode will disregard attribution score and prune as instructed.
+
+        Returns:
+        --------
+            pruned network : MultKAN
+
+        Example1
+        --------
+        >>> # automatic
+        # >>> from yms_kan import *
+        # >>> model = KAN(width=[3,5,1], grid=5, k=3, noise_scale=0.3, seed=2)
+        # >>> f = lambda x: 1 * x[:,[0]]**2 + 0.3 * x[:,[1]]**2 + 0.0 * x[:,[2]]**2
+        # >>> dataset = create_dataset(f, n_var=3)
+        # >>> model.fit(dataset, opt='LBFGS', steps=20, lamb=0.001);
+        # >>> model.plot()
+        # >>> model = model.prune_input()
+        # >>> model.plot()
+
+        Example2
+        --------
+        >>> # automatic
+        # >>> from yms_kan import *
+        # >>> model = KAN(width=[3,5,1], grid=5, k=3, noise_scale=0.3, seed=2)
+        # >>> f = lambda x: 1 * x[:,[0]]**2 + 0.3 * x[:,[1]]**2 + 0.0 * x[:,[2]]**2
+        # >>> dataset = create_dataset(f, n_var=3)
+        # >>> model.fit(dataset, opt='LBFGS', steps=20, lamb=0.001);
+        # >>> model.plot()
+        # >>> model = model.prune_input(active_inputs=[0,1])
+        # >>> model.plot()
+        """
+        if active_inputs is None:
+            self.attribute()
+            input_score = self.node_scores[0]
+            input_mask = input_score > threshold
+            print('keep:', input_mask.tolist())
+            input_id = torch.where(input_mask == True)[0]
+
+        else:
+            input_id = torch.tensor(active_inputs, dtype=torch.long).to(self.device)
+
+        model2 = MultKAN(copy.deepcopy(self.width), grid=self.grid, k=self.k, base_fun=self.base_fun,
+                         mult_arity=self.mult_arity, ckpt_path=self.ckpt_path, auto_save=True, first_init=False,
+                         state_id=self.state_id, round=self.round).to(self.device)
+        model2.load_state_dict(self.state_dict())
+
+        model2.act_fun[0] = model2.act_fun[0].get_subset(input_id, torch.arange(self.width_out[1]))
+        model2.symbolic_fun[0] = self.symbolic_fun[0].get_subset(input_id, torch.arange(self.width_out[1]))
+
+        model2.cache_data = self.cache_data
+        model2.acts = None
+
+        model2.width[0] = [len(input_id), 0]
+        model2.input_id = input_id
+
+        if log_history:
+            self.log_history('prune_input')
+            model2.state_id += 1
+
+        return model2
+
+    def remove_edge(self, l, i, j, log_history=True):
+        """
+        remove activtion phi(l,i,j) (set its mask to zero)
+        """
+        self.act_fun[l].mask[i][j] = 0.
+        if log_history:
+            self.log_history('remove_edge')
+
+    def remove_node(self, l, i, mode='all', log_history=True):
+        """
+        remove neuron (l,i) (set the masks of all incoming and outgoing activation functions to zero)
+        """
+        if mode == 'down':
+            self.act_fun[l - 1].mask[:, i] = 0.
+            self.symbolic_fun[l - 1].mask[i, :] *= 0.
+
+        elif mode == 'up':
+            self.act_fun[l].mask[i, :] = 0.
+            self.symbolic_fun[l].mask[:, i] *= 0.
+
+        else:
+            self.remove_node(l, i, mode='up')
+            self.remove_node(l, i, mode='down')
+
+        if log_history:
+            self.log_history('remove_node')
+
+    def attribute(self, l=None, i=None, out_score=None, plot=True):
+        """
+        get attribution scores
+
+        Args:
+        -----
+            l : None or int
+                layer index
+            i : None or int
+                neuron index
+            out_score : None or 1D torch.float
+                specify output scores
+            plot : bool
+                when plot = True, display the bar show
+
+        Returns:
+        --------
+            attribution scores
+
+        Example
+        -------
+        # >>> from yms_kan import *
+        # >>> model = KAN(width=[3,5,1], grid=5, k=3, noise_scale=0.3, seed=2)
+        # >>> f = lambda x: 1 * x[:,[0]]**2 + 0.3 * x[:,[1]]**2 + 0.0 * x[:,[2]]**2
+        # >>> dataset = create_dataset(f, n_var=3)
+        # >>> model.fit(dataset, opt='LBFGS', steps=20, lamb=0.001);
+        # >>> model.attribute()
+        # >>> model.feature_score
+        """
+        # output (out_dim, in_dim)
+
+        if l is not None:
+            self.attribute()
+            out_score = self.node_scores[l]
+
+        if self.acts == None:
+            self.get_act()
+
+        def score_node2subnode(node_score, width, mult_arity, out_dim):
+
+            assert np.sum(width) == node_score.shape[1]
+            if isinstance(mult_arity, int):
+                n_subnode = width[0] + mult_arity * width[1]
+            else:
+                n_subnode = width[0] + int(np.sum(mult_arity))
+
+            #subnode_score_leaf = torch.zeros(out_dim, n_subnode).requires_grad_(True)
+            #subnode_score = subnode_score_leaf.clone()
+            #subnode_score[:,:width[0]] = node_score[:,:width[0]]
+            subnode_score = node_score[:, :width[0]]
+            if isinstance(mult_arity, int):
+                #subnode_score[:,width[0]:] = node_score[:,width[0]:][:,:,None].expand(out_dim, node_score[width[0]:].shape[0], mult_arity).reshape(out_dim,-1)
+                subnode_score = torch.cat([subnode_score, node_score[:, width[0]:][:, :, None].expand(out_dim,
+                                                                                                      node_score[:,
+                                                                                                      width[0]:].shape[
+                                                                                                          1],
+                                                                                                      mult_arity).reshape(
+                    out_dim, -1)], dim=1)
+            else:
+                acml = width[0]
+                for i in range(len(mult_arity)):
+                    #subnode_score[:, acml:acml+mult_arity[i]] = node_score[:, width[0]+i]
+                    subnode_score = torch.cat(
+                        [subnode_score, node_score[:, width[0] + i].expand(out_dim, mult_arity[i])], dim=1)
+                    acml += mult_arity[i]
+            return subnode_score
+
+        node_scores = []
+        subnode_scores = []
+        edge_scores = []
+
+        l_query = l
+        if l is None:
+            l_end = self.depth
+        else:
+            l_end = l
+
+        # back propagate from the queried layer
+        out_dim = self.width_in[l_end]
+        if out_score is None:
+            node_score = torch.eye(out_dim).requires_grad_(True)
+        else:
+            node_score = torch.diag(out_score).requires_grad_(True)
+        node_scores.append(node_score)
+
+        device = self.act_fun[0].grid.device
+
+        for l in range(l_end, 0, -1):
+
+            # node to subnode 
+            if isinstance(self.mult_arity, int):
+                subnode_score = score_node2subnode(node_score, self.width[l], self.mult_arity, out_dim=out_dim)
+            else:
+                mult_arity = self.mult_arity[l]
+                #subnode_score = score_node2subnode(node_score, self.width[l], mult_arity)
+                subnode_score = score_node2subnode(node_score, self.width[l], mult_arity, out_dim=out_dim)
+
+            subnode_scores.append(subnode_score)
+            # subnode to edge
+            #print(self.edge_actscale[l-1].device, subnode_score.device, self.subnode_actscale[l-1].device)
+            edge_score = torch.einsum('ij,ki,i->kij', self.edge_actscale[l - 1], subnode_score.to(device),
+                                      1 / (self.subnode_actscale[l - 1] + 1e-4))
+            edge_scores.append(edge_score)
+
+            # edge to node
+            node_score = torch.sum(edge_score, dim=1)
+            node_scores.append(node_score)
+
+        self.node_scores_all = list(reversed(node_scores))
+        self.edge_scores_all = list(reversed(edge_scores))
+        self.subnode_scores_all = list(reversed(subnode_scores))
+
+        self.node_scores = [torch.mean(l, dim=0) for l in self.node_scores_all]
+        self.edge_scores = [torch.mean(l, dim=0) for l in self.edge_scores_all]
+        self.subnode_scores = [torch.mean(l, dim=0) for l in self.subnode_scores_all]
+
+        # return
+        if l_query != None:
+            if i == None:
+                return self.node_scores_all[0]
+            else:
+
+                # plot
+                if plot:
+                    in_dim = self.width_in[0]
+                    plt.figure(figsize=(1 * in_dim, 3))
+                    plt.bar(range(in_dim), self.node_scores_all[0][i].cpu().detach().numpy())
+                    plt.xticks(range(in_dim));
+
+                return self.node_scores_all[0][i]
+
+    def node_attribute(self):
+        self.node_attribute_scores = []
+        for l in range(1, self.depth + 1):
+            node_attr = self.attribute(l)
+            self.node_attribute_scores.append(node_attr)
+
+    def feature_interaction(self, l, neuron_th=1e-2, feature_th=1e-2):
+        """
+        get feature interaction
+
+        Args:
+        -----
+            l : int
+                layer index
+            neuron_th : float
+                threshold to determine whether a neuron is active
+            feature_th : float
+                threshold to determine whether a feature is active
+
+        Returns:
+        --------
+            dictionary
+
+        Example
+        -------
+        # >>> from yms_kan import *
+        # >>> model = KAN(width=[3,5,1], grid=5, k=3, noise_scale=0.3, seed=2)
+        # >>> f = lambda x: 1 * x[:,[0]]**2 + 0.3 * x[:,[1]]**2 + 0.0 * x[:,[2]]**2
+        # >>> dataset = create_dataset(f, n_var=3)
+        # >>> model.fit(dataset, opt='LBFGS', steps=20, lamb=0.001);
+        # >>> model.attribute()
+        # >>> model.feature_interaction(1)
+        """
+        dic = {}
+        width = self.width_in[l]
+
+        for i in range(width):
+            score = self.attribute(l, i, plot=False)
+
+            if torch.max(score) > neuron_th:
+                features = tuple(torch.where(score > torch.max(score) * feature_th)[0].detach().numpy())
+                if features in dic.keys():
+                    dic[features] += 1
+                else:
+                    dic[features] = 1
+
+        return dic
+
+    def suggest_symbolic(self, l, i, j, a_range=(-10, 10), b_range=(-10, 10), lib=None, topk=5, verbose=True,
+                         r2_loss_fun=lambda x: np.log2(1 + 1e-5 - x), c_loss_fun=lambda x: x, weight_simple=0.8):
+        """
+        suggest symbolic function
+
+        Args:
+        -----
+            l : int
+                layer index
+            i : int
+                neuron index in layer l
+            j : int
+                neuron index in layer j
+            a_range : tuple
+                search range of a
+            b_range : tuple
+                search range of b
+            lib : list of str
+                library of candidate symbolic functions
+            topk : int
+                the number of top functions displayed
+            verbose : bool
+                if verbose = True, print more information
+            r2_loss_fun : functoon
+                function : r2 -> "bits"
+            c_loss_fun : fun
+                function : c -> 'bits'
+            weight_simple : float
+                the simplifty weight: the higher, more prefer simplicity over performance
+
+
+        Returns:
+        --------
+            best_name (str), best_fun (function), best_r2 (float), best_c (float)
+
+        Example
+        -------
+        # >>> from yms_kan import *
+        # >>> model = KAN(width=[2,1,1], grid=5, k=3, noise_scale=0.0, seed=0)
+        # >>> f = lambda x: torch.exp(torch.sin(torch.pi*x[:,[0]])+x[:,[1]]**2)
+        # >>> dataset = create_dataset(f, n_var=3)
+        # >>> model.fit(dataset, opt='LBFGS', steps=20, lamb=0.001);
+        # >>> model.suggest_symbolic(0,1,0)
+        """
+        r2s = []
+        cs = []
+
+        if lib == None:
+            symbolic_lib = SYMBOLIC_LIB
+        else:
+            symbolic_lib = {}
+            for item in lib:
+                symbolic_lib[item] = SYMBOLIC_LIB[item]
+
+        # getting r2 and complexities
+        for (name, content) in symbolic_lib.items():
+            r2 = self.fix_symbolic(l, i, j, name, a_range=a_range, b_range=b_range, verbose=False, log_history=False)
+            if r2 == -1e8:  # zero function
+                r2s.append(-1e8)
+            else:
+                r2s.append(r2.item())
+                self.unfix_symbolic(l, i, j, log_history=False)
+            c = content[2]
+            cs.append(c)
+
+        r2s = np.array(r2s)
+        cs = np.array(cs)
+        r2_loss = r2_loss_fun(r2s).astype('float')
+        cs_loss = c_loss_fun(cs)
+
+        loss = weight_simple * cs_loss + (1 - weight_simple) * r2_loss
+
+        sorted_ids = np.argsort(loss)[:topk]
+        r2s = r2s[sorted_ids][:topk]
+        cs = cs[sorted_ids][:topk]
+        r2_loss = r2_loss[sorted_ids][:topk]
+        cs_loss = cs_loss[sorted_ids][:topk]
+        loss = loss[sorted_ids][:topk]
+
+        topk = np.minimum(topk, len(symbolic_lib))
+
+        if verbose:
+            # print results in a dataframe
+            results = {}
+            results['function'] = [list(symbolic_lib.items())[sorted_ids[i]][0] for i in range(topk)]
+            results['fitting r2'] = r2s[:topk]
+            results['r2 loss'] = r2_loss[:topk]
+            results['complexity'] = cs[:topk]
+            results['complexity loss'] = cs_loss[:topk]
+            results['total loss'] = loss[:topk]
+
+            df = pd.DataFrame(results)
+            print(df)
+
+        best_name = list(symbolic_lib.items())[sorted_ids[0]][0]
+        best_fun = list(symbolic_lib.items())[sorted_ids[0]][1]
+        best_r2 = r2s[0]
+        best_c = cs[0]
+
+        return best_name, best_fun, best_r2, best_c;
+
+    def auto_symbolic(self, a_range=(-10, 10), b_range=(-10, 10), lib=None, verbose=1, weight_simple=0.8,
+                      r2_threshold=0.0):
+        """
+        automatic symbolic regression for all edges
+
+        Args:
+        -----
+            a_range : tuple
+                search range of a
+            b_range : tuple
+                search range of b
+            lib : list of str
+                library of candidate symbolic functions
+            verbose : int
+                larger verbosity => more verbosity
+            weight_simple : float
+                a weight that prioritizies simplicity (low complexity) over performance (high r2) - set to 0.0 to ignore complexity
+            r2_threshold : float
+                If r2 is below this threshold, the edge will not be fixed with any symbolic function - set to 0.0 to ignore this threshold
+        Returns:
+        --------
+            None
+
+        Example
+        -------
+        # >>> from yms_kan import *
+        # >>> model = KAN(width=[2,1,1], grid=5, k=3, noise_scale=0.0, seed=0)
+        # >>> f = lambda x: torch.exp(torch.sin(torch.pi*x[:,[0]])+x[:,[1]]**2)
+        # >>> dataset = create_dataset(f, n_var=3)
+        # >>> model.fit(dataset, opt='LBFGS', steps=20, lamb=0.001);
+        # >>> model.auto_symbolic()
+        """
+        for l in range(len(self.width_in) - 1):
+            for i in range(self.width_in[l]):
+                for j in range(self.width_out[l + 1]):
+                    if self.symbolic_fun[l].mask[j, i] > 0. and self.act_fun[l].mask[i][j] == 0.:
+                        print(f'skipping ({l},{i},{j}) since already symbolic')
+                    elif self.symbolic_fun[l].mask[j, i] == 0. and self.act_fun[l].mask[i][j] == 0.:
+                        self.fix_symbolic(l, i, j, '0', verbose=verbose > 1, log_history=False)
+                        print(f'fixing ({l},{i},{j}) with 0')
+                    else:
+                        name, fun, r2, c = self.suggest_symbolic(l, i, j, a_range=a_range, b_range=b_range, lib=lib,
+                                                                 verbose=False, weight_simple=weight_simple)
+                        if r2 >= r2_threshold:
+                            self.fix_symbolic(l, i, j, name, verbose=verbose > 1, log_history=False)
+                            if verbose >= 1:
+                                print(f'fixing ({l},{i},{j}) with {name}, r2={r2}, c={c}')
+                        else:
+                            print(
+                                f'For ({l},{i},{j}) the best fit was {name}, but r^2 = {r2} and this is lower than {r2_threshold}. This edge was omitted, keep training or try a different threshold.')
+
+        self.log_history('auto_symbolic')
+
+    def symbolic_formula(self, var=None, normalizer=None, output_normalizer=None):
+        """
+        get symbolic formula
+
+        Args:
+        -----
+            var : None or a list of sympy expression
+                input variables
+            normalizer : [mean, std]
+            output_normalizer : [mean, std]
+
+        Returns:
+        --------
+            None
+
+        Example
+        -------
+        # >>> from yms_kan import *
+        # >>> model = KAN(width=[2,1,1], grid=5, k=3, noise_scale=0.0, seed=0)
+        # >>> f = lambda x: torch.exp(torch.sin(torch.pi*x[:,[0]])+x[:,[1]]**2)
+        # >>> dataset = create_dataset(f, n_var=3)
+        # >>> model.fit(dataset, opt='LBFGS', steps=20, lamb=0.001);
+        # >>> model.auto_symbolic()
+        # >>> model.symbolic_formula()[0][0]
+        """
+
+        symbolic_acts = []
+        symbolic_acts_premult = []
+        x = []
+
+        def ex_round(ex1, n_digit):
+            ex2 = ex1
+            for a in sympy.preorder_traversal(ex1):
+                if isinstance(a, sympy.Float):
+                    ex2 = ex2.subs(a, round(a, n_digit))
+            return ex2
+
+        # define variables
+        if var is None:
+            for ii in range(1, self.width[0][0] + 1):
+                exec(f"x{ii} = sympy.Symbol('x_{ii}')")
+                exec(f"x.append(x{ii})")
+        elif isinstance(var[0], sympy.Expr):
+            x = var
+        else:
+            x = [sympy.symbols(var_) for var_ in var]
+
+        x0 = x
+
+        if normalizer is not None:
+            mean = normalizer[0]
+            std = normalizer[1]
+            x = [(x[i] - mean[i]) / std[i] for i in range(len(x))]
+
+        symbolic_acts.append(x)
+
+        for l in range(len(self.width_in) - 1):
+            num_sum = self.width[l + 1][0]
+            num_mult = self.width[l + 1][1]
+            y = []
+            for j in range(self.width_out[l + 1]):
+                yj = 0.
+                for i in range(self.width_in[l]):
+                    a, b, c, d = self.symbolic_fun[l].affine[j, i]
+                    sympy_fun = self.symbolic_fun[l].funs_sympy[j][i]
+                    try:
+                        yj += c * sympy_fun(a * x[i] + b) + d
+                    except:
+                        print('make sure all activations need to be converted to symbolic formulas first!')
+                        return
+                yj = self.subnode_scale[l][j] * yj + self.subnode_bias[l][j]
+                if simplify == True:
+                    y.append(sympy.simplify(yj))
+                else:
+                    y.append(yj)
+
+            symbolic_acts_premult.append(y)
+
+            mult = []
+            for k in range(num_mult):
+                if isinstance(self.mult_arity, int):
+                    mult_arity = self.mult_arity
+                else:
+                    mult_arity = self.mult_arity[l + 1][k]
+                for i in range(mult_arity - 1):
+                    if i == 0:
+                        mult_k = y[num_sum + 2 * k] * y[num_sum + 2 * k + 1]
+                    else:
+                        mult_k = mult_k * y[num_sum + 2 * k + i + 1]
+                mult.append(mult_k)
+
+            y = y[:num_sum] + mult
+
+            for j in range(self.width_in[l + 1]):
+                y[j] = self.node_scale[l][j] * y[j] + self.node_bias[l][j]
+
+            x = y
+            symbolic_acts.append(x)
+
+        if output_normalizer != None:
+            output_layer = symbolic_acts[-1]
+            means = output_normalizer[0]
+            stds = output_normalizer[1]
+
+            assert len(output_layer) == len(means), 'output_normalizer does not match the output layer'
+            assert len(output_layer) == len(stds), 'output_normalizer does not match the output layer'
+
+            output_layer = [(output_layer[i] * stds[i] + means[i]) for i in range(len(output_layer))]
+            symbolic_acts[-1] = output_layer
+
+        self.symbolic_acts = [[symbolic_acts[l][i] for i in range(len(symbolic_acts[l]))] for l in
+                              range(len(symbolic_acts))]
+        self.symbolic_acts_premult = [[symbolic_acts_premult[l][i] for i in range(len(symbolic_acts_premult[l]))] for l
+                                      in range(len(symbolic_acts_premult))]
+
+        out_dim = len(symbolic_acts[-1])
+        #return [symbolic_acts[-1][i] for i in range(len(symbolic_acts[-1]))], x0
+
+        if simplify:
+            return [symbolic_acts[-1][i] for i in range(len(symbolic_acts[-1]))], x0
+        else:
+            return [symbolic_acts[-1][i] for i in range(len(symbolic_acts[-1]))], x0
+
+    def expand_depth(self):
+        '''
+        expand network depth, add an indentity layer to the end. For usage, please refer to tutorials interp_3_KAN_compiler.ipynb.
+        
+        Args:
+        -----
+            var : None or a list of sympy expression
+                input variables
+            normalizer : [mean, std]
+            output_normalizer : [mean, std]
+            
+        Returns:
+        --------
+            None
+        '''
+        self.depth += 1
+
+        # add kanlayer, set mask to zero
+        dim_out = self.width_in[-1]
+        layer = KANLayer(dim_out, dim_out, num=self.grid, k=self.k)
+        layer.mask *= 0.
+        self.act_fun.append(layer)
+
+        self.width.append([dim_out, 0])
+        self.mult_arity.append([])
+
+        # add symbolic_kanlayer set mask to one. fun = identity on diagonal and zero for off-diagonal
+        layer = Symbolic_KANLayer(dim_out, dim_out)
+        layer.mask += 1.
+
+        for j in range(dim_out):
+            for i in range(dim_out):
+                if i == j:
+                    layer.fix_symbolic(i, j, 'x')
+                else:
+                    layer.fix_symbolic(i, j, '0')
+
+        self.symbolic_fun.append(layer)
+
+        self.node_bias.append(
+            torch.nn.Parameter(torch.zeros(dim_out, device=self.device)).requires_grad_(self.affine_trainable))
+        self.node_scale.append(
+            torch.nn.Parameter(torch.ones(dim_out, device=self.device)).requires_grad_(self.affine_trainable))
+        self.subnode_bias.append(
+            torch.nn.Parameter(torch.zeros(dim_out, device=self.device)).requires_grad_(self.affine_trainable))
+        self.subnode_scale.append(
+            torch.nn.Parameter(torch.ones(dim_out, device=self.device)).requires_grad_(self.affine_trainable))
+
+    def expand_width(self, layer_id, n_added_nodes, sum_bool=True, mult_arity=2):
+        '''
+        expand network width. For usage, please refer to tutorials interp_3_KAN_compiler.ipynb.
+        
+        Args:
+        -----
+            layer_id : int
+                layer index
+            n_added_nodes : init
+                the number of added nodes
+            sum_bool : bool
+                if sum_bool == True, added nodes are addition nodes; otherwise multiplication nodes
+            mult_arity : init
+                multiplication arity (the number of numbers to be multiplied)
+            
+        Returns:
+        --------
+            None
+        '''
+
+        def _expand(layer_id, n_added_nodes, sum_bool=True, mult_arity=2, added_dim='out'):
+            l = layer_id
+            in_dim = self.symbolic_fun[l].in_dim
+            out_dim = self.symbolic_fun[l].out_dim
+            if sum_bool:
+
+                if added_dim == 'out':
+                    new = Symbolic_KANLayer(in_dim, out_dim + n_added_nodes)
+                    old = self.symbolic_fun[l]
+                    in_id = np.arange(in_dim)
+                    out_id = np.arange(out_dim + n_added_nodes)
+
+                    for j in out_id:
+                        for i in in_id:
+                            new.fix_symbolic(i, j, '0')
+                    new.mask += 1.
+
+                    for j in out_id:
+                        for i in in_id:
+                            if j > n_added_nodes - 1:
+                                new.funs[j][i] = old.funs[j - n_added_nodes][i]
+                                new.funs_avoid_singularity[j][i] = old.funs_avoid_singularity[j - n_added_nodes][i]
+                                new.funs_sympy[j][i] = old.funs_sympy[j - n_added_nodes][i]
+                                new.funs_name[j][i] = old.funs_name[j - n_added_nodes][i]
+                                new.affine.data[j][i] = old.affine.data[j - n_added_nodes][i]
+
+                    self.symbolic_fun[l] = new
+                    self.act_fun[l] = KANLayer(in_dim, out_dim + n_added_nodes, num=self.grid, k=self.k)
+                    self.act_fun[l].mask *= 0.
+
+                    self.node_scale[l].data = torch.cat(
+                        [torch.ones(n_added_nodes, device=self.device), self.node_scale[l].data])
+                    self.node_bias[l].data = torch.cat(
+                        [torch.zeros(n_added_nodes, device=self.device), self.node_bias[l].data])
+                    self.subnode_scale[l].data = torch.cat(
+                        [torch.ones(n_added_nodes, device=self.device), self.subnode_scale[l].data])
+                    self.subnode_bias[l].data = torch.cat(
+                        [torch.zeros(n_added_nodes, device=self.device), self.subnode_bias[l].data])
+
+                if added_dim == 'in':
+                    new = Symbolic_KANLayer(in_dim + n_added_nodes, out_dim)
+                    old = self.symbolic_fun[l]
+                    in_id = np.arange(in_dim + n_added_nodes)
+                    out_id = np.arange(out_dim)
+
+                    for j in out_id:
+                        for i in in_id:
+                            new.fix_symbolic(i, j, '0')
+                    new.mask += 1.
+
+                    for j in out_id:
+                        for i in in_id:
+                            if i > n_added_nodes - 1:
+                                new.funs[j][i] = old.funs[j][i - n_added_nodes]
+                                new.funs_avoid_singularity[j][i] = old.funs_avoid_singularity[j][i - n_added_nodes]
+                                new.funs_sympy[j][i] = old.funs_sympy[j][i - n_added_nodes]
+                                new.funs_name[j][i] = old.funs_name[j][i - n_added_nodes]
+                                new.affine.data[j][i] = old.affine.data[j][i - n_added_nodes]
+
+                    self.symbolic_fun[l] = new
+                    self.act_fun[l] = KANLayer(in_dim + n_added_nodes, out_dim, num=self.grid, k=self.k)
+                    self.act_fun[l].mask *= 0.
+
+
+            else:
+
+                if isinstance(mult_arity, int):
+                    mult_arity = [mult_arity] * n_added_nodes
+
+                if added_dim == 'out':
+                    n_added_subnodes = np.sum(mult_arity)
+                    new = Symbolic_KANLayer(in_dim, out_dim + n_added_subnodes)
+                    old = self.symbolic_fun[l]
+                    in_id = np.arange(in_dim)
+                    out_id = np.arange(out_dim + n_added_nodes)
+
+                    for j in out_id:
+                        for i in in_id:
+                            new.fix_symbolic(i, j, '0')
+                    new.mask += 1.
+
+                    for j in out_id:
+                        for i in in_id:
+                            if j < out_dim:
+                                new.funs[j][i] = old.funs[j][i]
+                                new.funs_avoid_singularity[j][i] = old.funs_avoid_singularity[j][i]
+                                new.funs_sympy[j][i] = old.funs_sympy[j][i]
+                                new.funs_name[j][i] = old.funs_name[j][i]
+                                new.affine.data[j][i] = old.affine.data[j][i]
+
+                    self.symbolic_fun[l] = new
+                    self.act_fun[l] = KANLayer(in_dim, out_dim + n_added_subnodes, num=self.grid, k=self.k)
+                    self.act_fun[l].mask *= 0.
+
+                    self.node_scale[l].data = torch.cat(
+                        [self.node_scale[l].data, torch.ones(n_added_nodes, device=self.device)])
+                    self.node_bias[l].data = torch.cat(
+                        [self.node_bias[l].data, torch.zeros(n_added_nodes, device=self.device)])
+                    self.subnode_scale[l].data = torch.cat(
+                        [self.subnode_scale[l].data, torch.ones(n_added_subnodes, device=self.device)])
+                    self.subnode_bias[l].data = torch.cat(
+                        [self.subnode_bias[l].data, torch.zeros(n_added_subnodes, device=self.device)])
+
+                if added_dim == 'in':
+                    new = Symbolic_KANLayer(in_dim + n_added_nodes, out_dim)
+                    old = self.symbolic_fun[l]
+                    in_id = np.arange(in_dim + n_added_nodes)
+                    out_id = np.arange(out_dim)
+
+                    for j in out_id:
+                        for i in in_id:
+                            new.fix_symbolic(i, j, '0')
+                    new.mask += 1.
+
+                    for j in out_id:
+                        for i in in_id:
+                            if i < in_dim:
+                                new.funs[j][i] = old.funs[j][i]
+                                new.funs_avoid_singularity[j][i] = old.funs_avoid_singularity[j][i]
+                                new.funs_sympy[j][i] = old.funs_sympy[j][i]
+                                new.funs_name[j][i] = old.funs_name[j][i]
+                                new.affine.data[j][i] = old.affine.data[j][i]
+
+                    self.symbolic_fun[l] = new
+                    self.act_fun[l] = KANLayer(in_dim + n_added_nodes, out_dim, num=self.grid, k=self.k)
+                    self.act_fun[l].mask *= 0.
+
+        _expand(layer_id - 1, n_added_nodes, sum_bool, mult_arity, added_dim='out')
+        _expand(layer_id, n_added_nodes, sum_bool, mult_arity, added_dim='in')
+        if sum_bool:
+            self.width[layer_id][0] += n_added_nodes
+        else:
+            if isinstance(mult_arity, int):
+                mult_arity = [mult_arity] * n_added_nodes
+
+            self.width[layer_id][1] += n_added_nodes
+            self.mult_arity[layer_id] += mult_arity
+
+    def perturb(self, mag=1.0, mode='non-intrusive'):
+        """
+        preturb a network. For usage, please refer to tutorials interp_3_KAN_compiler.ipynb.
+
+        Args:
+        -----
+            mag : float
+                perturbation magnitude
+            mode : str
+                pertubatation mode, choices = {'non-intrusive', 'all', 'minimal'}
+
+        Returns:
+        --------
+            None
+        """
+        perturb_bool = {}
+
+        if mode == 'all':
+            perturb_bool['aa_a'] = True
+            perturb_bool['aa_i'] = True
+            perturb_bool['ai'] = True
+            perturb_bool['ia'] = True
+            perturb_bool['ii'] = True
+        elif mode == 'non-intrusive':
+            perturb_bool['aa_a'] = False
+            perturb_bool['aa_i'] = False
+            perturb_bool['ai'] = True
+            perturb_bool['ia'] = False
+            perturb_bool['ii'] = True
+        elif mode == 'minimal':
+            perturb_bool['aa_a'] = True
+            perturb_bool['aa_i'] = False
+            perturb_bool['ai'] = False
+            perturb_bool['ia'] = False
+            perturb_bool['ii'] = False
+        else:
+            raise Exception('mode not recognized, valid modes are \'all\', \'non-intrusive\', \'minimal\'.')
+
+        for l in range(self.depth):
+            funs_name = self.symbolic_fun[l].funs_name
+            for j in range(self.width_out[l + 1]):
+                for i in range(self.width_in[l]):
+                    out_array = list(np.array(self.symbolic_fun[l].funs_name)[j])
+                    in_array = list(np.array(self.symbolic_fun[l].funs_name)[:, i])
+                    out_active = len([i for i, x in enumerate(out_array) if x != "0"]) > 0
+                    in_active = len([i for i, x in enumerate(in_array) if x != "0"]) > 0
+                    dic = {True: 'a', False: 'i'}
+                    edge_type = dic[in_active] + dic[out_active]
+
+                    if l < self.depth - 1 or mode != 'non-intrusive':
+
+                        if edge_type == 'aa':
+                            if self.symbolic_fun[l].funs_name[j][i] == '0':
+                                edge_type += '_i'
+                            else:
+                                edge_type += '_a'
+
+                        if perturb_bool[edge_type]:
+                            self.act_fun[l].mask.data[i][j] = mag
+
+                    if l == self.depth - 1 and mode == 'non-intrusive':
+                        self.act_fun[l].mask.data[i][j] = torch.tensor(1.)
+                        self.act_fun[l].scale_base.data[i][j] = torch.tensor(0.)
+                        self.act_fun[l].scale_sp.data[i][j] = torch.tensor(0.)
+
+        self.get_act(self.cache_data)
+
+        self.log_history('perturb')
+
+    def module(self, start_layer, chain):
+        """
+        specify network modules
+
+        Args:
+        -----
+            start_layer : int
+                the earliest layer of the module
+            chain : str
+                specify neurons in the module
+
+        Returns:
+        --------
+            None
+        """
+        #chain = '[-1]->[-1,-2]->[-1]->[-1]'
+        groups = chain.split('->')
+        n_total_layers = len(groups) // 2
+        #start_layer = 0
+
+        for l in range(n_total_layers):
+            current_layer = cl = start_layer + l
+            id_in = [int(i) for i in groups[2 * l][1:-1].split(',')]
+            id_out = [int(i) for i in groups[2 * l + 1][1:-1].split(',')]
+
+            in_dim = self.width_in[cl]
+            out_dim = self.width_out[cl + 1]
+            id_in_other = list(set(range(in_dim)) - set(id_in))
+            id_out_other = list(set(range(out_dim)) - set(id_out))
+            self.act_fun[cl].mask.data[np.ix_(id_in_other, id_out)] = 0.
+            self.act_fun[cl].mask.data[np.ix_(id_in, id_out_other)] = 0.
+            self.symbolic_fun[cl].mask.data[np.ix_(id_out, id_in_other)] = 0.
+            self.symbolic_fun[cl].mask.data[np.ix_(id_out_other, id_in)] = 0.
+
+        self.log_history('module')
+
+    def tree(self, x=None, in_var=None, style='tree', sym_th=1e-3, sep_th=1e-1, skip_sep_test=False, verbose=False):
+        """
+        turn KAN into a tree
+        """
+        if x == None:
+            x = self.cache_data
+        plot_tree(self, x, in_var=in_var, style=style, sym_th=sym_th, sep_th=sep_th, skip_sep_test=skip_sep_test,
+                  verbose=verbose)
+
+    def speed(self, compile=False):
+        """
+        turn on KAN's speed mode
+        """
+        self.symbolic_enabled = False
+        self.save_act = False
+        self.auto_save = False
+        if compile == True:
+            return torch.compile(self)
+        else:
+            return self
+
+    def get_act(self, x=None):
+        """
+        collect intermidate activations
+        """
+        if isinstance(x, dict):
+            x = x['train_input']
+        if x is None:
+            if self.cache_data != None:
+                x = self.cache_data
+            else:
+                raise Exception("missing input data x")
+        save_act = self.save_act
+        self.save_act = True
+        self.forward(x)
+        self.save_act = save_act
+
+    def get_fun(self, l, i, j):
+        """
+        get function (l,i,j)
+        """
+        inputs = self.spline_preacts[l][:, j, i].cpu().detach().numpy()
+        outputs = self.spline_postacts[l][:, j, i].cpu().detach().numpy()
+        # they are not ordered yet
+        rank = np.argsort(inputs)
+        inputs = inputs[rank]
+        outputs = outputs[rank]
+        plt.figure(figsize=(3, 3))
+        plt.plot(inputs, outputs, marker="o")
+        return inputs, outputs
+
+    def history(self, k='all'):
+        """
+        get history
+        """
+        with open(self.ckpt_path + '/history.txt', 'r') as f:
+            data = f.readlines()
+            n_line = len(data)
+            if k == 'all':
+                k = n_line
+
+            data = data[-k:]
+            for line in data:
+                print(line[:-1])
+
+    @property
+    def n_edge(self):
+        """
+        the number of active edges
+        """
+        depth = len(self.act_fun)
+        complexity = 0
+        for l in range(depth):
+            complexity += torch.sum(self.act_fun[l].mask > 0.)
+        return complexity.item()
+
+    def evaluate(self, dataset):
+        evaluation = {'test_loss': torch.sqrt(
+            torch.mean((self.forward(dataset['test_input']) - dataset['test_label']) ** 2)).item(),
+                      'n_edge': self.n_edge, 'n_grid': self.grid}
+        # add other metrics (maybe accuracy)
+        return evaluation
+
+    def swap(self, l, i1, i2, log_history=True):
+
+        self.act_fun[l - 1].swap(i1, i2, mode='out')
+        self.symbolic_fun[l - 1].swap(i1, i2, mode='out')
+        self.act_fun[l].swap(i1, i2, mode='in')
+        self.symbolic_fun[l].swap(i1, i2, mode='in')
+
+        def swap_(data, i1, i2):
+            data[i1], data[i2] = data[i2], data[i1]
+
+        swap_(self.node_scale[l - 1].data, i1, i2)
+        swap_(self.node_bias[l - 1].data, i1, i2)
+        swap_(self.subnode_scale[l - 1].data, i1, i2)
+        swap_(self.subnode_bias[l - 1].data, i1, i2)
+
+        if log_history:
+            self.log_history('swap')
+
+    @property
+    def connection_cost(self):
+
+        cc = 0.
+        for t in self.edge_scores:
+            def get_coordinate(n):
+                return torch.linspace(0, 1, steps=n + 1, device=self.device)[:n] + 1 / (2 * n)
+
+            in_dim = t.shape[0]
+            x_in = get_coordinate(in_dim)
+
+            out_dim = t.shape[1]
+            x_out = get_coordinate(out_dim)
+
+            dist = torch.abs(x_in[:, None] - x_out[None, :])
+            cc += torch.sum(dist * t)
+
+        return cc
+
+    def auto_swap_l(self, l):
+
+        num = self.width_in[1]
+        for i in range(num):
+            ccs = []
+            for j in range(num):
+                self.swap(l, i, j, log_history=False)
+                self.get_act()
+                self.attribute()
+                cc = self.connection_cost.detach().clone()
+                ccs.append(cc)
+                self.swap(l, i, j, log_history=False)
+            j = torch.argmin(torch.tensor(ccs))
+            self.swap(l, i, j, log_history=False)
+
+    def auto_swap(self):
+        """
+        automatically swap neurons such as connection costs are minimized
+        """
+        depth = self.depth
+        for l in range(1, depth):
+            self.auto_swap_l(l)
+
+        self.log_history('auto_swap')
+
+
+KAN = MultKAN