athena__orthogonal_nop

Module containing implementation of an Orthogonal Neural Operator layer

This module implements the Orthogonal Neural Operator (ONO) from "Improved Operator Learning by Orthogonal Attention" (Luo et al., 2024).

The ONO layer uses an orthogonal attention kernel to approximate the integral operator. It combines: 1. A learned orthogonal basis for efficient attention (k << N) 2. A spectral pathway through the orthogonal basis 3. A local affine bypass

The layer computes: $\mathbf{v} = \sigma\!\bigl( \mathbf{W}_V\,\mathbf{\Phi}\,(\mathbf{\Phi}^T\,\mathbf{u}) + \mathbf{W}\,\mathbf{u} + \mathbf{b}\bigr)$

where $\mathbf{\Phi} \in \mathbb{R}^{n_{in} \times k}$ is obtained by QR/Gram-Schmidt orthogonalisation of learnable basis weights $\mathbf{B}$ .

Parameters (learnable): - $\mathbf{R} \in \mathbb{R}^{k \times k}$ spectral mixing - $\mathbf{B} \in \mathbb{R}^{n_{in} \times k}$ basis (orthogonalised) - $\mathbf{W} \in \mathbb{R}^{n_{out} \times n_{in}}$ bypass - $\mathbf{b} \in \mathbb{R}^{n_{out}}$ bias (optional)

The spectral path is: decode( R * encode(u) ) encode = Phi^T @ u [k, batch] mix = R @ encoded [k, batch] decode = Phi @ mix [n_in, batch]

Then a linear projection to the output: W_out @ decode -> [n_out, batch] Plus bypass: W @ u -> [n_out, batch]

Uses

- diffstruc
- athena__misc_types
- coreutils
- athena__diffstruc_extd
- athena__base_layer

Interfaces

public interface orthogonal_nop_block_type

public module function layer_setup(num_outputs, num_basis, num_inputs, use_bias, activation, kernel_initialiser, bias_initialiser, verbose) result(layer)

Arguments

Type	Intent	Optional		Name
integer,	intent(in)		::	num_outputs
integer,	intent(in)		::	num_basis
integer,	intent(in),	optional	::	num_inputs
logical,	intent(in),	optional	::	use_bias
class(*),	intent(in),	optional	::	activation
class(*),	intent(in),	optional	::	kernel_initialiser
class(*),	intent(in),	optional	::	bias_initialiser
integer,	intent(in),	optional	::	verbose

Return Value type(orthogonal_nop_block_type)

Derived Types

type, public, extends(learnable_layer_type) :: orthogonal_nop_block_type

Type for an Orthogonal Neural Operator layer

Components

Type	Visibility	Attributes		Name		Initial
class(base_actv_type),	public,	allocatable	::	activation			Activation function
class(base_init_type),	public,	allocatable	::	bias_init			Initialisers for kernel and bias
character(len=14),	public		::	bias_initialiser	=	''	Initialisers for kernel and bias
integer,	public,	allocatable, dimension(:)	::	bias_shape			Shape of biases
type(graph_type),	public,	allocatable, dimension(:)	::	graph			Graph structure of input data
integer,	public		::	id			Unique identifier
logical,	public		::	inference	=	.false.	Inference mode
integer,	public		::	input_rank	=	0	Rank of input data
integer,	public,	allocatable, dimension(:)	::	input_shape			Input shape
class(base_init_type),	public,	allocatable	::	kernel_init			Initialisers for kernel and bias
character(len=14),	public		::	kernel_initialiser	=	''	Initialisers for kernel and bias
character(len=:),	public,	allocatable	::	name			Layer name
integer,	public		::	num_basis	=	0	Number of orthogonal basis functions (k)
integer,	public		::	num_inputs	=	0	Number of inputs (discretisation points)
integer,	public		::	num_outputs	=	0	Number of outputs (discretisation points)
integer,	public		::	num_params	=	0	Number of learnable parameters
class(array_type),	public,	allocatable, dimension(:,:)	::	output			Output
integer,	public		::	output_rank	=	0	Rank of output data
integer,	public,	allocatable, dimension(:)	::	output_shape			Output shape
type(array_type),	public,	allocatable, dimension(:)	::	params			Learnable parameters
character(len=20),	public		::	subtype	=	repeat(" ", 20)
character(len=4),	public		::	type	=	'base'	Layer type
logical,	public		::	use_bias	=	.false.	Layer has bias
logical,	public		::	use_graph_input	=	.false.	Use graph input
logical,	public		::	use_graph_output	=	.false.	Use graph output
integer,	public,	allocatable, dimension(:,:)	::	weight_shape			Shape of weights
type(array_type),	public,	dimension(1)	::	z			Temporary array for pre-activation values

Constructor

