API Reference

tnlearn.VecSymRegressor	Symbolic Regression with Regressor Class
tnlearn.MLPClassifier	MLPClassifier Class Implementation
tnlearn.MLPRegressor	MLPRegressor Class Implementation
tnlearn.LLMSymRegressor	LLM-based Symbolic Regression Class
tnlearn.RLRegressor	Reinforcement Learning Regressor Class

tnlearn.VecSymRegressor

• Description:

  tnlearn.VecSymRegressor implements the symbolic regression algorithm through the Regressor class, enabling the evolution of mathematical expressions to fit given data. The class provides methods for generating random expressions, evaluating their fitness, and evolving expressions through mutation and crossover. It aims to find the best-fitting mathematical model for a given dataset.

  tnlearn.VecSymRegressor : class tnlearn.VecSymRegressor (random_state=100, pop_size=5000, max_generations=20, tournament_size=10, x_pct=0.7, xover_pct=0.3, save=True, operations=None)

• How to initialize your VecSymRegressor :

  random_state：Seed for random number generation (default: 100)

  pop_size：Population size for genetic algorithm (default: 5000)

  max_generations：Maximum generations for genetic algorithm (default: 20)

  tournament_size：Size of tournament selection (default: 10)

  x_pct：Probability of selecting a variable node during random program generation (default: 0.7)

  xover_pct：Crossover probability during offspring generation (default: 0.3)

  save：Flag for saving (default: False)

  False: Don’t need to save the progress

  True: If save option is enabled, open a file called “log.txt” to log the progress

  operations：Set of operations to be used in program generation (default: None)

  None: Use the default mathematical operations for the algorithm, like the following code:

  (
  {"func": operator.add, "arg_count": 2, "format_str": "({} + {})"},
  {"func": operator.sub, "arg_count": 2, "format_str": "({} - {})"},
  {"func": operator.mul, "arg_count": 2, "format_str": "({} * {})"},
  {"func": operator.neg, "arg_count": 1, "format_str": "-({})"},
  )
  

  You can also define operations in a similar way.

• Here is an example of using VecSymRegressor quickly:

  from tnlearn import Regressor
  
  >> neuron = Regressor()
  >> neuron.fit(X_train, y_label)
  >> print('*' * 20)
  >> print(neuron.neuron)
  
  ********************
  6@x**2 + 1.46@x - 0.1316
  

  Note

  Format X_train and y_label as numpy arrays.

  In the above example, the shapes of X_train and y_label are (600, 10) and (600, 1) respectively.

tnlearn.MLPClassifier

• Description:

  tnlearn.MLPClassifier implements the class MLPClassifier, which extends the functionality of the base class to build and train a Multi-layer Perceptron (MLP) model. MLPClassifier is designed to allow easy customization of the neural network structure, activation functions, and loss function used during training. It incorporates device selection to leverage available GPU resources, ensuring efficient computation. The class covers essential methods for model training, evaluation, and prediction, making it a flexible tool for supervised learning tasks in PyTorch.

  tnlearn.MLPClassifier : class tnlearn.MLPClassifier (neurons=’x’, layers_list=None, activation_funcs=None, loss_function=None, optimizer_name=’adam’, random_state=1, max_iter=300, batch_size=128, valid_size=0.2, lr=0.01, visual=False, save=False, visual_interval=100, gpu=None, interval=None, scheduler=None, l1_reg=None, l2_reg=None)

• How to initialize your tnlearn.MLPClassifier :

  neurons: Neuronal expression (default: ‘x’)

  Users can pass the results of tnlearn.VecSymRegressor to neurons

  layers_list: List of neuron counts for each hidden layer (default: [50, 30, 10])

  activation_funcs: Activation functions (default: None)

  Users can choose different activation functions: 'relu', 'leakyrelu', 'sigmoid', 'tanh', 'softmax'

  loss_function: Loss function for the training process (default: None)

  Users can choose different loss functions: 'mse', 'l1', 'crossentropy', 'bce'

  optimizer_name: Name of the optimizer algorithm (default: ‘adam’)

  Users can choose different optimizers: 'adam', 'sgd', 'rmsprop', 'adamw'

  random_state: Seed for random number generators for reproducibility (default: 1)

  max_iter: Maximum number of training iterations (default: 300)

  batch_size: Number of samples per batch during training (default: 128)

  valid_size: Fraction of training data used for validation (default: 0.2)

  lr: Learning rate for the optimizer (default: 0.01)

  visual: Boolean indicating if training visualization is to be shown (default: False)

  save: Indicates if the training figure should be saved (default: False)

  visual_interval: Interval at which training visualization is updated (default: 100)

  gpu: Specifies GPU configuration for training (default: None)

  None: Not use GPU

  An Integer (e.g. 1): How many GPUs you want to use

  interval: Interval of screen output during training (default: None)

  scheduler: Learning rate scheduler (default: None)

  None: Not use any learning rate adjustment strategy

  {‘step_size’: 30, ‘gamma’: 0.2}: Use lr_sceduler.StepLR() with “step_size = 30” and “gamma = 0.2”

  l1_reg: L1 regularization term (default: None)

  None: Not use L1 regularization

  True: Use L1 regularization

  l2_reg: L2 regularization term (default: None)

  None: Not use L2 regularization

  True: Use L2 regularization

• Here is an example of using MLPClassifier quickly:

  from tnlearn import MLPClassifier
  from sklearn.datasets import make_classification
  from sklearn.model_selection import train_test_split
  from tnlearn import VecSymRegressor
  
  X, y = make_classification(n_samples=200, random_state=1)
  X_train, X_test, y_train, y_test = train_test_split(X, y, stratify=y, random_state=1)
  
  neuron = VecSymRegressor()
  
  neuron.fit(X_train, y_train)
  print('*' * 20)
  print(neuron.neuron)
  
  layers_list = [10, 10, 10]
  clf = MLPClassifier()
  
  clf.fit(X_train, y_train)
  print(f"Score: {clf.score(X_test, y_test)}")
  
  clf.save(path='my_model_dir', filename='mlp_classifier.pth')
  clf.load(path='my_model_dir', filename='mlp_classifier.pth', input_dim=X_train.shape[1], output_dim=1)
  clf.fit(X_train, y_train)
  

tnlearn.MLPRegressor

• Description:

  tnlearn.MLPRegressor embodies the MLPRegressor class, designed as an extension of BaseModel within a custom machine learning framework. It facilitates the creation and training of a multilayer perceptron (MLP) for regression tasks. Features of the class include the flexibility to define custom neural network architectures through parameters such as neuronal expression and layer structure, the utilization of various activation functions, and the incorporation of modern optimization algorithms and regularization techniques. Moreover, the class is equipped with optional GPU support for enhanced computational efficiency, as well as functionalities for training visualization, validation, and model performance evaluation. This representation of an MLP is tailored to adapt to an array of regression problems while ensuring ease of use and extensibility.

  tnlearn.MLPRegressor : class tnlearn.MLPRegressor (neurons=’x’, layers_list=None, activation_funcs=None, loss_function=None, optimizer_name=’adam’, random_state=1, max_iter=300, batch_size=128, valid_size=0.2, lr=0.01, visual=False, save=False, visual_interval=100, gpu=None, interval=None, scheduler=None, l1_reg=None, l2_reg=None)

