athena__diffstruc_extd_submodule

Submodule containing autodiff operations for the Graph Neural Operator

Provides two differentiable operations:

gno_kernel_eval — evaluates the kernel MLP on each edge: $\kappa(\Delta x) = V \, \mathrm{relu}(U \Delta x + b_u) + b_v$ left_operand → edge_features [d, num_edges] right_operand → packed kernel params [Hd + H + FH + F, 1] where F = F_out * F_in output → [F_out * F_in, num_edges] per-edge kernel values
gno_aggregate — aggregates messages using per-edge kernels: $m_i = \sum_{j \in \mathcal{N}(i)} \kappa_{ij} \, h_j$ left_operand → features [F_in, num_vertices] right_operand → edge_kernels [F_out * F_in, num_edges] output → [F_out, num_vertices]

gno_aggregate stores adj_ia and adj_ja on the result for use in the backward pass. Metadata (d, H, F_in, F_out) is stored in indices of the kernel evaluation result.

Uses

- Ancestors:
- athena__diffstruc_extd

Functions

function get_partial_gno_agg_features(this, upstream_grad) result(output)

Gradient of gno_aggregate w.r.t. features (left operand)

Arguments

Type	Intent	Optional	Attributes		Name
class(array_type),	intent(inout)			::	this	Forward result node containing saved operands
type(array_type),	intent(in)			::	upstream_grad	Upstream gradient tensor

Return Value type(array_type)

Gradient tensor for node features

function get_partial_gno_agg_kernels(this, upstream_grad) result(output)

Gradient of gno_aggregate w.r.t. edge_kernels (right operand)

Arguments

Type	Intent	Optional	Attributes		Name
class(array_type),	intent(inout)			::	this	Forward result node containing saved operands
type(array_type),	intent(in)			::	upstream_grad	Upstream gradient tensor

Return Value type(array_type)

Gradient tensor for edge kernels

function get_partial_gno_kernel_coords(this, upstream_grad) result(output)

Gradient of gno_kernel_eval w.r.t. edge features (left operand)

Arguments

Type	Intent	Optional	Attributes		Name
class(array_type),	intent(inout)			::	this	Forward result node containing saved operands
type(array_type),	intent(in)			::	upstream_grad	Upstream gradient tensor

Return Value type(array_type)

Gradient tensor for coordinates

function get_partial_gno_kernel_params(this, upstream_grad) result(output)

Gradient of gno_kernel_eval w.r.t. kernel_params (right operand)

Arguments

Type	Intent	Optional	Attributes		Name
class(array_type),	intent(inout)			::	this	Forward result node containing saved operands
type(array_type),	intent(in)			::	upstream_grad	Upstream gradient tensor

Return Value type(array_type)

Gradient tensor for packed kernel parameters

function get_partial_lno_decode_poles(this, upstream_grad) result(output)

Gradient of lno_decode with respect to poles.

Arguments

Type	Intent	Optional	Attributes		Name
class(array_type),	intent(inout)			::	this	Forward result node containing saved operands
type(array_type),	intent(in)			::	upstream_grad	Upstream gradient tensor

Return Value type(array_type)

Gradient tensor for poles

function get_partial_lno_decode_spectral(this, upstream_grad) result(output)

Gradient of lno_decode with respect to spectral input.

Arguments

Type	Intent	Optional	Attributes		Name
class(array_type),	intent(inout)			::	this	Forward result node containing saved operands
type(array_type),	intent(in)			::	upstream_grad	Upstream gradient tensor

Return Value type(array_type)

Gradient tensor for spectral input

function get_partial_lno_encode_input(this, upstream_grad) result(output)

Gradient of lno_encode with respect to input.

Arguments

Type	Intent	Optional	Attributes		Name
class(array_type),	intent(inout)			::	this	Forward result node containing saved operands
type(array_type),	intent(in)			::	upstream_grad	Upstream gradient tensor

Return Value type(array_type)

Gradient tensor for input

function get_partial_lno_encode_poles(this, upstream_grad) result(output)

Gradient of lno_encode with respect to poles.

Arguments

