build_neural_operator_onnx_expanded_nop Function

public function build_neural_operator_onnx_expanded_nop(prefix, nodes, num_nodes, inits, num_inits) result(layer)

Build one neural operator layer from an expanded-ONNX node cluster.

Arguments

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

Layer node prefix (e.g. layer3)

type(onnx_node_type), intent(in) :: nodes(:)

Parsed ONNX nodes

integer, intent(in) :: num_nodes

Number of valid node entries

type(onnx_initialiser_type), intent(in) :: inits(:)

Parsed ONNX initialisers

integer, intent(in) :: num_inits

Number of valid initialiser entries

Return Value class(base_layer_type), allocatable

Constructed neural operator layer


Source Code

  function build_neural_operator_onnx_expanded_nop( &
       prefix, nodes, num_nodes, inits, num_inits) result(layer)
    !! Build one neural operator layer from an expanded-ONNX node cluster.
    use athena__neural_operator_layer, only: neural_operator_layer_type
    implicit none

    ! Arguments
    character(*), intent(in) :: prefix
    !! Layer node prefix (e.g. layer3)
    type(onnx_node_type), intent(in) :: nodes(:)
    !! Parsed ONNX nodes
    integer, intent(in) :: num_nodes
    !! Number of valid node entries
    type(onnx_initialiser_type), intent(in) :: inits(:)
    !! Parsed ONNX initialisers
    integer, intent(in) :: num_inits
    !! Number of valid initialiser entries
    class(base_layer_type), allocatable :: layer
    !! Constructed neural operator layer

    ! Local variables
    type(neural_operator_layer_type) :: typed_layer
    !! Concrete layer object before up-casting
    integer :: matmul_idx, mul_idx, add1_idx
    !! Node indices for the neural operator decomposition
    integer :: w_idx, wk_idx, b_idx
    !! Initialiser indices used to populate the layer parameters
    integer :: num_inputs, num_outputs
    !! Reconstructed layer dimensions
    logical :: use_bias
    !! Whether the graph includes a bias add
    character(64) :: activation_name
    !! Activation reconstructed from the tail of the graph

    matmul_idx = find_onnx_expanded_node_by_suffix(nodes, num_nodes, prefix, &
         'MatMul')
    mul_idx = find_onnx_expanded_node_by_suffix(nodes, num_nodes, prefix, &
         'Mul')
    add1_idx = find_onnx_expanded_node_by_suffix(nodes, num_nodes, prefix, &
         'Add_1')

    if(matmul_idx .le. 0 .or. mul_idx .le. 0)then
       call stop_program('Neural operator ONNX cluster is incomplete for ' // &
            trim(prefix))
    end if

    w_idx = find_node_initialiser_index(nodes(matmul_idx), inits, num_inits)
    wk_idx = find_node_initialiser_index(nodes(mul_idx), inits, num_inits)

    if(min(w_idx, wk_idx) .le. 0)then
       call stop_program('Neural operator ONNX parameters are missing for ' // &
            trim(prefix))
    end if

    num_outputs = inits(w_idx)%dims(1)
    num_inputs = inits(w_idx)%dims(2)
    use_bias = add1_idx .gt. 0
    activation_name = detect_onnx_expanded_nop_activation( &
         prefix, nodes, num_nodes)

    typed_layer = neural_operator_layer_type( &
         num_outputs=num_outputs, num_inputs=num_inputs, &
         use_bias=use_bias, activation=trim(activation_name))

    call load_onnx_expanded_matrix_param( &
         typed_layer%params(1), inits(w_idx), num_outputs, num_inputs)
    typed_layer%params(2)%val(:,1) = inits(wk_idx)%data

    if(use_bias)then
       b_idx = find_node_initialiser_index(nodes(add1_idx), inits, num_inits)
       if(b_idx .le. 0)then
          call stop_program('Neural operator bias initialiser missing for ' // &
               trim(prefix))
       end if
       typed_layer%params(3)%val(:,1) = inits(b_idx)%data
    end if

    allocate(layer, source=typed_layer)

  end function build_neural_operator_onnx_expanded_nop