• How to initialize your tnlearn.MLPRegressor :

  neurons: Neuronal expression (default: ‘x’)

  Users can pass the results of tnlearn.VecSymRegressor to neurons

  layers_list: List of neuron counts for each hidden layer (default: [50, 30, 10])

  activation_funcs: Activation functions (default: None)

  Users can choose different activation functions: 'relu', 'leakyrelu', 'sigmoid', 'tanh', 'softmax'

  loss_function: Loss function for the training process (default: None)

  Users can choose different loss functions: 'mse', 'l1', 'crossentropy', 'bce'

  optimizer_name: Name of the optimizer algorithm (default: ‘adam’)

  Users can choose different optimizers: 'adam', 'sgd', 'rmsprop', 'adamw'

  random_state: Seed for random number generators for reproducibility (default: 1)

  max_iter: Maximum number of training iterations (default: 300)

  batch_size: Number of samples per batch during training (default: 128)

  valid_size: Fraction of training data used for validation (default: 0.2)

  lr: Learning rate for the optimizer (default: 0.01)

  visual: Boolean indicating if training visualization is to be shown (default: False)

  save: Indicates if the training figure should be saved (default: False)

  visual_interval: Interval at which training visualization is updated (default: 100)

  gpu: Specifies GPU configuration for training (default: None)

  None: Not use GPU

  An Integer (e.g. 1): How many GPUs you want to use

  interval: Interval of screen output during training (default: None)

  scheduler: Learning rate scheduler (default: None)

  None: Not use any learning rate adjustment strategy

  {‘step_size’: 30, ‘gamma’: 0.2}: Use lr_sceduler.StepLR() with “step_size = 30” and “gamma = 0.2”

  l1_reg: L1 regularization term (default: None)

  None: Not use L1 regularization

  True: Use L1 regularization

  l2_reg: L2 regularization term (default: None)

  None: Not use L2 regularization

  True: Use L2 regularization

• Here is an example of using MLPRegressor quickly:

  from tnlearn import MLPRegressor
  from sklearn.datasets import make_classification
  from sklearn.model_selection import train_test_split
  from tnlearn import VecSymRegressor
  
  X, y = make_classification(n_samples=200, random_state=1)
  X_train, X_test, y_train, y_test = train_test_split(X, y, stratify=y, random_state=1)
  
  neuron = VecSymRegressor()
  
  neuron.fit(X_train, y_train)
  print('*' * 20)
  print(neuron.neuron)
  
  layers_list = [10, 10, 10]
  clf = MLPRegressor()
  
  clf.fit(X_train, y_train)
  print(f"Score: {clf.score(X_test, y_test)}")
  
  clf.save(path='my_model_dir', filename='mlp_regressor.pth')
  clf.load(path='my_model_dir', filename='mlp_regressor.pth', input_dim=X_train.shape[1], output_dim=1)
  clf.fit(X_train, y_train)
  

tnlearn.LLMSymRegressor

• Description:

  tnlearn.LLMSymRegressor implements an LLM-based symbolic regression agent. It combines LLM‑generated equation skeletons with numerical optimization to discover interpretable expressions. The discovered expression can be directly used as a neuron formula in MLPRegressor. The agent supports multiple LLM providers (DeepSeek, SiliconFlow, Ollama, BLT, CSTCloud) and uses an experience replay buffer to iteratively improve candidate equations.

  tnlearn.LLMSymRegressor : class tnlearn.LLMSymRegressor (llm_config=None, max_iterations=20, samples_per_iteration=8, background=None, verbose=True, max_params=10, n_restarts=5, bfgs_maxiter=200, extra_prompt=None, exp_dir=None, save=False, random_state=None)

• How to initialize your LLMSymRegressor :

  llm_config: Dictionary configuring the LLM provider (default: None)

  Must contain a model key in the format 'provider/model', e.g. 'deepseek/deepseek-chat'. Optional keys: 'api_key' (string or dict), 'base_url', 'temperature' (float, default 0.6), 'max_tokens' (int, default 1024), 'top_p' (float, default 0.3). If no api_key is given, the corresponding environment variable is used (e.g. DEEPSEEK_API_KEY, SILICONFLOW_API_KEY).

  max_iterations: Number of evolutionary iterations (rounds of LLM sampling + evaluation) (default: 20)

  samples_per_iteration: Number of candidate equations generated per iteration (default: 8)

  background: Optional physical background description to guide the LLM (default: None)

  verbose: Verbosity level. If int: 0=quiet, 1=basic progress, 2=detailed debug; if bool: True -> 1, False -> 0 (default: True)

  max_params: Maximum number of coefficients (params) allowed in the equation (default: 10)

  n_restarts: Number of multi‑start BFGS runs used during evaluation (default: 5)

  bfgs_maxiter: Maximum BFGS iterations per optimization run (default: 200)

  extra_prompt: Additional user-provided text appended to the LLM prompt after the equation template (default: None)

  exp_dir: Directory for saving experiment logs and the best equation (if save=True). If not given and save=False, no folder is created (default: None)

  save: If True, save experiment logs and the best equation to exp_dir; if False, keep everything in memory (default: False)

  random_state: Seed for random number generators for reproducibility (default: None)

• Supported LLM Providers

  The following provider strings are recognised in the model field (case‑insensitive):

  Provider	Example model	Environment variable for API key
  DeepSeek	deepseek/deepseek-chat	DEEPSEEK_API_KEY
  SiliconFlow	siliconflow/Qwen/Qwen3-8B	SILICONFLOW_API_KEY
  Ollama (local)	ollama/llama3.1:8b	(not required)
  BLT	blt/gpt-4	BLT_API_KEY
  CSTCloud	cstcloud/gpt-oss-120b	CSTCLOUD_API_KEY

• Attributes after fitting

  After calling fit(), the following attributes are available:

  • best_equation_: the best discovered equation body (a string starting with return ...).

  • best_score_: the highest negative MSE (the evaluation score, higher = better).

  • best_params_: a numpy array of optimized coefficients for params[0], params[1], ….

  • The method get_neuron_formula() returns a string with numeric coefficients, ready to be used as the neurons argument in MLPRegressor.

• Here is an example of using LLMSymRegressor quickly:

  from tnlearn import LLMSymRegressor, MLPRegressor
  from sklearn.datasets import make_regression
  from sklearn.model_selection import train_test_split
  
  X, y = make_regression(n_samples=200, random_state=1)
  X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=1)
  
  llm_config = {'model': 'deepseek/deepseek-chat'}   # set DEEPSEEK_API_KEY environment variable
  
  neuron = LLMSymRegressor(llm_config=llm_config, max_iterations=5)
  neuron.fit(X_train, y_train)
  
  clf = MLPRegressor(neurons=reg.get_neuron_formula(), layers_list=[50,30,10])
  clf.fit(X_train, y_train)
  clf.predict(X_test)
  

  Note

  The fit method automatically performs BFGS parameter optimization. The resulting formula from get_neuron_formula() contains a numeric expression (e.g. 3.0@x**2 + 2.0@x + 0.1) that can be passed directly to the neurons parameter of MLPRegressor.

tnlearn.RLRegressor

• Description:

  tnlearn.RLRegressor implements a reinforcement learning (policy gradient) agent that discovers a vectorized (homogeneous) symbolic expression of the form:

  \[f(\mathbf{x}) = b + \sum_{k} c_k \sum_{j=1}^{d} \phi_k(x_j)\]

  where \(\mathbf{x} \in \mathbb{R}^d\) is a multi‑feature input vector, \(b\) is an intercept, \(\phi_k\) are basis functions selected from a user‑defined set, and \(c_k\) are coefficients fitted via Ridge regression. The agent learns to select a subset of basis functions (up to max_terms) to maximise the validation R². The discovered expression can be directly used as a neuron formula in MLPRegressor.

  tnlearn.RLRegressor : class tnlearn.RLRegressor (basis_mode=’trigonometric’, max_terms=3, max_power=5, alpha=0.1, random_state=42, max_episodes=100, val_split=0.2, lr_rl=1e-3, gamma=0.99, hidden_dim=64, verbose=True)