public module function layer_setup (num_outputs, num_basis, num_inputs, use_bias, activation, kernel_initialiser, bias_initialiser, verbose)

Finalizations Procedures

final :: finalise_ono

Type-Bound Procedures

procedure, public :: add_t_t => add_learnable	Add two layers
procedure, public, pass(this) :: build_from_onnx => build_from_onnx_base	Build layer from ONNX node and initialiser
procedure, public, pass(this) :: emit_onnx_graph_inputs => emit_onnx_graph_inputs_base	Emit graph input tensor declarations for this layer
procedure, public, pass(this) :: emit_onnx_nodes => emit_onnx_nodes_base	Emit ONNX JSON nodes for this layer (format-aware and polymorphic)
procedure, public, pass(this) :: extract_output => extract_output_base	Extract the output of the layer as a standard real array
procedure, public, pass(this) :: forward => forward_ono
procedure, public, pass(this) :: forward_eval => forward_eval_base	Forward pass of layer and return output for evaluation
procedure, public, pass(this) :: get_attributes => get_attributes_ono
procedure, public, pass(this) :: get_bases => get_bases_ono
procedure, public, pass(this) :: get_gradients	Get parameter gradients of layer
procedure, public, pass(this) :: get_num_params => get_num_params_ono
procedure, public, pass(this) :: get_orthogonality_metric
procedure, public, pass(this) :: get_params	Get learnable parameters of layer
procedure, public, pass(this) :: init => init_ono
procedure, public, pass(this) :: nullify_graph => nullify_graph_base	Nullify the forward pass data of the layer to free memory Read more…
generic, public :: operator(+) => add_t_t	Operator overloading for addition
procedure, public, pass(this) :: print => print_base	Print the layer to a file with additional information
procedure, public, pass(this) :: print_to_unit => print_to_unit_ono
procedure, public, pass(this) :: read => read_ono
procedure, public, pass(this) :: reduce => reduce_learnable	Merge another learnable layer into this one
procedure, public, pass(this) :: set_gradients	Set learnable parameters of layer
procedure, public, pass(this) :: set_graph => set_graph_base	Set the graph structure of the input data !! this is adjacency and edge weighting
procedure, public, pass(this) :: set_hyperparams => set_hyperparams_ono
procedure, public, pass(this) :: set_params	Set learnable parameters of layer
procedure, public, pass(this) :: set_rank => set_rank_base	Set the input and output ranks of the layer
procedure, public, pass(this) :: set_shape => set_shape_base	Set the input shape of the layer

Functions

public function get_attributes_ono(this) result(attributes)

Return list of ONO attributes for ONNX export

Arguments

Type	Intent	Optional	Attributes		Name
class(orthogonal_nop_block_type),	intent(in)			::	this	Instance of the ONO block

Return Value type(onnx_attribute_type), allocatable, dimension(:)

List of attributes for ONNX export

public function get_bases_ono(this) result(phi)

Orthogonalise the basis matrix B using modified Gram-Schmidt

Arguments

Type	Intent	Optional	Attributes		Name
class(orthogonal_nop_block_type),	intent(in)			::	this	Layer instance providing basis parameters

Return Value type(array_type)

Orthogonalised basis matrix packed in an array_type

public pure function get_num_params_ono(this) result(num_params)

Return the number of learnable parameters for the block

Arguments

Type	Intent	Optional	Attributes		Name
class(orthogonal_nop_block_type),	intent(in)			::	this	Layer instance

Return Value integer

Total number of learnable parameters

public function get_orthogonality_metric(this) result(metric)

Compute max(|Phi^T @ Phi - I|) as a measure of basis orthogonality

Arguments

Type	Intent	Optional	Attributes		Name
class(orthogonal_nop_block_type),	intent(in)			::	this	Layer instance providing basis parameters

Return Value real(kind=real32)

Maximum absolute deviation from orthogonality

public module function layer_setup(num_outputs, num_basis, num_inputs, use_bias, activation, kernel_initialiser, bias_initialiser, verbose) result(layer)

Arguments

Type	Intent	Optional		Name
integer,	intent(in)		::	num_outputs	Number of output features
integer,	intent(in)		::	num_basis	Number of orthogonal basis vectors
integer,	intent(in),	optional	::	num_inputs	Number of input features when known at construction time
logical,	intent(in),	optional	::	use_bias	Whether to allocate a bias term
class(*),	intent(in),	optional	::	activation	Activation function specification
class(*),	intent(in),	optional	::	kernel_initialiser	Kernel and bias initialiser specifications
class(*),	intent(in),	optional	::	bias_initialiser	Kernel and bias initialiser specifications
integer,	intent(in),	optional	::	verbose	Verbosity level

Return Value type(orthogonal_nop_block_type)

Constructed orthogonal neural operator block

public function read_orthogonal_nop_block(unit, verbose) result(layer)

