athena__onnx_msgpass_utils Module

Shared ONNX builder helpers for message-passing layers.

This module factors out the repeated edge-index extraction, scatter accumulation, weight export, and output naming logic used by the Duvenaud and Kipf message-passing layers.


Uses


Variables

Type Visibility Attributes Name Initial
character(len=*), private, parameter :: onnx_axis0_attr = '        "attribute": [{"name": "axis", "i": "0", "type": "INT"}]'
character(len=*), private, parameter :: onnx_cast_float_attr = '        "attribute": [{"name": "to", "i": "1", "type": "INT"}]'
character(len=*), private, parameter :: onnx_cast_int64_attr = '        "attribute": [{"name": "to", "i": "7", "type": "INT"}]'
character(len=*), private, parameter :: onnx_concat_axis0_attr = '        "attribute": [{"name": "axis", "i": "0", "type": "INT"}]'
character(len=*), private, parameter :: onnx_concat_axis1_attr = '        "attribute": [{"name": "axis", "i": "1", "type": "INT"}]'
character(len=*), private, parameter :: onnx_reduce_sum_attr = '        "attribute": [{"name": "keepdims", "i": "0", '//'"type": "INT"}]'
character(len=*), private, parameter :: onnx_scatter_add_attr = '        "attribute": ['//'{"name": "axis", "i": "0", "type": "INT"}, '//'{"name": "reduction", "s": "YWRk", "type": "STRING"}]'
character(len=*), private, parameter :: onnx_softmax_axis0_attr = '        "attribute": [{"name": "axis", "i": "0", "type": "INT"}]'
character(len=*), private, parameter :: onnx_transpose_10_attr = '        "attribute": [{"name": "perm", "ints": ["1", "0"], '//'"type": "INTS"}]'

Subroutines

public subroutine emit_edge_index_component(tp, edge_index_in, index_name, tag, component_out, nodes, num_nodes)

Gather one edge_index row and squeeze it into a vector.

Arguments

Type IntentOptional Attributes Name
character(len=*), intent(in) :: tp
character(len=*), intent(in) :: edge_index_in
character(len=*), intent(in) :: index_name
character(len=*), intent(in) :: tag
character(len=128), intent(out) :: component_out
type(onnx_node_type), intent(inout), dimension(:) :: nodes
integer, intent(inout) :: num_nodes

public subroutine emit_msgpass_graph_inputs(prefix, input_shape, graph_inputs, num_inputs)

Emit the standard graph input tensors used by message-passing layers.

Read more…

Arguments

Type IntentOptional Attributes Name
character(len=*), intent(in) :: prefix

Input name prefix (e.g. "input_1")
integer, intent(in), dimension(:) :: input_shape

Layer input shape [num_vertex_features, num_edge_features]
type(onnx_tensor_type), intent(inout), dimension(:) :: graph_inputs

Accumulator for graph inputs
integer, intent(inout) :: num_inputs

Current number of graph inputs

public subroutine emit_output_identity(prefix, source_name, activation_name, nodes, num_nodes)

Emit a final Identity node using the standard ATHENA output naming.

Arguments

Type IntentOptional Attributes Name
character(len=*), intent(in) :: prefix

Layer node prefix
character(len=*), intent(in) :: source_name

Source tensor to rename
character(len=*), intent(in) :: activation_name

Final activation name used in the exported output suffix
type(onnx_node_type), intent(inout), dimension(:) :: nodes

Accumulator for ONNX nodes
integer, intent(inout) :: num_nodes

Current number of nodes

public subroutine emit_scatter_aggregator(tp, vertex_in, target_in, message_in, feature_dim, nodes, num_nodes, inits, num_inits, aggr_out)

Emit the zero-initialise, expand, and scatter-add aggregation block.

Arguments

Type IntentOptional Attributes Name
character(len=*), intent(in) :: tp
character(len=*), intent(in) :: vertex_in
character(len=*), intent(in) :: target_in
character(len=*), intent(in) :: message_in
integer, intent(in) :: feature_dim
type(onnx_node_type), intent(inout), dimension(:) :: nodes
integer, intent(inout) :: num_nodes
type(onnx_initialiser_type), intent(inout), dimension(:) :: inits
integer, intent(inout) :: num_inits
character(len=128), intent(out) :: aggr_out

public subroutine emit_weight_initialiser_2d(name, nrows, ncols, weight_data, inits, num_inits)

Store a 2D weight matrix as an ONNX initialiser in row-major order.

Arguments

Type IntentOptional Attributes Name
character(len=*), intent(in) :: name
integer, intent(in) :: nrows
integer, intent(in) :: ncols
real(kind=real32), intent(in) :: weight_data(:)
type(onnx_initialiser_type), intent(inout), dimension(:) :: inits
integer, intent(inout) :: num_inits

public subroutine emit_weight_initialiser_3d(name, nslices, nrows, ncols, weight_data, inits, num_inits)

Store a stacked bank of 2D weight matrices as one ONNX tensor.

Arguments

Type IntentOptional Attributes Name
character(len=*), intent(in) :: name
integer, intent(in) :: nslices
integer, intent(in) :: nrows
integer, intent(in) :: ncols
real(kind=real32), intent(in) :: weight_data(:)
type(onnx_initialiser_type), intent(inout), dimension(:) :: inits
integer, intent(inout) :: num_inits

public subroutine get_timestep_output_name(prefix, t, activation_name, inactive_suffix, activation_suffix, output)

Build the canonical ONNX output name for one exported timestep.

Arguments

Type IntentOptional Attributes Name
character(len=*), intent(in) :: prefix
integer, intent(in) :: t
character(len=*), intent(in) :: activation_name
character(len=*), intent(in) :: inactive_suffix
character(len=*), intent(in) :: activation_suffix
character(len=128), intent(out) :: output

private subroutine add_graph_input_tensor(graph_inputs, num_inputs, name, elem_type, dim1, dim_param1, dim2, dim_param2)

Add one graph input tensor declaration to the ONNX input list.

Arguments

Type IntentOptional Attributes Name
type(onnx_tensor_type), intent(inout), dimension(:) :: graph_inputs
integer, intent(inout) :: num_inputs
character(len=*), intent(in) :: name
integer, intent(in) :: elem_type
integer, intent(in) :: dim1
character(len=*), intent(in) :: dim_param1
integer, intent(in), optional :: dim2
character(len=*), intent(in), optional :: dim_param2