• How to initialize your RLRegressor :

  basis_mode : determines the set of available basis functions (default: ‘trigonometric’)

  • 'polynomial' : only polynomials (constant, \(x\), \(x^2\), …, \(x^{\text{max_power}}\))

  • 'trigonometric' : polynomials plus \(\sin(x)\) and \(\cos(x)\)

  • 'all' : polynomials plus \(\sin(x)\), \(\cos(x)\), \(\exp(x)\), and \(\log(|x|)\) (with a small safety offset)

  max_terms : maximum number of basis functions allowed in the expression (default: 3)

  max_power : maximum exponent for polynomial terms (i.e., \(x^2\) to \(x^{\text{max_power}}\)). The linear term \(x\) and the constant are always available (default: 5)

  alpha : regularisation strength for Ridge regression (default: 0.1)

  random_state : random seed for reproducibility (default: 42)

  max_episodes : number of training episodes (default: 100)

  val_split : fraction of training data used as validation for reward computation (default: 0.2)

  lr_rl : learning rate for the policy network (default: 1e-3)

  gamma : discount factor for reward calculation (default: 0.99)

  hidden_dim : number of neurons in the policy network’s hidden layers (default: 64)

  verbose : if True, print progress updates during training (default: True)

• Supported basis functions (depending on basis_mode)

  • Polynomial terms (always available): constant (implicit via intercept), \(x\), \(x^2\), …, \(x^{\text{max_power}}\)

  • Trigonometric terms (if basis_mode is 'trigonometric' or 'all'): \(\sin(x)\), \(\cos(x)\)

  • Exponential/Logarithmic (only in 'all' mode): \(\exp(x)\), \(\log(|x| + 1e-8)\)

• Attributes after fitting

  • best_expr : the best discovered symbolic expression (with numeric coefficients and intercept), e.g., "0.01 + 2.01@(x**2) + 2.99@x + 0.48@(sin(x))".

  • best_score : the best validation R² score achieved.

  • neuron : alias for best_expr, compatible with the neurons parameter of MLPRegressor.

• Here is an example of using RLRegressor quickly:

  from tnlearn import RLRegressor, MLPRegressor
  import numpy as np
  from sklearn.model_selection import train_test_split
  
  # Generate multi‑feature data with a homogeneous pattern
  np.random.seed(42)
  X = np.random.uniform(-2, 2, (500, 3))
  y = 2.0 * np.sum(X**2, axis=1) + 3.0 * np.sum(X, axis=1) + 0.5 * np.sum(np.sin(X), axis=1)
  y += 0.1 * np.random.randn(500)
  
  X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3)
  
  rl = RLRegressor(basis_mode='trigonometric', max_terms=4, max_episodes=150, verbose=True)
  rl.fit(X_train, y_train)
  
  expr = rl.get_neuron()
  mlp = MLPRegressor(neurons=expr, layers_list=[30,20])
  mlp.fit(X_train, y_train)
  mlp.predict(X_test)
  

  Warning

  When basis_mode='all', terms like exp(x) and log(x) may cause numerical instability for large inputs. Ensure your data is appropriately scaled.

tnlearn.modules – Task‑based Neural Modules

The tnlearn.modules subpackage provides a family of neural network layers that extend standard PyTorch layers (linear, convolutional, and recurrent) by allowing the user to specify a symbolic expression as the neuron’s aggregation function. This enables the construction of task‑specific neurons directly within popular architectures.

All modules follow a unified mathematical framework: the input (or hidden state) is transformed by a set of user‑defined basis functions \(f_i\), and the results are linearly combined with learnable weights \(W_i\):

\[\text{output} = \sum_{i} W_i \cdot f_i(\text{input}) + b\]

where \(\cdot\) denotes the appropriate linear operation (matrix multiplication, convolution, or transposed convolution). Each basis function has its own learnable weight, providing both flexibility and interpretability.

Example of symbolic expression interpretation

Suppose you set symbolic_expression = '2@x + 4@x**3 + torch.sin(x)'. This indicates the use of three basis functions:

• \(f_1(x) = x\)

• \(f_2(x) = x^3\)

• \(f_3(x) = \sin(x)\)

The layer will learn three weight matrices \(W_1, W_2, W_3\) (or convolution kernels) and compute:

\[\text{output} = W_1 \cdot x + W_2 \cdot x^3 + W_3 \cdot \sin(x) + b\]

where \(x\) is the input tensor. All operations (\(x\), \(x^3\), \(\sin(x)\)) are element‑wise over the input features (or spatial dimensions for convolutions). The coefficients 2@ and 4@ are ignored because the layer learns its own weights; they are only for parsing convenience.

TNLinear

tnlearn.modules.TNLinear : class TNLinear (in_features, out_features, symbolic_expression=’x’, bias=True, device=None, dtype=None)

• Description:

  A custom fully‑connected layer with task‑based neurons, where the neuron’s aggregation function is defined by a symbolic expression. The input is first transformed by each basis function in the expression, and the results are linearly combined with learnable weights:

  \[y = \sum_{i} W_i \, f_i(x) + b\]

  The coefficients (e.g., 3@ and 4@ in '3@x + 4@x**3 + torch.sin(x)') are ignored; the layer learns its own weight for each basis function. When symbolic_expression='x', this layer reduces to a standard affine transformation (equivalent to nn.Linear).

  This module supports TensorFloat32.

• Parameters

  in_features : size of each input sample (default: required)

  out_features : size of each output sample (default: required)

  symbolic_expression : string defining the basis functions, e.g., 'x + torch.sin(x) + x**3' or '3@x + 4@x**3 + torch.sin(x)'. The variable x represents the input. The layer learns its own weights for each basis function, so numeric coefficients are ignored. (default: ‘x’)

  bias : If set to False, the layer will not learn an additive bias. (default: True)

  device : the desired device of the parameters (default: None)

  dtype : the desired dtype of the parameters (default: None)

• Shape

  • Input: \((*, H_{in})\) where \(H_{in} = \text{in_features}\) and \(*\) means any number of dimensions.

  • Output: \((*, H_{out})\) where all but the last dimension are the same shape as the input and \(H_{out} = \text{out_features}\).

• Variables

  • weight : the learnable weights of the module of shape \((\text{out_features}, \text{in_features})\) per basis function. Initialized from \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) where \(k = \frac{1}{\text{in_features}}\).

  • bias : the learnable bias of the module of shape \((\text{out_features})\). If bias is True, initialized from \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) where \(k = \frac{1}{\text{in_features}}\).

• Examples

  from tnlearn import TNLinear
  import torch
  
  # Example 1: using multiple basis functions
  m1 = TNLinear(20, 30, symbolic_expression='x + torch.sin(x)')
  input1 = torch.randn(128, 20)
  output1 = m1(input1)
  
  # Example 2: using a symbolic expression with coefficients (ignored)
  m2 = TNLinear(20, 30, symbolic_expression='3@x + 4@x**3 + torch.sin(x)')
  input2 = torch.randn(128, 20)
  output2 = m2(input2)
  print(output2.size())   # torch.Size([128, 30])
  

TNConv1d

tnlearn.modules.TNConv1d : class TNConv1d (in_channels, out_channels, kernel_size, stride=1, padding=0, symbolic_expression=’x’, groups=1, dilation=1, padding_mode=’zeros’, bias=True, device=None, dtype=None)