Type	Intent	Optional	Attributes		Name
class(array_type),	intent(inout)			::	this	Forward result node containing saved operands
type(array_type),	intent(in)			::	upstream_grad	Upstream gradient tensor

Return Value type(array_type)

Gradient tensor for poles

function get_partial_ono_decode_basis(this, upstream_grad) result(output)

Gradient of ono_decode with respect to basis weights.

Arguments

Type	Intent	Optional	Attributes		Name
class(array_type),	intent(inout)			::	this	Forward result node containing saved operands
type(array_type),	intent(in)			::	upstream_grad	Upstream gradient tensor

Return Value type(array_type)

Gradient tensor for basis weights

function get_partial_ono_decode_mixed(this, upstream_grad) result(output)

Gradient of ono_decode with respect to mixed input.

Arguments

Type	Intent	Optional	Attributes		Name
class(array_type),	intent(inout)			::	this	Forward result node containing saved operands
type(array_type),	intent(in)			::	upstream_grad	Upstream gradient tensor

Return Value type(array_type)

Gradient tensor for mixed input

function get_partial_ono_encode_basis(this, upstream_grad) result(output)

Gradient of ono_encode with respect to basis weights.

Arguments

Type	Intent	Optional	Attributes		Name
class(array_type),	intent(inout)			::	this	Forward result node containing saved operands
type(array_type),	intent(in)			::	upstream_grad	Upstream gradient tensor

Return Value type(array_type)

Gradient tensor for basis weights

function get_partial_ono_encode_input(this, upstream_grad) result(output)

Gradient of ono_encode with respect to input.

Arguments

Type	Intent	Optional	Attributes		Name
class(array_type),	intent(inout)			::	this	Forward result node containing saved operands
type(array_type),	intent(in)			::	upstream_grad	Upstream gradient tensor

Return Value type(array_type)

Gradient tensor for input

Subroutines

pure subroutine get_partial_elem_scale_input_val(this, upstream_grad, output)

d(out)/d(input): upstream * scale (broadcast scale along samples)

Arguments

Type	Intent	Attributes		Name
class(array_type),	intent(in)		::	this	Forward result node containing saved operands
real(kind=real32),	intent(in),	dimension(:,:)	::	upstream_grad	Upstream gradient values
real(kind=real32),	intent(out),	dimension(:,:)	::	output	Output gradient values for input

pure subroutine get_partial_elem_scale_scale_val(this, upstream_grad, output)

d(out)/d(scale): upstream * input (element-wise, per sample)

Arguments

Type	Intent	Attributes		Name
class(array_type),	intent(in)		::	this	Forward result node containing saved operands
real(kind=real32),	intent(in),	dimension(:,:)	::	upstream_grad	Upstream gradient values
real(kind=real32),	intent(out),	dimension(:,:)	::	output	Output gradient values for scale tensor

pure subroutine get_partial_gno_agg_features_val(this, upstream_grad, output)

In-place gradient w.r.t. features

Arguments

Type	Intent	Attributes		Name
class(array_type),	intent(in)		::	this	Forward result node containing saved operands
real(kind=real32),	intent(in),	dimension(:,:)	::	upstream_grad	Upstream gradient values
real(kind=real32),	intent(out),	dimension(:,:)	::	output	Output gradient values for node features

pure subroutine get_partial_gno_agg_kernels_val(this, upstream_grad, output)

In-place gradient w.r.t. edge_kernels

Arguments

Type	Intent	Attributes		Name
class(array_type),	intent(in)		::	this	Forward result node containing saved operands
real(kind=real32),	intent(in),	dimension(:,:)	::	upstream_grad	Upstream gradient values
real(kind=real32),	intent(out),	dimension(:,:)	::	output	Output gradient values for edge kernels

pure subroutine get_partial_gno_kernel_coords_val(this, upstream_grad, output)

In-place gradient w.r.t. edge features

Arguments

Type	Intent	Attributes		Name
class(array_type),	intent(in)		::	this	Forward result node containing saved operands
real(kind=real32),	intent(in),	dimension(:,:)	::	upstream_grad	Upstream gradient values
real(kind=real32),	intent(out),	dimension(:,:)	::	output	Output gradient values for coordinates

