module athena__neural_operator_layer !! Module containing implementation of a simple neural operator layer !! !! This module implements a neural operator layer that approximates a !! discretized integral operator. The layer combines a standard affine !! transform (local component) with a mean-field integral operator !! (global/non-local component): !! !! \[ \mathbf{v} = \sigma\!\left(\mathbf{W}\mathbf{u} !! + \mathbf{w}_k \langle\mathbf{u}\rangle + \mathbf{b}\right) \] !! !! where: !! - \(\mathbf{u} \in \mathbb{R}^{n_{in}}\) is the input (discretised function) !! - \(\mathbf{W} \in \mathbb{R}^{n_{out} \times n_{in}}\) are the local weights !! - \(\mathbf{w}_k \in \mathbb{R}^{n_{out}}\) are the integral kernel weights !! - \(\langle\mathbf{u}\rangle = \frac{1}{n_{in}}\sum_j u_j\) is the input mean !! - \(\mathbf{b} \in \mathbb{R}^{n_{out}}\) is the bias !! - \(\sigma\) is the activation function !! !! The global mean \(\langle\mathbf{u}\rangle\) acts as a rank-1 approximation !! to a continuous integral operator \(\mathcal{K}[u](x) = \kappa(x)\int u\,\mathrm{d}y\), !! where \(\mathbf{w}_k\) discretises \(\kappa\). Using this layer stacked !! in sequence provides a resolution-invariant building block similar to the !! graph neural operator family. !! !! Number of parameters: !! - with bias: \(n_{out}(n_{in} + 1) + n_{out} = n_{out}(n_{in}+2)\) !! - without bias: \(n_{out}(n_{in} + 1)\) use coreutils, only: real32, stop_program use athena__base_layer, only: learnable_layer_type, base_layer_type use athena__misc_types, only: base_actv_type, base_init_type, & onnx_attribute_type, & onnx_node_type, onnx_initialiser_type, onnx_tensor_type use athena__onnx_nop_utils, only: emit_nop_input_transpose, & emit_nop_output_tail, emit_float_initialiser, emit_matrix_initialiser use diffstruc, only: array_type, matmul, mean, operator(+) use athena__initialiser_data, only: data_init_type implicit none private public :: neural_operator_layer_type public :: read_neural_operator_layer type, extends(learnable_layer_type) :: neural_operator_layer_type !! Type for a neural operator layer integer :: num_inputs !! Number of inputs (discretisation points of the input function) integer :: num_outputs !! Number of outputs (discretisation points of the output function) type(array_type), dimension(1) :: z !! Temporary array for pre-activation values (forward propagation) contains procedure, pass(this) :: get_num_params => get_num_params_neural_operator !! Get the number of parameters for the neural operator layer procedure, pass(this) :: set_hyperparams => set_hyperparams_neural_operator !! Set the hyperparameters for the neural operator layer procedure, pass(this) :: init => init_neural_operator !! Initialise the neural operator layer procedure, pass(this) :: print_to_unit => print_to_unit_neural_operator !! Print the layer to a file procedure, pass(this) :: read => read_neural_operator !! Read the layer from a file procedure, pass(this) :: forward => forward_neural_operator !! Forward propagation procedure, pass(this) :: get_attributes => get_attributes_neural_operator !! Get layer attributes for ONNX export procedure, pass(this) :: emit_onnx_nodes => & emit_onnx_nodes_neural_operator !! Emit format-aware ONNX nodes for the layer final :: finalise_neural_operator !! Finalise neural operator layer end type neural_operator_layer_type interface neural_operator_layer_type !! Interface for setting up the neural operator layer module function layer_setup( & num_outputs, num_inputs, use_bias, & activation, & kernel_initialiser, bias_initialiser, verbose & ) result(layer) !! Setup a neural operator layer integer, intent(in) :: num_outputs !! Number of outputs integer, optional, intent(in) :: num_inputs !! Number of inputs logical, optional, intent(in) :: use_bias !! Whether to use bias class(*), optional, intent(in) :: activation !! Activation function class(*), optional, intent(in) :: kernel_initialiser, bias_initialiser !! Kernel and bias initialisers integer, optional, intent(in) :: verbose !! Verbosity level type(neural_operator_layer_type) :: layer !! Instance of the neural operator layer end function layer_setup end interface neural_operator_layer_type contains !############################################################################### subroutine finalise_neural_operator(this) !! Finalise neural operator layer implicit none ! Arguments type(neural_operator_layer_type), intent(inout) :: this !! Instance of the neural operator layer if(allocated(this%input_shape)) deallocate(this%input_shape) if(allocated(this%output)) deallocate(this%output) if(this%z(1)%allocated) call this%z(1)%deallocate() end subroutine finalise_neural_operator !############################################################################### !##############################################################################! ! * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * ! !##############################################################################! !############################################################################### pure function get_num_params_neural_operator(this) result(num_params) !! Get the number of parameters for the neural operator layer !! !! Parameters consist of: !! - W matrix : num_outputs * num_inputs !! - W_k vector : num_outputs (integral kernel coupling) !! - b vector : num_outputs (if use_bias) implicit none ! Arguments class(neural_operator_layer_type), intent(in) :: this !! Instance of the neural operator layer integer :: num_params !! Number of parameters ! W: n_out * n_in, W_k: n_out, b: n_out (if use_bias) num_params = this%num_outputs * this%num_inputs + this%num_outputs if(this%use_bias) num_params = num_params + this%num_outputs end function get_num_params_neural_operator !############################################################################### !##############################################################################! ! * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * ! !##############################################################################! !############################################################################### module function layer_setup( & num_outputs, num_inputs, & use_bias, & activation, & kernel_initialiser, bias_initialiser, verbose & ) result(layer) !! Setup a neural operator layer use athena__activation, only: activation_setup use athena__initialiser, only: initialiser_setup implicit none ! Arguments integer, intent(in) :: num_outputs !! Number of outputs integer, optional, intent(in) :: num_inputs !! Number of inputs logical, optional, intent(in) :: use_bias !! Whether to use bias class(*), optional, intent(in) :: activation !! Activation function class(*), optional, intent(in) :: kernel_initialiser, bias_initialiser !! Kernel and bias initialisers integer, optional, intent(in) :: verbose !! Verbosity level type(neural_operator_layer_type) :: layer !! Instance of the neural operator layer ! Local variables integer :: verbose_ = 0 !! Verbosity level logical :: use_bias_ = .true. !! Whether to use bias class(base_actv_type), allocatable :: activation_ !! Activation function class(base_init_type), allocatable :: kernel_initialiser_, bias_initialiser_ !! Kernel and bias initialisers if(present(verbose)) verbose_ = verbose !--------------------------------------------------------------------------- ! Set use_bias !--------------------------------------------------------------------------- if(present(use_bias)) use_bias_ = use_bias !--------------------------------------------------------------------------- ! Set activation function !--------------------------------------------------------------------------- if(present(activation))then activation_ = activation_setup(activation) else activation_ = activation_setup("none") end if !--------------------------------------------------------------------------- ! Define kernel and bias initialisers !--------------------------------------------------------------------------- if(present(kernel_initialiser))then kernel_initialiser_ = initialiser_setup(kernel_initialiser) end if if(present(bias_initialiser))then bias_initialiser_ = initialiser_setup(bias_initialiser) end if !--------------------------------------------------------------------------- ! Set hyperparameters !--------------------------------------------------------------------------- call layer%set_hyperparams( & num_outputs = num_outputs, & use_bias = use_bias_, & activation = activation_, & kernel_initialiser = kernel_initialiser_, & bias_initialiser = bias_initialiser_, & verbose = verbose_ & ) !--------------------------------------------------------------------------- ! Initialise layer shape if num_inputs is provided !--------------------------------------------------------------------------- if(present(num_inputs)) call layer%init(input_shape=[num_inputs]) end function layer_setup !############################################################################### !############################################################################### subroutine set_hyperparams_neural_operator( & this, num_outputs, & use_bias, & activation, & kernel_initialiser, bias_initialiser, & verbose & ) !! Set the hyperparameters for the neural operator layer use athena__activation, only: activation_setup use athena__initialiser, only: get_default_initialiser, initialiser_setup implicit none ! Arguments class(neural_operator_layer_type), intent(inout) :: this !! Instance of the neural operator layer integer, intent(in) :: num_outputs !! Number of outputs logical, intent(in) :: use_bias !! Whether to use bias class(base_actv_type), allocatable, intent(in) :: activation !! Activation function class(base_init_type), allocatable, intent(in) :: & kernel_initialiser, bias_initialiser !! Kernel and bias initialisers integer, optional, intent(in) :: verbose !! Verbosity level ! Local variables character(len=256) :: buffer this%name = "neural_operator" this%type = "nop" this%input_rank = 1 this%output_rank = 1 this%use_bias = use_bias this%num_outputs = num_outputs if(allocated(this%activation)) deallocate(this%activation) if(.not.allocated(activation))then this%activation = activation_setup("none") else allocate(this%activation, source=activation) end if if(allocated(this%kernel_init)) deallocate(this%kernel_init) if(.not.allocated(kernel_initialiser))then buffer = get_default_initialiser(this%activation%name) this%kernel_init = initialiser_setup(buffer) else allocate(this%kernel_init, source=kernel_initialiser) end if if(allocated(this%bias_init)) deallocate(this%bias_init) if(.not.allocated(bias_initialiser))then buffer = get_default_initialiser( & this%activation%name, & is_bias=.true. & ) this%bias_init = initialiser_setup(buffer) else if(allocated(this%bias_init)) deallocate(this%bias_init) allocate(this%bias_init, source=bias_initialiser) end if if(present(verbose))then if(abs(verbose).gt.0)then write(*,'("NEURAL_OPERATOR activation function: ",A)') & trim(this%activation%name) write(*,'("NEURAL_OPERATOR kernel initialiser: ",A)') & trim(this%kernel_init%name) write(*,'("NEURAL_OPERATOR bias initialiser: ",A)') & trim(this%bias_init%name) end if end if end subroutine set_hyperparams_neural_operator !############################################################################### !############################################################################### subroutine init_neural_operator(this, input_shape, verbose) !! Initialise neural operator layer implicit none ! Arguments class(neural_operator_layer_type), intent(inout) :: this !! Instance of the neural operator layer integer, dimension(:), intent(in) :: input_shape !! Input shape integer, optional, intent(in) :: verbose !! Verbosity level ! Local variables integer :: num_inputs !! Effective fan-in for initialisation integer :: verbose_ = 0 !--------------------------------------------------------------------------- ! Initialise optional arguments !--------------------------------------------------------------------------- if(present(verbose)) verbose_ = verbose !--------------------------------------------------------------------------- ! Initialise number of inputs !--------------------------------------------------------------------------- if(.not.allocated(this%input_shape)) call this%set_shape(input_shape) this%num_inputs = this%input_shape(1) this%output_shape = [this%num_outputs] this%num_params = this%get_num_params() !--------------------------------------------------------------------------- ! Allocate parameters ! ! params(1): W (n_out x n_in) - local transform weights ! params(2): W_k (n_out x 1) - integral kernel coupling weights ! params(3): b (n_out) - bias [only when use_bias=.true.] !--------------------------------------------------------------------------- allocate(this%weight_shape(2,1)) this%weight_shape(:,1) = [ this%num_outputs, this%num_inputs ] if(this%use_bias)then this%bias_shape = [ this%num_outputs ] allocate(this%params(3)) else allocate(this%params(2)) end if ! W: local transform (n_out x n_in) call this%params(1)%allocate([this%weight_shape(:,1), 1]) call this%params(1)%set_requires_grad(.true.) this%params(1)%fix_pointer = .true. this%params(1)%is_sample_dependent = .false. this%params(1)%is_temporary = .false. ! W_k: integral kernel coupling (n_out x 1) call this%params(2)%allocate([this%num_outputs, 1, 1]) call this%params(2)%set_requires_grad(.true.) this%params(2)%fix_pointer = .true. this%params(2)%is_sample_dependent = .false. this%params(2)%is_temporary = .false. num_inputs = this%num_inputs if(this%use_bias)then num_inputs = this%num_inputs + 1 call this%params(3)%allocate([this%bias_shape, 1]) call this%params(3)%set_requires_grad(.true.) this%params(3)%fix_pointer = .true. this%params(3)%is_sample_dependent = .false. this%params(3)%is_temporary = .false. end if !--------------------------------------------------------------------------- ! Initialise W with kernel initialiser !--------------------------------------------------------------------------- call this%kernel_init%initialise( & this%params(1)%val(:,1), & fan_in = num_inputs, fan_out = this%num_outputs, & spacing = [ this%num_outputs ] & ) !--------------------------------------------------------------------------- ! Initialise W_k with kernel initialiser (smaller scale), treating it as ! a rank-1 integral correction so fan_in=1 !--------------------------------------------------------------------------- call this%kernel_init%initialise( & this%params(2)%val(:,1), & fan_in = num_inputs, fan_out = this%num_outputs, & spacing = [ this%num_outputs ] & ) !--------------------------------------------------------------------------- ! Initialise bias if used !--------------------------------------------------------------------------- if(this%use_bias)then call this%bias_init%initialise( & this%params(3)%val(:,1), & fan_in = num_inputs, fan_out = this%num_outputs & ) end if !--------------------------------------------------------------------------- ! Allocate output and pre-activation arrays !--------------------------------------------------------------------------- if(allocated(this%output)) deallocate(this%output) allocate(this%output(1,1)) if(this%z(1)%allocated) call this%z(1)%deallocate() end subroutine init_neural_operator !############################################################################### !##############################################################################! ! * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * ! !##############################################################################! !############################################################################### subroutine print_to_unit_neural_operator(this, unit) !! Print neural operator layer to unit use coreutils, only: to_upper implicit none ! Arguments class(neural_operator_layer_type), intent(in) :: this !! Instance of the neural operator layer integer, intent(in) :: unit !! File unit ! Write hyperparameters !--------------------------------------------------------------------------- write(unit,'(3X,"NUM_INPUTS = ",I0)') this%num_inputs write(unit,'(3X,"NUM_OUTPUTS = ",I0)') this%num_outputs write(unit,'(3X,"USE_BIAS = ",L1)') this%use_bias if(this%activation%name .ne. 'none')then call this%activation%print_to_unit(unit) end if ! Write weights, kernel coupling, and optional bias !--------------------------------------------------------------------------- write(unit,'("WEIGHTS")') write(unit,'(5(E16.8E2))') this%params(1)%val(:,1) ! W write(unit,'(5(E16.8E2))') this%params(2)%val(:,1) ! W_k if(this%use_bias)then write(unit,'(5(E16.8E2))') this%params(3)%val(:,1) ! b end if write(unit,'("END WEIGHTS")') end subroutine print_to_unit_neural_operator !############################################################################### !############################################################################### subroutine read_neural_operator(this, unit, verbose) !! Read neural operator layer from file use athena__tools_infile, only: assign_val, assign_vec, move use coreutils, only: to_lower, to_upper, icount use athena__activation, only: read_activation use athena__initialiser, only: initialiser_setup implicit none ! Arguments class(neural_operator_layer_type), intent(inout) :: this !! Instance of the neural operator layer integer, intent(in) :: unit !! Unit number integer, optional, intent(in) :: verbose !! Verbosity level ! Local variables integer :: stat !! Status of read integer :: verbose_ = 0 !! Verbosity level integer :: i, j, k, c, itmp1, iline, num_params !! Loop variables and temporary integers integer :: num_inputs, num_outputs !! Number of inputs and outputs logical :: use_bias = .true. !! Whether to use bias character(14) :: kernel_initialiser_name='', bias_initialiser_name='' !! Initialiser names character(20) :: activation_name='' !! Activation function name class(base_actv_type), allocatable :: activation !! Activation function class(base_init_type), allocatable :: kernel_initialiser, bias_initialiser !! Initialisers character(256) :: buffer, tag, err_msg !! Buffer, tag, and error message real(real32), allocatable, dimension(:) :: data_list !! Data list integer :: param_line, final_line !! Parameter line numbers ! Initialise optional arguments !--------------------------------------------------------------------------- if(present(verbose)) verbose_ = verbose ! Loop over tags in layer card !--------------------------------------------------------------------------- iline = 0 param_line = 0 final_line = 0 tag_loop: do ! Check for end of file !------------------------------------------------------------------------ read(unit,'(A)',iostat=stat) buffer if(stat.ne.0)then write(err_msg,'("file encountered error (EoF?) before END ",A)') & to_upper(this%name) call stop_program(err_msg) return end if if(trim(adjustl(buffer)).eq."") cycle tag_loop ! Check for end of layer card !------------------------------------------------------------------------ if(trim(adjustl(buffer)).eq."END "//to_upper(trim(this%name)))then final_line = iline backspace(unit) exit tag_loop end if iline = iline + 1 tag=trim(adjustl(buffer)) if(scan(buffer,"=").ne.0) tag=trim(tag(:scan(tag,"=")-1)) ! Read parameters from file !------------------------------------------------------------------------ select case(trim(tag)) case("NUM_INPUTS") call assign_val(buffer, num_inputs, itmp1) case("NUM_OUTPUTS") call assign_val(buffer, num_outputs, itmp1) case("USE_BIAS") call assign_val(buffer, use_bias, itmp1) case("ACTIVATION") iline = iline - 1 backspace(unit) activation = read_activation(unit, iline) case("KERNEL_INITIALISER", "KERNEL_INIT", "KERNEL_INITIALIZER") call assign_val(buffer, kernel_initialiser_name, itmp1) case("BIAS_INITIALISER", "BIAS_INIT", "BIAS_INITIALIZER") call assign_val(buffer, bias_initialiser_name, itmp1) case("WEIGHTS") kernel_initialiser_name = 'zeros' bias_initialiser_name = 'zeros' param_line = iline case default ! Skip lines that only contain numbers (e.g. scientific notation) if(scan(to_lower(trim(adjustl(buffer))),& 'abcdfghijklmnopqrstuvwxyz').eq.0)then cycle tag_loop elseif(tag(:3).eq.'END')then cycle tag_loop end if write(err_msg,'("Unrecognised line in input file: ",A)') & trim(adjustl(buffer)) call stop_program(err_msg) return end select end do tag_loop kernel_initialiser = initialiser_setup(kernel_initialiser_name) bias_initialiser = initialiser_setup(bias_initialiser_name) ! Set hyperparameters and initialise layer !--------------------------------------------------------------------------- call this%set_hyperparams( & num_outputs = num_outputs, & use_bias = use_bias, & activation = activation, & kernel_initialiser = kernel_initialiser, & bias_initialiser = bias_initialiser, & verbose = verbose_ & ) call this%init(input_shape=[num_inputs]) ! Read weights if WEIGHTS card was found !--------------------------------------------------------------------------- if(param_line.eq.0)then write(0,*) "WARNING: WEIGHTS card in " // trim(this%name) // " not found" else call move(unit, param_line - iline, iostat=stat) ! Read W (num_inputs * num_outputs elements) num_params = this%num_inputs * this%num_outputs allocate(data_list(num_params), source=0._real32) c = 1 k = 1 data_concat_loop: do while(c.le.num_params) read(unit,'(A)',iostat=stat) buffer if(stat.ne.0) exit data_concat_loop k = icount(buffer) read(buffer,*,iostat=stat) (data_list(j),j=c,c+k-1) c = c + k end do data_concat_loop this%params(1)%val(:,1) = data_list deallocate(data_list) ! Read W_k (num_outputs elements) allocate(data_list(this%num_outputs), source=0._real32) c = 1 k = 1 data_concat_loop2: do while(c.le.this%num_outputs) read(unit,'(A)',iostat=stat) buffer if(stat.ne.0) exit data_concat_loop2 k = icount(buffer) read(buffer,*,iostat=stat) (data_list(j),j=c,c+k-1) c = c + k end do data_concat_loop2 this%params(2)%val(:,1) = data_list(1:this%num_outputs) deallocate(data_list) ! Read b (num_outputs elements, only if use_bias) if(use_bias)then allocate(data_list(num_outputs), source=0._real32) c = 1 k = 1 data_concat_loop3: do while(c.le.num_outputs) read(unit,'(A)',iostat=stat) buffer if(stat.ne.0) exit data_concat_loop3 k = icount(buffer) read(buffer,*,iostat=stat) (data_list(j),j=c,c+k-1) c = c + k end do data_concat_loop3 this%params(3)%val(:,1) = data_list(1:num_outputs) deallocate(data_list) end if ! Check for END WEIGHTS tag !------------------------------------------------------------------------ read(unit,'(A)') buffer if(trim(adjustl(buffer)).ne."END WEIGHTS")then write(0,*) trim(adjustl(buffer)) call stop_program("END WEIGHTS not where expected") return end if end if !--------------------------------------------------------------------------- ! Check for end of layer card !--------------------------------------------------------------------------- call move(unit, final_line - iline, iostat=stat) read(unit,'(A)') buffer if(trim(adjustl(buffer)).ne."END "//to_upper(trim(this%name)))then write(0,*) trim(adjustl(buffer)) write(err_msg,'("END ",A," not where expected")') to_upper(this%name) call stop_program(err_msg) return end if end subroutine read_neural_operator !############################################################################### !############################################################################### function read_neural_operator_layer(unit, verbose) result(layer) !! Read neural operator layer from file and return as base_layer_type implicit none ! Arguments integer, intent(in) :: unit !! Unit number integer, optional, intent(in) :: verbose !! Verbosity level class(base_layer_type), allocatable :: layer !! Allocated layer instance ! Local variables integer :: verbose_ = 0 if(present(verbose)) verbose_ = verbose allocate(layer, source=neural_operator_layer_type(num_outputs=0)) call layer%read(unit, verbose=verbose_) end function read_neural_operator_layer !############################################################################### !##############################################################################! ! * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * ! !##############################################################################! !############################################################################### subroutine forward_neural_operator(this, input) !! Forward propagation for the neural operator layer !! !! Computes: !! v = sigma( W * u + W_k * mean(u) + b ) !! !! where mean(u) is the global mean of the input (scalar per sample), !! approximating the integral operator. implicit none ! Arguments class(neural_operator_layer_type), intent(inout) :: this !! Instance of the neural operator layer class(array_type), dimension(:,:), intent(in) :: input !! Input values ! Local variables type(array_type), pointer :: ptr, ptr_mean, ptr_kern ! Local transform: W · u → shape [n_out] !--------------------------------------------------------------------------- ptr => matmul(this%params(1), input(1,1)) ! Integral (mean-field) term: W_k · mean(u) → shape [n_out] ! mean(input, dim=1) reduces over all spatial elements, giving a scalar ! per batch sample (shape [1]). matmul then expands W_k ([n_out x 1]) ! by this scalar to produce a [n_out] correction vector. !--------------------------------------------------------------------------- ptr_mean => mean(input(1,1), dim=1) ptr_kern => matmul(this%params(2), ptr_mean) ! Combine local + integral terms !--------------------------------------------------------------------------- ptr => ptr + ptr_kern ! Add bias if used !--------------------------------------------------------------------------- if(this%use_bias)then ptr => ptr + this%params(3) end if ! Apply activation function !--------------------------------------------------------------------------- call this%output(1,1)%zero_grad() if(trim(this%activation%name) .eq. "none")then call this%output(1,1)%assign_and_deallocate_source(ptr) else call this%z(1)%zero_grad() call this%z(1)%assign_and_deallocate_source(ptr) this%z(1)%is_temporary = .false. ptr => this%activation%apply(this%z(1)) call this%output(1,1)%assign_and_deallocate_source(ptr) end if this%output(1,1)%is_temporary = .false. end subroutine forward_neural_operator !############################################################################### !############################################################################### function get_attributes_neural_operator(this) result(attributes) !! Return list of neural operator attributes for ONNX export implicit none ! Arguments class(neural_operator_layer_type), intent(in) :: this !! Instance of the neural operator layer type(onnx_attribute_type), allocatable, dimension(:) :: attributes !! List of attributes for ONNX export ! Local variables character(32) :: buffer !! Buffer for integer-to-string conversion allocate(attributes(4)) write(buffer, '(I0)') this%num_inputs attributes(1) = onnx_attribute_type( & name='num_inputs', type='int', val=trim(buffer)) write(buffer, '(I0)') this%num_outputs attributes(2) = onnx_attribute_type( & name='num_outputs', type='int', val=trim(buffer)) if(this%use_bias)then buffer = '1' else buffer = '0' end if attributes(3) = onnx_attribute_type( & name='use_bias', type='int', val=trim(buffer)) attributes(4) = onnx_attribute_type( & name='activation', type='string', val=trim(this%activation%name)) end function get_attributes_neural_operator !############################################################################### !############################################################################### subroutine emit_onnx_nodes_neural_operator( & this, prefix, nodes, num_nodes, max_nodes, inits, num_inits, & max_inits, input_name, is_last_layer, format) !! Emit decomposed standard ONNX nodes for a Neural Operator layer. !! !! Forward: v = sigma(W * u + w_k * mean(u) + b) implicit none ! Arguments class(neural_operator_layer_type), intent(in) :: this !! Neural operator layer instance character(*), intent(in) :: prefix !! Layer name prefix type(onnx_node_type), intent(inout), dimension(:) :: nodes !! Node accumulator integer, intent(inout) :: num_nodes !! Node counter integer, intent(in) :: max_nodes !! Node limit type(onnx_initialiser_type), intent(inout), dimension(:) :: inits !! Initialiser accumulator integer, intent(inout) :: num_inits !! Initialiser counter integer, intent(in) :: max_inits !! Initialiser limit character(*), optional, intent(in) :: input_name !! Name of the input tensor logical, optional, intent(in) :: is_last_layer !! Whether this is the last layer integer, optional, intent(in) :: format !! Export format selector ! Local variables integer :: n character(128) :: w_name, wk_name, b_name character(128) :: trans_in_out, mm_w_out, reduce_out character(128) :: mul_out, add_out, add_b_out, final_output, & output_source character(4096) :: reduce_attr integer :: format_ format_ = 1 if(present(format)) format_ = format if(format_ .ne. 2) return if(.not.present(input_name)) return if(.not.present(is_last_layer)) return write(w_name, '(A,".W")') trim(prefix) write(wk_name, '(A,".w_k")') trim(prefix) write(b_name, '(A,".b")') trim(prefix) write(trans_in_out, '("/",A,"/Transpose_output_0")') trim(prefix) write(mm_w_out, '("/",A,"/MatMul_output_0")') trim(prefix) write(reduce_out, '("/",A,"/ReduceMean_output_0")') trim(prefix) write(mul_out, '("/",A,"/Mul_output_0")') trim(prefix) write(add_out, '("/",A,"/Add_output_0")') trim(prefix) write(add_b_out, '("/",A,"/Add_1_output_0")') trim(prefix) reduce_attr = ' "attribute": [{"name": "axes", "ints": ' // & '["0"], "type": "INTS"}, {"name": "keepdims", "i": "1", ' // & '"type": "INT"}]' ! Transpose(input) call emit_nop_input_transpose(trim(prefix), trim(input_name), nodes, & num_nodes, trim(trans_in_out)) ! MatMul(W, x_t) num_nodes = num_nodes + 1 write(nodes(num_nodes)%name, '("/",A,"/MatMul")') trim(prefix) nodes(num_nodes)%op_type = 'MatMul' allocate(nodes(num_nodes)%inputs(2)) nodes(num_nodes)%inputs(1) = trim(w_name) nodes(num_nodes)%inputs(2) = trim(trans_in_out) allocate(nodes(num_nodes)%outputs(1)) nodes(num_nodes)%outputs(1) = trim(mm_w_out) nodes(num_nodes)%attributes_json = '' ! ReduceMean(x_t, axis=0) num_nodes = num_nodes + 1 write(nodes(num_nodes)%name, '("/",A,"/ReduceMean")') trim(prefix) nodes(num_nodes)%op_type = 'ReduceMean' allocate(nodes(num_nodes)%inputs(1)) nodes(num_nodes)%inputs(1) = trim(trans_in_out) allocate(nodes(num_nodes)%outputs(1)) nodes(num_nodes)%outputs(1) = trim(reduce_out) nodes(num_nodes)%attributes_json = reduce_attr ! Mul(w_k, mean) num_nodes = num_nodes + 1 write(nodes(num_nodes)%name, '("/",A,"/Mul")') trim(prefix) nodes(num_nodes)%op_type = 'Mul' allocate(nodes(num_nodes)%inputs(2)) nodes(num_nodes)%inputs(1) = trim(wk_name) nodes(num_nodes)%inputs(2) = trim(reduce_out) allocate(nodes(num_nodes)%outputs(1)) nodes(num_nodes)%outputs(1) = trim(mul_out) nodes(num_nodes)%attributes_json = '' ! Add(local, kernel) num_nodes = num_nodes + 1 write(nodes(num_nodes)%name, '("/",A,"/Add")') trim(prefix) nodes(num_nodes)%op_type = 'Add' allocate(nodes(num_nodes)%inputs(2)) nodes(num_nodes)%inputs(1) = trim(mm_w_out) nodes(num_nodes)%inputs(2) = trim(mul_out) allocate(nodes(num_nodes)%outputs(1)) nodes(num_nodes)%outputs(1) = trim(add_out) nodes(num_nodes)%attributes_json = '' ! Add(combined, bias) if(this%use_bias)then num_nodes = num_nodes + 1 write(nodes(num_nodes)%name, '("/",A,"/Add_1")') trim(prefix) nodes(num_nodes)%op_type = 'Add' allocate(nodes(num_nodes)%inputs(2)) nodes(num_nodes)%inputs(1) = trim(add_out) nodes(num_nodes)%inputs(2) = trim(b_name) allocate(nodes(num_nodes)%outputs(1)) nodes(num_nodes)%outputs(1) = trim(add_b_out) nodes(num_nodes)%attributes_json = '' end if if(this%use_bias)then output_source = add_b_out else output_source = add_out end if call emit_nop_output_tail(trim(prefix), trim(this%activation%name), & is_last_layer, trim(output_source), nodes, num_nodes, final_output) ! Initialisers n = this%num_outputs * this%num_inputs call emit_matrix_initialiser(trim(w_name), this%params(1)%val(:,1), & this%num_outputs, this%num_inputs, inits, num_inits) ! w_k: mean-field kernel [n_out, 1] call emit_float_initialiser(trim(wk_name), this%params(2)%val(:,1), & [this%num_outputs, 1], inits, num_inits) ! b: bias [n_out, 1] if(this%use_bias)then call emit_float_initialiser(trim(b_name), this%params(3)%val(:,1), & [this%num_outputs, 1], inits, num_inits) end if end subroutine emit_onnx_nodes_neural_operator !############################################################################### end module athena__neural_operator_layer