• Description:

  Applies a 1D convolution over an input signal composed of several input planes, where the input is first transformed by a set of basis functions defined by symbolic_expression:

  \[\text{out}(N_i, C_{\text{out}_j}) = \text{bias}(C_{\text{out}_j}) + \sum_{k=0}^{C_{\text{in}}-1} \sum_{m} W_m \star f_m(\text{input}(N_i, k))\]

  where \(\star\) is the valid cross-correlation operator, and \(f_m\) are the basis functions.

  All parameters are identical to those of torch.nn.Conv1d, with the addition of symbolic_expression.

• Parameters

  in_channels : number of channels in the input image (default: required)

  out_channels : number of channels produced by the convolution (default: required)

  kernel_size : size of the convolving kernel (default: required)

  stride : stride of the convolution (default: 1)

  padding : padding added to both sides of the input (default: 0)

  symbolic_expression : basis functions to apply element‑wise to the input (same syntax as TNLinear) (default: ‘x’)

  groups : number of blocked connections from input channels to output channels (default: 1)

  dilation : spacing between kernel elements (default: 1)

  padding_mode : 'zeros', 'reflect', 'replicate' or 'circular' (default: ‘zeros’)

  bias : if True, adds a learnable bias to the output (default: True)

  device : the desired device of the parameters (default: None)

  dtype : the desired dtype of the parameters (default: None)

• Shape

  • Input: \((N, C_{in}, L_{in})\)

  • Output: \((N, C_{out}, L_{out})\) where

    \[\begin{aligned} L_{out} = \Big\lfloor & \frac{L_{in} + 2 \times \text{padding} - \text{dilation} \times (\text{kernel_size} - 1) - 1}{\text{stride}} \Big\rfloor + 1 \end{aligned}\]

• Example

  from tnlearn import TNConv1d
  import torch
  
  m = TNConv1d(16, 33, 3, stride=2, symbolic_expression='x + torch.sin(x)')
  input = torch.randn(20, 16, 50)
  output = m(input)
  

TNConv2d

tnlearn.modules.TNConv2d : class TNConv2d (in_channels, out_channels, kernel_size, stride=1, padding=0, symbolic_expression=’x’, groups=1, dilation=1, padding_mode=’zeros’, bias=True, device=None, dtype=None)

• Description:

  Applies a 2D convolution over an input signal composed of several input planes, where the input is first transformed by a set of basis functions defined by symbolic_expression:

  \[\text{out}(N_i, C_{\text{out}_j}) = \text{bias}(C_{\text{out}_j}) + \sum_{k=0}^{C_{\text{in}}-1} \sum_{m} W_m \star f_m(\text{input}(N_i, k))\]

  where \(\star\) is the valid 2D cross-correlation operator, and \(f_m\) are the basis functions.

  All parameters are identical to those of torch.nn.Conv2d, with the addition of symbolic_expression.

• Parameters

  in_channels : number of channels in the input image (default: required)

  out_channels : number of channels produced by the convolution (default: required)

  kernel_size : size of the convolving kernel (default: required)

  stride : stride of the convolution (default: 1)

  padding : padding added to all four sides of the input (default: 0)

  symbolic_expression : basis functions to apply element‑wise to the input (same syntax as TNLinear) (default: ‘x’)

  groups : number of blocked connections from input channels to output channels (default: 1)

  dilation : spacing between kernel elements (default: 1)

  padding_mode : 'zeros', 'reflect', 'replicate' or 'circular' (default: ‘zeros’)

  bias : if True, adds a learnable bias to the output (default: True)

  device : the desired device of the parameters (default: None)

  dtype : the desired dtype of the parameters (default: None)

• Shape

  • Input: \((N, C_{in}, H_{in}, W_{in})\)

  • Output: \((N, C_{out}, H_{out}, W_{out})\) where

    \[\begin{aligned} H_{out} = \Big\lfloor & \frac{H_{in} + 2 \times \text{padding}[0] - \text{dilation}[0] \times (\text{kernel_size}[0] - 1) - 1}{\text{stride}[0]} \Big\rfloor + 1 \end{aligned}\]
    \[\begin{aligned} W_{out} = \Big\lfloor & \frac{W_{in} + 2 \times \text{padding}[1] - \text{dilation}[1] \times (\text{kernel_size}[1] - 1) - 1}{\text{stride}[1]} \Big\rfloor + 1 \end{aligned}\]

• Example

  from tnlearn import TNConv2d
  import torch
  
  m = TNConv2d(16, 33, 3, stride=2, symbolic_expression='x + 0.5@x**2')
  input = torch.randn(20, 16, 50, 100)
  output = m(input)
  

TNConv3d

tnlearn.modules.TNConv3d : class TNConv3d (in_channels, out_channels, kernel_size, stride=1, padding=0, symbolic_expression=’x’, groups=1, dilation=1, padding_mode=’zeros’, bias=True, device=None, dtype=None)

• Description:

  Applies a 3D convolution over an input signal composed of several input planes, where the input is first transformed by a set of basis functions defined by symbolic_expression:

  \[\text{out}(N_i, C_{\text{out}_j}) = \text{bias}(C_{\text{out}_j}) + \sum_{k=0}^{C_{\text{in}}-1} \sum_{m} W_m \star f_m(\text{input}(N_i, k))\]

  where \(\star\) is the valid 3D cross-correlation operator, and \(f_m\) are the basis functions.

  All parameters are identical to those of torch.nn.Conv3d, with the addition of symbolic_expression.

• Parameters

  in_channels : number of channels in the input image (default: required)

  out_channels : number of channels produced by the convolution (default: required)

  kernel_size : size of the convolving kernel (default: required)

  stride : stride of the convolution (default: 1)

  padding : padding added to all sides of the input (default: 0)

  symbolic_expression : basis functions to apply element‑wise to the input (same syntax as TNLinear) (default: ‘x’)

  groups : number of blocked connections from input channels to output channels (default: 1)

  dilation : spacing between kernel elements (default: 1)

  padding_mode : 'zeros', 'reflect', 'replicate' or 'circular' (default: ‘zeros’)

  bias : if True, adds a learnable bias to the output (default: True)

  device : the desired device of the parameters (default: None)

  dtype : the desired dtype of the parameters (default: None)

• Shape

  • Input: \((N, C_{in}, D_{in}, H_{in}, W_{in})\)

  • Output: \((N, C_{out}, D_{out}, H_{out}, W_{out})\) where

    \[\begin{aligned} D_{out} = \Big\lfloor & \frac{D_{in} + 2 \times \text{padding}[0] - \text{dilation}[0] \times (\text{kernel_size}[0] - 1) - 1}{\text{stride}[0]} \Big\rfloor + 1 \end{aligned}\]
    \[\begin{aligned} H_{out} = \Big\lfloor & \frac{H_{in} + 2 \times \text{padding}[1] - \text{dilation}[1] \times (\text{kernel_size}[1] - 1) - 1}{\text{stride}[1]} \Big\rfloor + 1 \end{aligned}\]
    \[\begin{aligned} W_{out} = \Big\lfloor & \frac{W_{in} + 2 \times \text{padding}[2] - \text{dilation}[2] \times (\text{kernel_size}[2] - 1) - 1}{\text{stride}[2]} \Big\rfloor + 1 \end{aligned}\]

• Example

  from tnlearn import TNConv3d
  import torch
  
  m = TNConv3d(16, 33, 3, stride=2, symbolic_expression='x + torch.cos(x)')
  input = torch.randn(20, 16, 10, 50, 100)
  output = m(input)
  

TNConvTranspose1d

tnlearn.modules.TNConvTranspose1d : class TNConvTranspose1d (in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, symbolic_expression=’x’, groups=1, bias=True, dilation=1, padding_mode=’zeros’, device=None, dtype=None)