pure subroutine get_partial_gno_kernel_params_val(this, upstream_grad, output)

In-place gradient w.r.t. packed kernel params

Arguments

Type	Intent	Attributes		Name
class(array_type),	intent(in)		::	this	Forward result node containing saved operands
real(kind=real32),	intent(in),	dimension(:,:)	::	upstream_grad	Upstream gradient values
real(kind=real32),	intent(out),	dimension(:,:)	::	output	Output gradient values for packed kernel parameters

pure subroutine get_partial_lno_decode_poles_val(this, upstream_grad, output)

dL/dmu_m per sample: output[m,s] = sum_i upstream[i,s](-tau_i)exp(-mu_mtau_i)x[m,s]

Arguments

Type	Intent	Attributes		Name
class(array_type),	intent(in)		::	this
real(kind=real32),	intent(in),	dimension(:,:)	::	upstream_grad
real(kind=real32),	intent(out),	dimension(:,:)	::	output

pure subroutine get_partial_lno_decode_spectral_val(this, upstream_grad, output)

dL/dx = D^T @ upstream [M, batch]

Arguments

Type	Intent	Attributes		Name
class(array_type),	intent(in)		::	this
real(kind=real32),	intent(in),	dimension(:,:)	::	upstream_grad
real(kind=real32),	intent(out),	dimension(:,:)	::	output

pure subroutine get_partial_lno_encode_input_val(this, upstream_grad, output)

dL/du = E^T @ upstream [n_in, batch]

Arguments

Type	Intent	Attributes		Name
class(array_type),	intent(in)		::	this
real(kind=real32),	intent(in),	dimension(:,:)	::	upstream_grad
real(kind=real32),	intent(out),	dimension(:,:)	::	output

pure subroutine get_partial_lno_encode_poles_val(this, upstream_grad, output)

dL/dmu_m per sample: output[m,s] = upstream[m,s] * sum_j (-t_j) * exp(-mu_m*t_j) * u[j,s]

Arguments

Type	Intent	Attributes		Name
class(array_type),	intent(in)		::	this
real(kind=real32),	intent(in),	dimension(:,:)	::	upstream_grad
real(kind=real32),	intent(out),	dimension(:,:)	::	output

pure subroutine get_partial_ono_decode_basis_val(this, upstream_grad, output)

dL/dB per sample through Gram-Schmidt backward.

Arguments

Type	Intent	Attributes		Name
class(array_type),	intent(in)		::	this
real(kind=real32),	intent(in),	dimension(:,:)	::	upstream_grad
real(kind=real32),	intent(out),	dimension(:,:)	::	output

pure subroutine get_partial_ono_decode_mixed_val(this, upstream_grad, output)

dL/dx = Q^T @ upstream [k, batch]

Arguments

Type	Intent	Attributes		Name
class(array_type),	intent(in)		::	this
real(kind=real32),	intent(in),	dimension(:,:)	::	upstream_grad
real(kind=real32),	intent(out),	dimension(:,:)	::	output

pure subroutine get_partial_ono_encode_basis_val(this, upstream_grad, output)

dL/dB per sample through Gram-Schmidt backward.

Arguments

Type	Intent	Attributes		Name
class(array_type),	intent(in)		::	this
real(kind=real32),	intent(in),	dimension(:,:)	::	upstream_grad
real(kind=real32),	intent(out),	dimension(:,:)	::	output

pure subroutine get_partial_ono_encode_input_val(this, upstream_grad, output)

dL/du = Q @ upstream [n, batch]

Arguments

Type	Intent	Attributes		Name
class(array_type),	intent(in)		::	this
real(kind=real32),	intent(in),	dimension(:,:)	::	upstream_grad
real(kind=real32),	intent(out),	dimension(:,:)	::	output

Module Functions

module function elem_scale(input, scale) result(c)

Element-wise scaling with explicit support for sample-independent scale.

Arguments