Read an orthogonal neural operator block from file and return it

Arguments

Type	Intent	Optional	Attributes		Name
integer,	intent(in)			::	unit	Input unit number
integer,	intent(in),	optional		::	verbose	Verbosity level

Return Value class(base_layer_type), allocatable

Allocated base-layer instance containing the result

Subroutines

public subroutine finalise_ono(this)

Finalise the orthogonal neural operator block

Arguments

Type	Intent	Optional	Attributes		Name
type(orthogonal_nop_block_type),	intent(inout)			::	this	Layer instance to release

public subroutine forward_ono(this, input)

Forward propagation for the Orthogonal Neural Operator layer

Arguments

Type	Intent	Optional	Attributes		Name
class(orthogonal_nop_block_type),	intent(inout)			::	this	Layer instance to execute
class(array_type),	intent(in),		dimension(:,:)	::	input	Input batch tensor collection

public subroutine init_ono(this, input_shape, verbose)

Initialise parameter storage and output buffers for the block

Arguments

Type	Intent	Optional	Attributes		Name
class(orthogonal_nop_block_type),	intent(inout)			::	this	Layer instance to initialise
integer,	intent(in),		dimension(:)	::	input_shape	Input shape used to infer num_inputs
integer,	intent(in),	optional		::	verbose	Verbosity level

public subroutine print_to_unit_ono(this, unit)

Print orthogonal neural operator settings and parameters to a unit

Arguments

Type	Intent	Optional	Attributes		Name
class(orthogonal_nop_block_type),	intent(in)			::	this	Layer instance to print
integer,	intent(in)			::	unit	Output unit number

public subroutine read_ono(this, unit, verbose)

Arguments

Type	Intent	Optional		Name
class(orthogonal_nop_block_type),	intent(inout)		::	this	Layer instance to populate from file data
integer,	intent(in)		::	unit	Input unit number
integer,	intent(in),	optional	::	verbose	Verbosity level

public subroutine set_hyperparams_ono(this, num_outputs, num_basis, use_bias, activation, kernel_initialiser, bias_initialiser, verbose)

Arguments

Type	Intent	Optional	Attributes		Name
class(orthogonal_nop_block_type),	intent(inout)			::	this	Layer instance to configure
integer,	intent(in)			::	num_outputs	Number of output features
integer,	intent(in)			::	num_basis	Number of orthogonal basis vectors
logical,	intent(in)			::	use_bias	Whether to use a bias term
class(base_actv_type),	intent(in),		allocatable	::	activation	Activation function object
class(base_init_type),	intent(in),		allocatable	::	kernel_initialiser	Kernel and bias initialiser objects
class(base_init_type),	intent(in),		allocatable	::	bias_initialiser	Kernel and bias initialiser objects
integer,	intent(in),	optional		::	verbose	Verbosity level

athena__orthogonal_nop_block Module

Contents

Interfaces

Derived Types

Functions

Subroutines

Uses

Interfaces

public interface orthogonal_nop_block_type

public module function layer_setup(num_outputs, num_basis, num_inputs, use_bias, activation, kernel_initialiser, bias_initialiser, verbose) result(layer) Implementation →

Arguments

Return Value type(orthogonal_nop_block_type)

Derived Types

type, public, extends(learnable_layer_type) :: orthogonal_nop_block_type

Components

Constructor

Finalizations Procedures

Type-Bound Procedures

Functions

public function get_attributes_ono(this) result(attributes)

Arguments

Return Value type(onnx_attribute_type), allocatable, dimension(:)

public function get_bases_ono(this) result(phi)

Arguments

Return Value type(array_type)

public pure function get_num_params_ono(this) result(num_params)

Arguments

Return Value integer

public function get_orthogonality_metric(this) result(metric)

Arguments

Return Value real(kind=real32)

public module function layer_setup(num_outputs, num_basis, num_inputs, use_bias, activation, kernel_initialiser, bias_initialiser, verbose) result(layer)

Arguments

Return Value type(orthogonal_nop_block_type)

public function read_orthogonal_nop_block(unit, verbose) result(layer)

Arguments

Return Value class(base_layer_type), allocatable

Subroutines

public subroutine finalise_ono(this)

Arguments

public subroutine forward_ono(this, input)

Arguments

public subroutine init_ono(this, input_shape, verbose)

Arguments

public subroutine print_to_unit_ono(this, unit)

Arguments

public subroutine read_ono(this, unit, verbose)

Arguments

public subroutine set_hyperparams_ono(this, num_outputs, num_basis, use_bias, activation, kernel_initialiser, bias_initialiser, verbose)

Arguments

public module function layer_setup(num_outputs, num_basis, num_inputs, use_bias, activation, kernel_initialiser, bias_initialiser, verbose) result(layer)