• Description:

  Applies a 1D transposed convolution operator over an input signal composed of several input planes, where the input is first transformed by a set of basis functions defined by symbolic_expression:

  \[\text{out}(N_i, C_{\text{out}_j}) = \text{bias}(C_{\text{out}_j}) + \sum_{k=0}^{C_{\text{in}}-1} \sum_{m} W_m \circledast f_m(\text{input}(N_i, k))\]

  where \(\circledast\) is the transposed convolution operator, and \(f_m\) are the basis functions.

  All parameters are identical to those of torch.nn.ConvTranspose1d, with the addition of symbolic_expression.

• Parameters

  in_channels : number of channels in the input image (default: required)

  out_channels : number of channels produced by the convolution (default: required)

  kernel_size : size of the convolving kernel (default: required)

  stride : stride of the convolution (default: 1)

  padding : dilation * (kernel_size - 1) - padding zero-padding will be added to both sides of the input (default: 0)

  output_padding : additional size added to one side of the output shape (default: 0)

  symbolic_expression : basis functions to apply element‑wise to the input (same syntax as TNLinear) (default: ‘x’)

  groups : number of blocked connections from input channels to output channels (default: 1)

  bias : if True, adds a learnable bias to the output (default: True)

  dilation : spacing between kernel elements (default: 1)

  padding_mode : 'zeros', 'reflect', 'replicate' or 'circular' (default: ‘zeros’)

  device : the desired device of the parameters (default: None)

  dtype : the desired dtype of the parameters (default: None)

• Shape

  • Input: \((N, C_{in}, L_{in})\)

  • Output: \((N, C_{out}, L_{out})\) where

    \[\begin{split}\begin{aligned} L_{out} = &(L_{in} - 1) \times \text{stride} - 2 \times \text{padding} \\ &+ \text{dilation} \times (\text{kernel_size} - 1) + \text{output_padding} + 1 \end{aligned}\end{split}\]

• Example

  from tnlearn import TNConvTranspose1d
  import torch
  
  m = TNConvTranspose1d(16, 33, 3, stride=2, symbolic_expression='x**2')
  input = torch.randn(20, 16, 50)
  output = m(input)
  

TNConvTranspose2d

tnlearn.modules.TNConvTranspose2d : class TNConvTranspose2d (in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, symbolic_expression=’x’, groups=1, bias=True, dilation=1, padding_mode=’zeros’, device=None, dtype=None)

• Description:

  Applies a 2D transposed convolution operator over an input signal composed of several input planes, where the input is first transformed by a set of basis functions defined by symbolic_expression:

  \[\text{out}(N_i, C_{\text{out}_j}) = \text{bias}(C_{\text{out}_j}) + \sum_{k=0}^{C_{\text{in}}-1} \sum_{m} W_m \circledast f_m(\text{input}(N_i, k))\]

  where \(\circledast\) is the transposed convolution operator, and \(f_m\) are the basis functions.

  All parameters are identical to those of torch.nn.ConvTranspose2d, with the addition of symbolic_expression.

• Parameters

  in_channels : number of channels in the input image (default: required)

  out_channels : number of channels produced by the convolution (default: required)

  kernel_size : size of the convolving kernel (default: required)

  stride : stride of the convolution (default: 1)

  padding : dilation * (kernel_size - 1) - padding zero-padding will be added to both sides of each spatial dimension (default: 0)

  output_padding : additional size added to one side of each spatial dimension of the output shape (default: 0)

  symbolic_expression : basis functions to apply element‑wise to the input (same syntax as TNLinear) (default: ‘x’)

  groups : number of blocked connections from input channels to output channels (default: 1)

  bias : if True, adds a learnable bias to the output (default: True)

  dilation : spacing between kernel elements (default: 1)

  padding_mode : 'zeros', 'reflect', 'replicate' or 'circular' (default: ‘zeros’)

  device : the desired device of the parameters (default: None)

  dtype : the desired dtype of the parameters (default: None)

• Shape

  • Input: \((N, C_{in}, H_{in}, W_{in})\)

  • Output: \((N, C_{out}, H_{out}, W_{out})\) where

    \[\begin{split}\begin{aligned} H_{out} = &(H_{in} - 1) \times \text{stride}[0] - 2 \times \text{padding}[0] \\ &+ \text{dilation}[0] \times (\text{kernel_size}[0] - 1) + \text{output_padding}[0] + 1 \end{aligned}\end{split}\]
    \[\begin{split}\begin{aligned} W_{out} = &(W_{in} - 1) \times \text{stride}[1] - 2 \times \text{padding}[1] \\ &+ \text{dilation}[1] \times (\text{kernel_size}[1] - 1) + \text{output_padding}[1] + 1 \end{aligned}\end{split}\]

• Example

  from tnlearn import TNConvTranspose2d
  import torch
  
  m = TNConvTranspose2d(16, 33, 3, stride=2, symbolic_expression='x + torch.sin(x)')
  input = torch.randn(20, 16, 50, 100)
  output = m(input)
  

TNConvTranspose3d

tnlearn.modules.TNConvTranspose3d : class TNConvTranspose3d (in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, symbolic_expression=’x’, groups=1, bias=True, dilation=1, padding_mode=’zeros’, device=None, dtype=None)

• Description:

  Applies a 3D transposed convolution operator over an input signal composed of several input planes, where the input is first transformed by a set of basis functions defined by symbolic_expression:

  \[\text{out}(N_i, C_{\text{out}_j}) = \text{bias}(C_{\text{out}_j}) + \sum_{k=0}^{C_{\text{in}}-1} \sum_{m} W_m \circledast f_m(\text{input}(N_i, k))\]

  where \(\circledast\) is the transposed convolution operator, and \(f_m\) are the basis functions.

  All parameters are identical to those of torch.nn.ConvTranspose3d, with the addition of symbolic_expression.

• Parameters

  in_channels : number of channels in the input image (default: required)

  out_channels : number of channels produced by the convolution (default: required)

  kernel_size : size of the convolving kernel (default: required)

  stride : stride of the convolution (default: 1)

  padding : dilation * (kernel_size - 1) - padding zero-padding will be added to both sides of each spatial dimension (default: 0)

  output_padding : additional size added to one side of each spatial dimension of the output shape (default: 0)

  symbolic_expression : basis functions to apply element‑wise to the input (same syntax as TNLinear) (default: ‘x’)

  groups : number of blocked connections from input channels to output channels (default: 1)

  bias : if True, adds a learnable bias to the output (default: True)

  dilation : spacing between kernel elements (default: 1)

  padding_mode : 'zeros', 'reflect', 'replicate' or 'circular' (default: ‘zeros’)

  device : the desired device of the parameters (default: None)

  dtype : the desired dtype of the parameters (default: None)

• Shape

  • Input: \((N, C_{in}, D_{in}, H_{in}, W_{in})\)

  • Output: \((N, C_{out}, D_{out}, H_{out}, W_{out})\) where

    \[\begin{split}\begin{aligned} D_{out} = &(D_{in} - 1) \times \text{stride}[0] - 2 \times \text{padding}[0] \\ &+ \text{dilation}[0] \times (\text{kernel_size}[0] - 1) + \text{output_padding}[0] + 1 \end{aligned}\end{split}\]
    \[\begin{split}\begin{aligned} H_{out} = &(H_{in} - 1) \times \text{stride}[1] - 2 \times \text{padding}[1] \\ &+ \text{dilation}[1] \times (\text{kernel_size}[1] - 1) + \text{output_padding}[1] + 1 \end{aligned}\end{split}\]
    \[\begin{split}\begin{aligned} W_{out} = &(W_{in} - 1) \times \text{stride}[2] - 2 \times \text{padding}[2] \\ &+ \text{dilation}[2] \times (\text{kernel_size}[2] - 1) + \text{output_padding}[2] + 1 \end{aligned}\end{split}\]

