write_onnx Module Subroutine

  module subroutine write_onnx(file, network, format)
    !! Export a network to ONNX JSON format.
    !!
    !! GNN layers are exported as standard ONNX operators plus metadata
    !! needed to reconstruct the original ATHENA layer on import.
    use athena__onnx_utils, only: write_json_nodes, write_json_initialisers, &
         write_json_tensors
    implicit none

    ! Arguments
    class(network_type), intent(in) :: network
    !! Instance of the network
    character(*), intent(in) :: file
    !! Output file name
    class(*), optional, intent(in) :: format
    !! Export format: 'athena_abstract' (default) or 'onnx_expanded'

    ! Local variables
    type(onnx_node_type), allocatable :: nodes(:)
    !! Exported ONNX nodes
    type(onnx_initialiser_type), allocatable :: inits(:)
    !! Exported ONNX initialisers
    type(onnx_tensor_type), allocatable :: graph_inputs(:)
    !! Graph input tensor specifications
    type(onnx_tensor_type), allocatable :: graph_outputs(:)
    !! Graph output tensor specifications
    character(4096), allocatable :: gnn_metadata(:)
    !! Metadata entries required to reconstruct ATHENA GNN layers
    integer :: num_nodes, num_inits, num_inputs, num_outputs, num_gnn_meta
    !! Numbers of populated export records
    integer :: max_nodes, max_inits
    !! Pre-allocation sizes for node and initialiser storage
    integer :: ifmt
    !! Integer selector for export format


    !---------------------------------------------------------------------------
    ! Validate the export format and convert to integer selector
    !---------------------------------------------------------------------------
    ifmt = resolve_onnx_export_format(format)
    if(ifmt .eq. 0) return


    !--------------------------------------------------------------------------
    ! Initialise export storage
    !--------------------------------------------------------------------------
    call initialise_export_storage( &
         network, nodes, inits, graph_inputs, graph_outputs, gnn_metadata, &
         max_nodes, max_inits)

    num_nodes = 0
    num_inits = 0
    num_inputs = 0
    num_outputs = 0
    num_gnn_meta = 0


    !--------------------------------------------------------------------------
    ! Collect graph content
    !--------------------------------------------------------------------------
    call collect_export_nodes( &
         network, ifmt, nodes, num_nodes, max_nodes, &
         inits, num_inits, max_inits, &
         gnn_metadata, num_gnn_meta)
    call build_graph_inputs( &
         network, ifmt, graph_inputs, num_inputs)
    call build_graph_outputs( &
         network, ifmt, graph_outputs, num_outputs)


    !--------------------------------------------------------------------------
    ! Write the JSON model
    !--------------------------------------------------------------------------
    call write_onnx_json_file( &
         file, ifmt, nodes, num_nodes, inits, num_inits, &
         graph_inputs, num_inputs, graph_outputs, num_outputs, &
         gnn_metadata, num_gnn_meta)

  end subroutine write_onnx