Type	Intent	Optional	Attributes		Name
class(array_type),	intent(in),		target	::	input	Input tensor [n, batch]
class(array_type),	intent(in),		target	::	scale	Scale tensor [n, 1]

Return Value type(array_type), pointer

Scaled output tensor

module function gno_aggregate(features, edge_kernels, adj_ia, adj_ja, F_in, F_out) result(c)

Aggregate neighbour messages using pre-computed per-edge kernels.

Arguments

Type	Intent	Attributes		Name
class(array_type),	intent(in),	target	::	features	Node features [F_in, num_vertices]
class(array_type),	intent(in),	target	::	edge_kernels	Per-edge kernel values [F_out*F_in, num_edges]
integer,	intent(in),	dimension(:)	::	adj_ia	CSR row pointers
integer,	intent(in),	dimension(:,:)	::	adj_ja	CSR column indices
integer,	intent(in)		::	F_in	Feature dimensions
integer,	intent(in)		::	F_out	Feature dimensions

Return Value type(array_type), pointer

Aggregated node output tensor

module function gno_kernel_eval(coords, kernel_params, adj_ia, adj_ja, coord_dim, kernel_hidden, F_in, F_out) result(c)

Evaluate the GNO kernel MLP on every directed edge in the graph.

Arguments

Type	Intent	Attributes		Name
class(array_type),	intent(in),	target	::	coords	Edge features / relative coordinates [d, num_edges]
class(array_type),	intent(in),	target	::	kernel_params	Packed kernel parameters [Hd + H + FH + F, 1]
integer,	intent(in),	dimension(:)	::	adj_ia	CSR row pointers (size num_vertices + 1)
integer,	intent(in),	dimension(:,:)	::	adj_ja	CSR column indices (adj_ja(1,:) = neighbour index)
integer,	intent(in)		::	coord_dim	Metadata for unpacking kernel_params
integer,	intent(in)		::	kernel_hidden	Metadata for unpacking kernel_params
integer,	intent(in)		::	F_in	Metadata for unpacking kernel_params
integer,	intent(in)		::	F_out	Metadata for unpacking kernel_params

Return Value type(array_type), pointer

Output per-edge kernel values

module function lno_decode(spectral, poles, num_outputs, num_modes) result(c)

Decode through the Laplace basis built from learnable poles.

Arguments

Type	Intent	Attributes		Name
class(array_type),	intent(in),	target	::	spectral	Spectral tensor [M, batch]
class(array_type),	intent(in),	target	::	poles	Learnable poles [M, 1]
integer,	intent(in)		::	num_outputs	Output dimension and number of modes
integer,	intent(in)		::	num_modes	Output dimension and number of modes

Return Value type(array_type), pointer

Decoded output tensor

module function lno_encode(input, poles, num_inputs, num_modes) result(c)

Encode input through the Laplace basis built from learnable poles.

Arguments

Type	Intent	Attributes		Name
class(array_type),	intent(in),	target	::	input	Input signal tensor [n_in, batch]
class(array_type),	intent(in),	target	::	poles	Learnable poles [M, 1]
integer,	intent(in)		::	num_inputs	Input dimension and number of modes
integer,	intent(in)		::	num_modes	Input dimension and number of modes

Return Value type(array_type), pointer

Encoded output tensor

module function ono_decode(mixed, basis_weights, num_inputs, num_basis) result(c)

Decode through an orthogonalised basis.

Arguments

Type	Intent	Attributes		Name
class(array_type),	intent(in),	target	::	mixed	Mixed spectral tensor [k, batch]
class(array_type),	intent(in),	target	::	basis_weights	Flattened basis matrix parameters [n*k, 1]
integer,	intent(in)		::	num_inputs	Output dimension and basis size
integer,	intent(in)		::	num_basis	Output dimension and basis size

Return Value type(array_type), pointer

Decoded output tensor

module function ono_encode(input, basis_weights, num_inputs, num_basis) result(c)

Encode input through an orthogonalised basis.

Arguments