• Example

  from tnlearn import TNConvTranspose3d
  import torch
  
  m = TNConvTranspose3d(16, 33, 3, stride=2, symbolic_expression='x**2 + torch.cos(x)')
  input = torch.randn(20, 16, 10, 50, 100)
  output = m(input)
  

TNRNN

tnlearn.modules.TNRNN : class TNRNN (input_size, hidden_size, num_layers=1, nonlinearity=’tanh’, bias=True, batch_first=False, dropout=0.0, bidirectional=False, symbolic_expression=’x’, device=None, dtype=None)

• Description:

  Applies a multi-layer Elman RNN with \(\tanh\) or \(\text{ReLU}\) non-linearity to an input sequence, where both the input‑to‑hidden and hidden‑to‑hidden transformations are augmented with basis functions defined by symbolic_expression.

  For each element in the input sequence, each layer computes:

  \[h_t = \text{act}\!\left( \sum_{i} W_{ih}^{(i)} f_i(x_t) + \sum_{j} W_{hh}^{(j)} f_j(h_{t-1}) + b \right)\]

  where \(f_i\) and \(f_j\) are the basis functions derived from symbolic_expression, applied to the input and the previous hidden state respectively. The layer learns separate weight matrices for each basis function.

  All parameters are identical to those of torch.nn.RNN, with the addition of symbolic_expression.

• Parameters

  input_size : the number of expected features in the input \(x\) (default: required)

  hidden_size : the number of features in the hidden state \(h\) (default: required)

  num_layers : number of recurrent layers (default: 1)

  nonlinearity : the non-linearity to use. Can be either 'tanh' or 'relu'. (default: ‘tanh’)

  bias : if False, the layer does not use bias weights (default: True)

  batch_first : if True, input and output tensors are provided as (batch, seq, feature) (default: False)

  dropout : if non-zero, introduces a Dropout layer on the outputs of each RNN layer except the last layer (default: 0.0)

  bidirectional : if True, becomes a bidirectional RNN (default: False)

  symbolic_expression : symbolic expression defining the basis functions for both input and hidden transformations (same syntax as TNLinear) (default: ‘x’)

  device : the desired device of the parameters (default: None)

  dtype : the desired dtype of the parameters (default: None)

• Example

  from tnlearn import TNRNN
  import torch
  
  rnn = TNRNN(10, 20, 2, symbolic_expression='x + torch.sin(x)')
  input = torch.randn(5, 3, 10)
  h0 = torch.randn(2, 3, 20)
  output, hn = rnn(input, h0)
  

TNLSTM

tnlearn.modules.TNLSTM : class TNLSTM (input_size, hidden_size, num_layers=1, bias=True, batch_first=False, dropout=0.0, bidirectional=False, symbolic_expression=’x’, device=None, dtype=None)

• Description:

  Applies a multi-layer long short-term memory (LSTM) RNN to an input sequence, where the input‑to‑hidden and hidden‑to‑hidden transformations are augmented with basis functions defined by symbolic_expression.

  For each element in the input sequence, each layer computes the LSTM equations with basis functions replacing the linear transformations.

  All parameters are identical to those of torch.nn.LSTM, with the addition of symbolic_expression.

• Parameters

  input_size : the number of expected features in the input \(x\) (default: required)

  hidden_size : the number of features in the hidden state \(h\) (default: required)

  num_layers : number of recurrent layers (default: 1)

  bias : if False, the layer does not use bias weights (default: True)

  batch_first : if True, input and output tensors are provided as (batch, seq, feature) (default: False)

  dropout : if non-zero, introduces a Dropout layer on the outputs of each LSTM layer except the last layer (default: 0.0)

  bidirectional : if True, becomes a bidirectional LSTM (default: False)

  symbolic_expression : symbolic expression defining the basis functions (same syntax as TNLinear) (default: ‘x’)

  device : the desired device of the parameters (default: None)

  dtype : the desired dtype of the parameters (default: None)

• Example

  from tnlearn import TNLSTM
  import torch
  
  lstm = TNLSTM(10, 20, 2, symbolic_expression='x + 0.5@x**2')
  input = torch.randn(5, 3, 10)
  h0 = torch.randn(2, 3, 20)
  c0 = torch.randn(2, 3, 20)
  output, (hn, cn) = lstm(input, (h0, c0))
  

TNGRU

tnlearn.modules.TNGRU : class TNGRU (input_size, hidden_size, num_layers=1, bias=True, batch_first=False, dropout=0.0, bidirectional=False, symbolic_expression=’x’, device=None, dtype=None)

• Description:

  Applies a multi-layer gated recurrent unit (GRU) RNN to an input sequence, where the input‑to‑hidden and hidden‑to‑hidden transformations are augmented with basis functions defined by symbolic_expression.

  For each element in the input sequence, each layer computes the GRU equations with basis functions replacing the linear transformations.

  All parameters are identical to those of torch.nn.GRU, with the addition of symbolic_expression.

• Parameters

  input_size : the number of expected features in the input \(x\) (default: required)

  hidden_size : the number of features in the hidden state \(h\) (default: required)

  num_layers : number of recurrent layers (default: 1)

  bias : if False, the layer does not use bias weights (default: True)

  batch_first : if True, input and output tensors are provided as (batch, seq, feature) (default: False)

  dropout : if non-zero, introduces a Dropout layer on the outputs of each GRU layer except the last layer (default: 0.0)

  bidirectional : if True, becomes a bidirectional GRU (default: False)

  symbolic_expression : symbolic expression defining the basis functions (same syntax as TNLinear) (default: ‘x’)

  device : the desired device of the parameters (default: None)

  dtype : the desired dtype of the parameters (default: None)

• Example

  from tnlearn import TNGRU
  import torch
  
  gru = TNGRU(10, 20, 2, symbolic_expression='x + torch.sin(x)')
  input = torch.randn(5, 3, 10)
  h0 = torch.randn(2, 3, 20)
  output, hn = gru(input, h0)
  

TNRNNCell

tnlearn.modules.TNRNNCell : class TNRNNCell (input_size, hidden_size, bias=True, nonlinearity=’tanh’, symbolic_expression=’x’, device=None, dtype=None)

• Description:

  An Elman RNN cell with \(\tanh\) or \(\text{ReLU}\) non-linearity, where both input‑to‑hidden and hidden‑to‑hidden transformations are augmented with basis functions defined by symbolic_expression.

  This is the single‑step version of TNRNN.

  All parameters are identical to those of torch.nn.RNNCell, with the addition of symbolic_expression.

• Parameters

  input_size : the number of expected features in the input \(x\) (default: required)

  hidden_size : the number of features in the hidden state \(h\) (default: required)

  bias : if False, the layer does not use bias weights (default: True)

  nonlinearity : the non-linearity to use. Can be either 'tanh' or 'relu'. (default: ‘tanh’)

  symbolic_expression : symbolic expression defining the basis functions (same syntax as TNLinear) (default: ‘x’)

  device : the desired device of the parameters (default: None)

  dtype : the desired dtype of the parameters (default: None)

• Shape

  • Input: \((N, H_{in})\) or \((H_{in})\) where \(H_{in} = \text{input_size}\)

  • Hidden: \((N, H_{out})\) or \((H_{out})\) where \(H_{out} = \text{hidden_size}\) (defaults to zero if not provided)

  • Output: \((N, H_{out})\) or \((H_{out})\) tensor containing the next hidden state

