Pytorch: Back-propagation of custom many-to-one nonlinearity

Question

I have code for an input-output nonlinear function that takes a list of inputs X and weights W and produces a single nonlinear output. I am interested in using this as my "neuron" and seeing if I can use back-propagation to train this. (Ideally I have a number of these neurons chained together, but a single one is fine for now.)

I've asked before if it is possible to do training on a many-to-one nonlinear function, and the answers seem to suggest that it is straightforward to use autograd to do the backpropagation.

import numpy as np
from scipy import integrate, special
from scipy.constants import epsilon_0

z_values =  np.linspace(1e-10, 1-1e-10, 100)

def readinKernel(wdummy, z, Ec, Ep, kval=1):
    return (Ec * kval * special.jv(0, Ec * kval * np.sqrt(np.outer(z, (1 - wdummy))))* Ep / np.cosh(np.arctanh(wdummy)))

def steep_sigmoid(x, k=50):    
    return 1.0 / (1.0 + np.exp(-k*x))

def readoutKernel(zdummy, z, B_in, Ec, kval=1):
    return (1 / np.sqrt(np.maximum(1e-10, np.subtract.outer(z, zdummy))) * 
            special.jv(1, 2 * Ec * kval * 
                       np.sqrt(np.maximum(1e-10, np.subtract.outer(z, zdummy)))) *
            Ec * kval *
            steep_sigmoid(np.subtract.outer(z, zdummy), 50) * 
            np.repeat(B_in, len(zdummy)).reshape(len(B_in), len(zdummy)))

def spinwave_recursive_calculation(B_in, z_values, Ec, Ep):

    wdummy_values = np.linspace(1e-10, 1-1e-10, 100)
    zdummy_values = np.linspace(1e-10, 1-1e-10, 100)

    readin_values = readinKernel(wdummy_values, z_values, Ec, Ep)
    readout_values = readoutKernel(zdummy_values, z_values, B_in, Ec)

    readin_integrals = np.trapz(readin_values, wdummy_values, axis=1)
    readout_integrals = np.trapz(readout_values, zdummy_values, axis=1)

    spinwave = readin_integrals - readout_integrals + B_in
    return spinwave

def input_output_nonlinearity(x, w):

    Bin = np.zeros(len(z_values))
    BoutMatrix = np.tile(Bin, [len(w),1])

    for i in range(len(w)):
        E_c_val = w[i]
        E_p_val = x[i]
        # Origionally this was a recursive funciton that returned a single output array, but now I have implemented it as a for loop that returns a matrix of outputs
        # This was to try to avoid a potential bug in pytorch which was recommended by a user DerekG. (I'm not sure my implementation fixes this issue though)
        # Bout = spinwave_recursive_calculation(Bin, z_values, E_c_val, E_p_val)
        BoutMatrix[i, :] = spinwave_recursive_calculation(Bin, z_values, E_c_val, E_p_val)
        Bin = BoutMatrix[i, :]
    output = print(np.sum(np.abs(Bout)))
    return output

x = np.array([1, 0, 0])
w = np.array([1, .5, 1])

input_output_nonlinearity(x, w)

Is it straightforward to do back-propagation with this code? Any ideas how to proceed with training?