Type	Intent	Attributes		Name
class(array_type),	intent(in),	target	::	input	Input tensor [n, batch]
class(array_type),	intent(in),	target	::	basis_weights	Flattened basis matrix parameters [n*k, 1]
integer,	intent(in)		::	num_inputs	Input dimension and basis size
integer,	intent(in)		::	num_basis	Input dimension and basis size

Return Value type(array_type), pointer

Encoded output tensor

athena__diffstruc_extd_submodule_nop Submodule

Contents

Functions

Subroutines

Module Functions

Uses

Ancestors:

Functions

function get_partial_gno_agg_features(this, upstream_grad) result(output)

Arguments

Return Value type(array_type)

function get_partial_gno_agg_kernels(this, upstream_grad) result(output)

Arguments

Return Value type(array_type)

function get_partial_gno_kernel_coords(this, upstream_grad) result(output)

Arguments

Return Value type(array_type)

function get_partial_gno_kernel_params(this, upstream_grad) result(output)

Arguments

Return Value type(array_type)

function get_partial_lno_decode_poles(this, upstream_grad) result(output)

Arguments

Return Value type(array_type)

function get_partial_lno_decode_spectral(this, upstream_grad) result(output)

Arguments

Return Value type(array_type)

function get_partial_lno_encode_input(this, upstream_grad) result(output)

Arguments

Return Value type(array_type)

function get_partial_lno_encode_poles(this, upstream_grad) result(output)

Arguments

Return Value type(array_type)

function get_partial_ono_decode_basis(this, upstream_grad) result(output)

Arguments

Return Value type(array_type)

function get_partial_ono_decode_mixed(this, upstream_grad) result(output)

Arguments

Return Value type(array_type)

function get_partial_ono_encode_basis(this, upstream_grad) result(output)

Arguments

Return Value type(array_type)

function get_partial_ono_encode_input(this, upstream_grad) result(output)

Arguments

Return Value type(array_type)

Subroutines

pure subroutine get_partial_elem_scale_input_val(this, upstream_grad, output)

Arguments

pure subroutine get_partial_elem_scale_scale_val(this, upstream_grad, output)

Arguments

pure subroutine get_partial_gno_agg_features_val(this, upstream_grad, output)

Arguments

pure subroutine get_partial_gno_agg_kernels_val(this, upstream_grad, output)

Arguments

pure subroutine get_partial_gno_kernel_coords_val(this, upstream_grad, output)

Arguments

pure subroutine get_partial_gno_kernel_params_val(this, upstream_grad, output)

Arguments

pure subroutine get_partial_lno_decode_poles_val(this, upstream_grad, output)

Arguments

pure subroutine get_partial_lno_decode_spectral_val(this, upstream_grad, output)

Arguments

pure subroutine get_partial_lno_encode_input_val(this, upstream_grad, output)

Arguments

pure subroutine get_partial_lno_encode_poles_val(this, upstream_grad, output)

Arguments

pure subroutine get_partial_ono_decode_basis_val(this, upstream_grad, output)

Arguments

pure subroutine get_partial_ono_decode_mixed_val(this, upstream_grad, output)

Arguments

pure subroutine get_partial_ono_encode_basis_val(this, upstream_grad, output)

Arguments

pure subroutine get_partial_ono_encode_input_val(this, upstream_grad, output)

Arguments

Module Functions

module function elem_scale(input, scale) result(c) Interface →

Arguments

Return Value type(array_type), pointer

module function gno_aggregate(features, edge_kernels, adj_ia, adj_ja, F_in, F_out) result(c) Interface →

Arguments

Return Value type(array_type), pointer

module function elem_scale(input, scale) result(c)

module function gno_aggregate(features, edge_kernels, adj_ia, adj_ja, F_in, F_out) result(c)

module function gno_kernel_eval(coords, kernel_params, adj_ia, adj_ja, coord_dim, kernel_hidden, F_in, F_out) result(c)

module function lno_decode(spectral, poles, num_outputs, num_modes) result(c)

module function lno_encode(input, poles, num_inputs, num_modes) result(c)

module function ono_decode(mixed, basis_weights, num_inputs, num_basis) result(c)

module function ono_encode(input, basis_weights, num_inputs, num_basis) result(c)