• Example

  from tnlearn import TNRNNCell
  import torch
  
  rnn = TNRNNCell(10, 20, symbolic_expression='x + torch.sin(x)')
  input = torch.randn(6, 3, 10)
  hx = torch.randn(3, 20)
  output = []
  for i in range(6):
      hx = rnn(input[i], hx)
      output.append(hx)
  

TNLSTMCell

tnlearn.modules.TNLSTMCell : class TNLSTMCell (input_size, hidden_size, bias=True, symbolic_expression=’x’, device=None, dtype=None)

• Description:

  A long short-term memory (LSTM) cell, where input‑to‑hidden and hidden‑to‑hidden transformations are augmented with basis functions defined by symbolic_expression.

  This is the single‑step version of TNLSTM.

  All parameters are identical to those of torch.nn.LSTMCell, with the addition of symbolic_expression.

• Parameters

  input_size : the number of expected features in the input \(x\) (default: required)

  hidden_size : the number of features in the hidden state \(h\) (default: required)

  bias : if False, the layer does not use bias weights (default: True)

  symbolic_expression : symbolic expression defining the basis functions (same syntax as TNLinear) (default: ‘x’)

  device : the desired device of the parameters (default: None)

  dtype : the desired dtype of the parameters (default: None)

• Shape

  • Input: \((N, H_{in})\) or \((H_{in})\) where \(H_{in} = \text{input_size}\)

  • Hidden: \((N, H_{out})\) or \((H_{out})\), Cell: \((N, H_{out})\) or \((H_{out})\) where \(H_{out} = \text{hidden_size}\) (defaults to zero if not provided)

  • Output: \((h', c')\) each of shape \((N, H_{out})\) or \((H_{out})\)

• Example

  from tnlearn import TNLSTMCell
  import torch
  
  rnn = TNLSTMCell(10, 20, symbolic_expression='x + 0.1@x**2')
  input = torch.randn(6, 3, 10)
  hx = torch.randn(3, 20)
  cx = torch.randn(3, 20)
  output = []
  for i in range(6):
      hx, cx = rnn(input[i], (hx, cx))
      output.append(hx)
  

TNGRUCell

tnlearn.modules.TNGRUCell : class TNGRUCell (input_size, hidden_size, bias=True, symbolic_expression=’x’, device=None, dtype=None)

• Description:

  A gated recurrent unit (GRU) cell, where input‑to‑hidden and hidden‑to‑hidden transformations are augmented with basis functions defined by symbolic_expression.

  This is the single‑step version of TNGRU.

  All parameters are identical to those of torch.nn.GRUCell, with the addition of symbolic_expression.

• Parameters

  input_size : the number of expected features in the input \(x\) (default: required)

  hidden_size : the number of features in the hidden state \(h\) (default: required)

  bias : if False, the layer does not use bias weights (default: True)

  symbolic_expression : symbolic expression defining the basis functions (same syntax as TNLinear) (default: ‘x’)

  device : the desired device of the parameters (default: None)

  dtype : the desired dtype of the parameters (default: None)

• Shape

  • Input: \((N, H_{in})\) or \((H_{in})\) where \(H_{in} = \text{input_size}\)

  • Hidden: \((N, H_{out})\) or \((H_{out})\) where \(H_{out} = \text{hidden_size}\) (defaults to zero if not provided)

  • Output: \((N, H_{out})\) or \((H_{out})\) tensor containing the next hidden state

• Example

  from tnlearn import TNGRUCell
  import torch
  
  rnn = TNGRUCell(10, 20, symbolic_expression='x + torch.cos(x)')
  input = torch.randn(6, 3, 10)
  hx = torch.randn(3, 20)
  output = []
  for i in range(6):
      hx = rnn(input[i], hx)
      output.append(hx)
  

TNTransformer

tnlearn.modules.TNTransformer : class TNTransformer (d_model=512, nhead=8, num_encoder_layers=6, num_decoder_layers=6, dim_feedforward=2048, dropout=0.1, activation=’relu’, custom_encoder=None, custom_decoder=None, layer_norm_eps=1e-5, batch_first=False, norm_first=False, bias=True, symbolic_expression=’x’, device=None, dtype=None)

• Description:

  A basic Transformer model where the feed‑forward networks and linear projections in encoder/decoder layers are replaced with TNLinear, allowing user‑defined basis functions (e.g., 'x + torch.sin(x)') inside the transformer blocks.

  This module implements the original Transformer architecture described in the Attention Is All You Need paper. The intent of this layer is as a reference implementation for foundational understanding and thus it contains only limited features relative to newer Transformer architectures. All parameters are identical to those of torch.nn.Transformer, with the addition of symbolic_expression.

• Parameters

  d_model : the number of expected features in the encoder/decoder inputs (default: 512)

  nhead : the number of heads in the multiheadattention models (default: 8)

  num_encoder_layers : the number of sub‑encoder‑layers in the encoder (default: 6)

  num_decoder_layers : the number of sub‑decoder‑layers in the decoder (default: 6)

  dim_feedforward : the dimension of the feedforward network model (default: 2048)

  dropout : the dropout value (default: 0.1)

  activation : the activation function of encoder/decoder intermediate layer, can be a string ('relu' or 'gelu') or a unary callable (default: ‘relu’)

  custom_encoder : custom encoder (default: None)

  custom_decoder : custom decoder (default: None)

  layer_norm_eps : the eps value in layer normalization components (default: 1e-5)

  batch_first : if True, input and output tensors are provided as (batch, seq, feature) (default: False)

  norm_first : if True, layer norm is done prior to attention and feedforward operations (default: False)

  bias : if set to False, Linear and LayerNorm layers will not learn an additive bias (default: True)

  symbolic_expression : symbolic expression defining the basis functions for all TNLinear layers inside the transformer (same syntax as TNLinear) (default: ‘x’)

  device : the desired device of the parameters (default: None)

  dtype : the desired dtype of the parameters (default: None)

• Example

  from tnlearn import TNTransformer
  import torch
  
  transformer = TNTransformer(
      d_model=512, nhead=8,
      num_encoder_layers=2, num_decoder_layers=2,
      symbolic_expression='x + torch.tanh(x)'
  )
  src = torch.randn(10, 32, 512)   # (seq, batch, feature)
  tgt = torch.randn(20, 32, 512)
  out = transformer(src, tgt)
  

  Note

  This class uses TNTransformerEncoder and TNTransformerDecoder internally. When batch_first=False (default), the input and output shapes are (seq, batch, feature). When batch_first=True, shapes are (batch, seq, feature).

TNTransformerEncoder

tnlearn.modules.TNTransformerEncoder : class TNTransformerEncoder (encoder_layer, num_layers, norm=None, enable_nested_tensor=True, mask_check=True)

• Description:

  A stack of N encoder layers, where each layer uses TNTransformerEncoderLayer. This module is the TNLearn counterpart of torch.nn.TransformerEncoder.

  TransformerEncoder is a stack of N encoder layers. The intent of this layer is as a reference implementation for foundational understanding.

  Warning

  All layers in the TransformerEncoder are initialized with the same parameters. It is recommended to manually initialize the layers after creating the TransformerEncoder instance.

• Parameters

  encoder_layer : an instance of the TNTransformerEncoderLayer class (required)

  num_layers : the number of sub‑encoder‑layers in the encoder (required)

  norm : the layer normalization component (optional)

  enable_nested_tensor : if True, input will automatically convert to nested tensor (and convert back on output). This will improve the overall performance when padding rate is high (default: True)

  mask_check : whether to check the mask for left‑aligned padding (default: True)

• Example

  from tnlearn import TNTransformerEncoderLayer, TNTransformerEncoder
  import torch
  
  encoder_layer = TNTransformerEncoderLayer(
      d_model=512, nhead=8,
      symbolic_expression='x + 0.5 * torch.sin(x)'
  )
  transformer_encoder = TNTransformerEncoder(encoder_layer, num_layers=2)
  src = torch.randn(10, 32, 512)
  out = transformer_encoder(src)
  

TNTransformerDecoder

tnlearn.modules.TNTransformerDecoder : class TNTransformerDecoder (decoder_layer, num_layers, norm=None)

• Description:

  A stack of N decoder layers, where each layer uses TNTransformerDecoderLayer. This module is the TNLearn counterpart of torch.nn.TransformerDecoder.

  TransformerDecoder is a stack of N decoder layers. The intent of this layer is as a reference implementation for foundational understanding.

  Warning

  All layers in the TransformerDecoder are initialized with the same parameters. It is recommended to manually initialize the layers after creating the TransformerDecoder instance.

• Parameters

  decoder_layer : an instance of the TNTransformerDecoderLayer class (required)

  num_layers : the number of sub‑decoder‑layers in the decoder (required)

  norm : the layer normalization component (optional)

• Example

  from tnlearn import TNTransformerDecoderLayer, TNTransformerDecoder
  import torch
  
  decoder_layer = TNTransformerDecoderLayer(
      d_model=512, nhead=8,
      symbolic_expression='x * torch.sigmoid(x)'
  )
  transformer_decoder = TNTransformerDecoder(decoder_layer, num_layers=2)
  tgt = torch.randn(20, 32, 512)
  memory = torch.randn(10, 32, 512)
  out = transformer_decoder(tgt, memory)
  

TNTransformerEncoderLayer

tnlearn.modules.TNTransformerEncoderLayer : class TNTransformerEncoderLayer (d_model, nhead, dim_feedforward=2048, dropout=0.1, activation=’relu’, layer_norm_eps=1e-5, batch_first=False, norm_first=False, bias=True, symbolic_expression=’x’, device=None, dtype=None)

• Description:

  TransformerEncoderLayer is made up of self‑attention and a feed‑forward network (FFN). The FFN uses TNLinear to support custom basis functions defined by symbolic_expression, while the self‑attention projections (query/key/value and output) remain as standard nn.Linear layers.

  This layer implements the original Transformer encoder block described in the Attention Is All You Need paper. All parameters are identical to those of torch.nn.TransformerEncoderLayer, with the addition of symbolic_expression.

• Parameters

  d_model : the number of expected features in the input (required)

  nhead : the number of heads in the multiheadattention models (required)

  dim_feedforward : the dimension of the feedforward network model (default: 2048)

  dropout : the dropout value (default: 0.1)

  activation : the activation function of the intermediate layer, can be a string ('relu' or 'gelu') or a unary callable (default: ‘relu’)

  layer_norm_eps : the eps value in layer normalization components (default: 1e-5)

  batch_first : if True, input and output tensors are provided as (batch, seq, feature) (default: False)

  norm_first : if True, layer norm is done prior to self attention and feedforward operations (default: False)

  bias : if set to False, Linear and LayerNorm layers will not learn an additive bias (default: True)

  symbolic_expression : symbolic expression defining the basis functions for the TNLinear layers in the FFN (same syntax as TNLinear) (default: ‘x’)

  device : the desired device of the parameters (default: None)

  dtype : the desired dtype of the parameters (default: None)

• Shape

  • Input: \((S, N, E)\) or \((N, S, E)\) if batch_first=True, where \(S\) is the source sequence length, \(N\) is the batch size, and \(E\) is the feature dimension.

  • Output: same shape as input.

• Example

  from tnlearn import TNTransformerEncoderLayer
  import torch
  
  encoder_layer = TNTransformerEncoderLayer(
      d_model=512, nhead=8,
      symbolic_expression='x + torch.sin(x)'
  )
  src = torch.randn(10, 32, 512)   # (seq, batch, feature)
  out = encoder_layer(src)
  

TNTransformerDecoderLayer

tnlearn.modules.TNTransformerDecoderLayer : class TNTransformerDecoderLayer (d_model, nhead, dim_feedforward=2048, dropout=0.1, activation=’relu’, layer_norm_eps=1e-5, batch_first=False, norm_first=False, bias=True, symbolic_expression=’x’, device=None, dtype=None)

• Description:

  TransformerDecoderLayer is made up of self‑attention, cross‑attention, and a feed‑forward network (FFN). The FFN uses TNLinear to support custom basis functions defined by symbolic_expression, while the attention projections (query/key/value and output) remain as standard nn.Linear layers.

  This layer implements the original Transformer decoder block described in the Attention Is All You Need paper. All parameters are identical to those of torch.nn.TransformerDecoderLayer, with the addition of symbolic_expression.

• Parameters

  d_model : the number of expected features in the input (required)

  nhead : the number of heads in the multiheadattention models (required)

  dim_feedforward : the dimension of the feedforward network model (default: 2048)

  dropout : the dropout value (default: 0.1)

  activation : the activation function of the intermediate layer, can be a string ('relu' or 'gelu') or a unary callable (default: ‘relu’)

  layer_norm_eps : the eps value in layer normalization components (default: 1e-5)

  batch_first : if True, input and output tensors are provided as (batch, seq, feature) (default: False)

  norm_first : if True, layer norm is done prior to self attention, cross attention and feedforward operations (default: False)

  bias : if set to False, Linear and LayerNorm layers will not learn an additive bias (default: True)

  symbolic_expression : symbolic expression defining the basis functions for the TNLinear layers in the FFN (same syntax as TNLinear) (default: ‘x’)

  device : the desired device of the parameters (default: None)

  dtype : the desired dtype of the parameters (default: None)

• Shape

  • Input: target \((T, N, E)\), memory \((S, N, E)\) (or with batch first if enabled).

  • Output: same shape as target.

• Example

  from tnlearn import TNTransformerDecoderLayer
  import torch
  
  decoder_layer = TNTransformerDecoderLayer(
      d_model=512, nhead=8,
      symbolic_expression='x**2 + torch.cos(x)'
  )
  tgt = torch.randn(20, 32, 512)   # (tgt_seq, batch, feature)
  memory = torch.randn(10, 32, 512)
  out = decoder_layer(tgt, memory)
  

Common Notes

• Basis function syntax: The expression is a Python string that can contain any valid PyTorch operations (e.g., torch.sin, torch.cos, torch.exp, **, *, etc.). The variable representing the input is always x. Coefficients (e.g., 0.5@x**2) are ignored because the layer learns its own weights; the @ symbol is only used for parsing.

• Weight sharing: Each basis function \(f_i\) has its own weight matrix/convolution kernel (or linear projection). Consequently, the number of trainable parameters scales linearly with the number of basis functions. This applies to all modules in tnlearn.modules, including linear, convolutional, recurrent, and transformer layers.

• Transformer-specific: In transformer encoder/decoder layers, only the feed‑forward network (FFN) uses TNLinear; the attention projections (query/key/value and output) remain standard nn.Linear layers for efficiency and stability.

• Save/load: All modules support standard PyTorch serialisation via torch.save() and torch.load(), including the dynamic function compilation.

• Performance: The forward pass evaluates each basis function in a loop. For large models or many basis functions, consider using simpler expressions to reduce computational and memory overhead.

• Integration: These modules are fully compatible with PyTorch’s optimizers, data parallelism, and mixed precision training.

Note

The RNN modules internally augment the input by concatenating the outputs of all basis functions along the feature dimension before passing to a native PyTorch RNN implementation. This provides high performance while still allowing arbitrary symbolic expressions.