#-------------------------------------------------------------
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#
#-------------------------------------------------------------

source("nn/layers/batch_norm2d_old.dml") as bn2d
source("nn/layers/conv2d_builtin.dml") as conv2d
source("nn/layers/relu.dml") as relu
source("nn/layers/max_pool2d_builtin.dml") as mp2d
source("nn/layers/global_avg_pool2d.dml") as ap2d
source("nn/layers/affine.dml") as fc
source("nn/layers/softmax.dml") as softmax

conv3x3_forward = function(matrix[double] X, matrix[double] W,
                           int C_in, int C_out, int Hin, int Win,
                           int strideh, int stridew)
    return (matrix[double] out, int Hout, int Wout) {
    /*
     * Simple 2D conv layer with 3x3 filter
     */
    # bias should not be applied
    bias = matrix(0, C_out, 1)
    [out, Hout, Wout] = conv2d::forward(X, W, bias, C_in, Hin, Win, 3, 3, strideh, stridew, 1, 1)
}

conv3x3_backward = function(matrix[double] dOut, int Hout, int Wout,
                            matrix[double] X, matrix[double] W,
                            int C_in, int C_out, int Hin, int Win,
                            int strideh, int stridew)
    return (matrix[double] dX, matrix[double] dW) {
    /*
     * Simple 2D conv layer with 3x3 filter
     */
    bias = matrix(0, C_out, 1)
    [dX, dW, db] = conv2d::backward(dOut, Hout, Wout, X, W, bias, C_in, Hin, Win, 3, 3, strideh, stridew, 1, 1)
}

conv1x1_forward = function(matrix[double] X, matrix[double] W,
                           int C_in, int C_out, int Hin, int Win,
                           int strideh, int stridew)
    return(matrix[double] out, int Hout, int Wout) {
    /*
     * Simple 2D conv layer with 1x1 filter
     */
    # bias should not be applied
    bias = matrix(0, C_out, 1)
    [out, Hout, Wout] = conv2d::forward(X, W, bias, C_in, Hin, Win, 1, 1, strideh, stridew, 0, 0)
}

conv1x1_backward = function(matrix[double] dOut, int Hout, int Wout,
                           matrix[double] X, matrix[double] W,
                           int C_in, int C_out, int Hin, int Win,
                           int strideh, int stridew)
    return(matrix[double] dX, matrix[double] dW) {
    /*
     * Simple 2D conv layer with 1x1 filter
     */
    bias = matrix(0, C_out, 1)
    [dX, dW, db] = conv2d::backward(dOut, Hout, Wout, X, W, bias, C_in, Hin, Win, 1, 1, strideh, stridew, 0, 0)
}

basic_block_forward = function(matrix[double] X, list[unknown] weights,
                             int C_in, int C_base, int Hin, int Win,
                             int strideh, int stridew, string mode,
                             list[unknown] ema_means_vars)
    return (matrix[double] out, int Hout, int Wout,
            list[unknown] ema_means_vars_upd,
            list[unknown] cached_out,
            list[unknown] cached_means_vars) {
    /*
     * Computes the forward pass for a basic residual block.
     * This basic residual block (with 2 3x3 conv layers of
     * same channel size) is used in the smaller ResNets 18
     * and 34.
     *
     * Inputs:
     * - X: Inputs, of shape (N, C_in*Hin*Win).
     * - weights: list of weights for all layers of res block
     *     with the following order/content:
     *   -> 1: Weights of conv 1, of shape (C_base, C_in*3*3).
     *   -> 2: Weights of batch norm 1, of shape (C_base, 1).
     *   -> 3: Bias of batch norm 1, of shape (C_base, 1).
     *   -> 4: Weights of conv 2, of shape (C_base, C_base*3*3).
     *   -> 5: Weights of batch norm 2, of shape (C_base, 1).
     *   -> 6: Bias of batch norm 2, of shape (C_base, 1).
     *   If the block should downsample X:
     *   -> 7: Weights of downsample conv, of shape (C_base, C_in*1*1).
     *   -> 8: Weights of downsample batch norm, of shape (C_base, 1).
     *   -> 9: Bias of downsample batch norm, of shape (C_base, 1).
     * - C_in: Number of input channels.
     * - C_base: Number of base channels for this block.
     * - Hin: Input height.
     * - Win: Input width.
     * - strideh: Stride over height (usually 1 or 2).
     * - stridew: Stride over width (usually same as strideh).
     * - mode: 'train' or 'test' to indicate if the model is currently
     *     being trained or tested for badge normalization layers.
     *     See badge_norm2d.dml docs for more info.
     * - ema_means_vars: List of exponential moving averages for mean
     *     and variance for badge normalization layers.
     *   -> 1: EMA for mean of badge norm 1, of shape (C_base, 1).
     *   -> 2: EMA for variance of badge norm 1, of shape (C_base, 1).
     *   -> 3: EMA for mean of badge norm 2, of shape (C_base, 1).
     *   -> 4: EMA for variance of badge norm 2, of shape (C_base, 1).
     *   If the block should downsample X:
     *   -> 5: EMA for mean of downs. badge norm, of shape (C_base, 1).
     *   -> 6: EMA for variance of downs. badge norm, of shape (C_base, 1).
     *
     * Outputs:
     * - out: Output, of shape (N, C_base*Hout*Wout).
     * - Hout: Output height.
     * - Wout: Output width.
     * - ema_means_vars_upd: List of updated exponential moving averages
     *     for mean and variance of badge normalization layers.
     * - cached_out: Outputs of each layer for computation of backward
     *     pass. Refer to the code for the order of elements.
     * - cached_means_vars: List of cached means and vars returned from
     *     each batch normalization layer. This is required for the
     *     backward pass of the network.
     */
    downsample = strideh > 1 | stridew > 1 | C_in != C_base
    # default values
    mu_bn = 0.1
    epsilon_bn = 1e-05

    # get all params
    W_conv1 = as.matrix(weights[1])
    gamma_bn1 = as.matrix(weights[2])
    beta_bn1 = as.matrix(weights[3])
    W_conv2 = as.matrix(weights[4])
    gamma_bn2 = as.matrix(weights[5])
    beta_bn2 = as.matrix(weights[6])

    ema_mean_bn1 = as.matrix(ema_means_vars[1])
    ema_var_bn1 = as.matrix(ema_means_vars[2])
    ema_mean_bn2 = as.matrix(ema_means_vars[3])
    ema_var_bn2 = as.matrix(ema_means_vars[4])

    if (downsample) {
        # gather params for downsampling
        W_conv3 = as.matrix(weights[7])
        gamma_bn3 = as.matrix(weights[8])
        beta_bn3 = as.matrix(weights[9])
        ema_mean_bn3 = as.matrix(ema_means_vars[5])
        ema_var_bn3 = as.matrix(ema_means_vars[6])
    }

    # RESIDUAL PATH
    # First convolutional layer
    [out_conv1, Hout_conv1, Wout_conv1] = conv3x3_forward(X, W_conv1, C_in, C_base, Hin, Win, strideh, stridew)
    [out_bn1, ema_mean_bn1_upd, ema_var_bn1_upd, c_m_bn1, c_v_bn1] = bn2d::forward(out_conv1, gamma_bn1, beta_bn1,
        C_base, Hout_conv1, Wout_conv1, mode, ema_mean_bn1, ema_var_bn1, mu_bn, epsilon_bn)
    out_re1 = relu::forward(out_bn1)

    # Second convolutional layer
    [out_conv2, Hout_conv2, Wout_conv2] = conv3x3_forward(out_re1, W_conv2, C_base, C_base, Hout_conv1, Wout_conv1, 1,
        1)
    [out_bn2, ema_mean_bn2_upd, ema_var_bn2_upd, c_m_bn2, c_v_bn2] = bn2d::forward(out_conv2, gamma_bn2, beta_bn2,
        C_base, Hout_conv2, Wout_conv2, mode, ema_mean_bn2, ema_var_bn2, mu_bn, epsilon_bn)

    # IDENTITY PATH
    if (downsample) {
        # downsample input
        [out_conv3, Hout_conv3, Wout_conv3] = conv1x1_forward(X, W_conv3, C_in, C_base, Hin, Win, strideh, stridew)
        [out_bn3, ema_mean_bn3_upd, ema_var_bn3_upd, c_m_bn3, c_v_bn3] = bn2d::forward(out_conv3, gamma_bn3, beta_bn3,
            C_base, Hout_conv3, Wout_conv3, mode, ema_mean_bn3, ema_var_bn3, mu_bn, epsilon_bn)
        identity = out_bn3
    } else {
        identity = X
    }

    out_re2 = relu::forward(out_bn2 + identity)

    ema_means_vars_upd = list(ema_mean_bn1_upd, ema_var_bn1_upd, ema_mean_bn2_upd, ema_var_bn2_upd)
    cached_out = list(X, Hin, Win, out_conv1, Hout_conv1, Wout_conv1, out_bn1, out_re1, out_conv2, Hout_conv2,
        Wout_conv2, out_bn2, out_re2)
    cached_means_vars = list(c_m_bn1, c_v_bn1, c_m_bn2, c_v_bn2)
    if (downsample) {
        ema_means_vars_upd = append(ema_means_vars_upd, ema_mean_bn3_upd)
        ema_means_vars_upd = append(ema_means_vars_upd, ema_var_bn3_upd)
        cached_out = append(cached_out, out_conv3)
        cached_out = append(cached_out, Hout_conv3)
        cached_out = append(cached_out, Wout_conv3)
        cached_out = append(cached_out, out_bn3)
        cached_means_vars = append(cached_means_vars, c_m_bn3)
        cached_means_vars = append(cached_means_vars, c_v_bn3)
    }

    out = out_re2
    Hout = Hout_conv2
    Wout = Wout_conv2
}

basic_block_backward = function(matrix[double] dOut, list[unknown] cached_out,
                                list[unknown] weights, int C_in, int C_base,
                                int strideh, int stridew,
                                list[unknown] cached_means_vars)
    return (matrix[double] dX, list[unknown] gradients) {
    downsample = strideh > 1 | stridew > 1 | C_in != C_base
    # default values
    mu_bn = 0.1
    epsilon_bn = 1e-05

    # get all params
    W_conv1 = as.matrix(weights[1])
    gamma_bn1 = as.matrix(weights[2])
    beta_bn1 = as.matrix(weights[3])
    W_conv2 = as.matrix(weights[4])
    gamma_bn2 = as.matrix(weights[5])
    beta_bn2 = as.matrix(weights[6])

    # cached outputs
    X = as.matrix(cached_out[1])
    Hin = as.integer(as.scalar(cached_out[2]))
    Win = as.integer(as.scalar(cached_out[3]))
    out_conv1 = as.matrix(cached_out[4])
    Hout_conv1 = as.integer(as.scalar(cached_out[5]))
    Wout_conv1 = as.integer(as.scalar(cached_out[6]))
    out_bn1 = as.matrix(cached_out[7])
    out_re1 = as.matrix(cached_out[8])
    out_conv2 = as.matrix(cached_out[9])
    Hout_conv2 = as.integer(as.scalar(cached_out[10]))
    Wout_conv2 = as.integer(as.scalar(cached_out[11]))
    out_bn2 = as.matrix(cached_out[12])
    out_re2 = as.matrix(cached_out[13])

    # cached means and vars
    mean_bn1 = as.matrix(cached_means_vars[1])
    var_bn1 = as.matrix(cached_means_vars[2])
    mean_bn2 = as.matrix(cached_means_vars[3])
    var_bn2 = as.matrix(cached_means_vars[4])

    if (downsample) {
        # gather params for downsampling
        W_conv3 = as.matrix(weights[7])
        gamma_bn3 = as.matrix(weights[8])
        beta_bn3 = as.matrix(weights[9])

        out_conv3 = as.matrix(cached_out[14])
        Hout_conv3 = as.integer(as.scalar(cached_out[15]))
        Wout_conv3 = as.integer(as.scalar(cached_out[16]))
        out_bn3 = as.matrix(cached_out[17])
        identity = out_bn3

        mean_bn3 = as.matrix(cached_means_vars[5])
        var_bn3 = as.matrix(cached_means_vars[6])
    } else {
        identity = X
    }

    dX = relu::backward(dOut, identity + out_bn2)

    # IDENTITY PATH
    if (downsample) {
        [dX_id, dgamma_bn3, dbeta_bn3] = bn2d::backward(dX, mean_bn3, var_bn3, out_conv3, gamma_bn3, C_base, Hout_conv3,
            Wout_conv3, epsilon_bn)
        [dX_id, dW_conv3] = conv1x1_backward(dX_id, Hout_conv3, Wout_conv3, X, W_conv3, C_in, C_base, Hin,
            Win, strideh, stridew)
    } else {
        dX_id = dX
    }

    # RESIDUAL PATH
    # Second convolutional layer
    [dX_res, dgamma_bn2, dbeta_bn2] = bn2d::backward(dX, mean_bn2, var_bn2, out_conv2, gamma_bn2, C_base, Hout_conv2,
        Wout_conv2, epsilon_bn)
    [dX_res, dW_conv2] = conv3x3_backward(dX_res, Hout_conv2, Wout_conv2, out_re1, W_conv2, C_base, C_base,
        Hout_conv1, Wout_conv1, 1, 1)

    # First convoluational layer
    dX_res = relu::backward(dX_res, out_bn1)
    [dX_res, dgamma_bn1, dbeta_bn1] = bn2d::backward(dX_res, mean_bn1, var_bn1, out_conv1, gamma_bn1, C_base,
        Hout_conv1, Wout_conv1, epsilon_bn)
    [dX_res, dW_conv1] = conv3x3_backward(dX_res, Hout_conv1, Wout_conv1, X, W_conv1, C_in, C_base, Hin, Win,
        strideh, stridew)

    gradients = list(dW_conv1, dgamma_bn1, dbeta_bn1, dW_conv2, dgamma_bn2, dbeta_bn2)
    if (downsample) {
        gradients = append(gradients, dW_conv3)
        gradients = append(gradients, dgamma_bn3)
        gradients = append(gradients, dbeta_bn3)
    }

    # Combine res path and id path
    dX = dX_res + dX_id
}

bottleneck_block_forward = function(matrix[double] X,
        list[unknown] weights, int C_in, int C_base, int Hin,
        int Win, int strideh, int stridew, string mode,
        list[unknown] ema_means_vars)
    return (matrix[double] out, int Hout, int Wout,
            list[unknown] ema_means_vars_upd,
            list[unknown] cached_out,
            list[unknown] cached_means_vars) {
    /*
     * Computes the forward pass for a bottleneck residual
     * block.
     * This residual block architecture is used in the
     * bigger ResNets. They consist of 3 convolutional
     * layers - 1x1, 3x3, 1x1. The last layer increases
     * the number of channels by a factor of 4 which is
     * downscaled by the first layer of the next res block
     * to always keep computational complexity at a
     * minimum.
     * The downsampling of the image dimensions (Hin & Win)
     * through the stride is placed at the 3x3 conv layer
     * instead of the first 1x1 layer (as proposed in "Deep
     * residual learning for image recognition") which is
     * called ResNet V1.5 and introduced in
     * https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
     *
     * Inputs:
     * - X: Inputs, of shape (N, C_in*Hin*Win).
     * - weights: list of weights for all layers of res block
     *     with the following order/content:
     *   -> 1: Weights of conv 1, of shape (C_base, C_in*1*1).
     *   -> 2: Weights of batch norm 1, of shape (C_base, 1).
     *   -> 3: Bias of batch norm 1, of shape (C_base, 1).
     *   -> 4: Weights of conv 2, of shape (C_base, C_base*3*3).
     *   -> 5: Weights of batch norm 2, of shape (C_base, 1).
     *   -> 6: Bias of batch norm 2, of shape (C_base, 1).
     *   -> 7: Weights of conv 3, of shape (4*C_base, C_base*1*1).
     *   -> 8: Weights of batch norm 3, of shape (4*C_base, 1).
     *   -> 9: Bias of batch norm 3, of shape (4*C_base, 1).
     *   If the block should downsample X:
     *   -> 10: Weights of downsample conv, of shape (4*C_base, C_in*1*1).
     *   -> 11: Weights of downsample batch norm, of shape (4*C_base, 1).
     *   -> 12: Bias of downsample batch norm, of shape (4*C_base, 1).
     * - C_in: Number of input channels.
     * - C_base: Number of base channels for this block.
     * - Hin: Input height.
     * - Win: Input width.
     * - strideh: Stride over height (usually 1 or 2).
     * - stridew: Stride over width (usually same as strideh).
     * - mode: 'train' or 'test' to indicate if the model is currently
     *     being trained or tested for badge normalization layers.
     *     See badge_norm2d.dml docs for more info.
     * - ema_means_vars: List of exponential moving averages for mean
     *     and variance for badge normalization layers.
     *   -> 1: EMA for mean of badge norm 1, of shape (C_base, 1).
     *   -> 2: EMA for variance of badge norm 1, of shape (C_base, 1).
     *   -> 3: EMA for mean of badge norm 2, of shape (C_base, 1).
     *   -> 4: EMA for variance of badge norm 2, of shape (C_base, 1).
     *   -> 5: EMA for mean of badge norm 3, of shape (4*C_base, 1).
     *   -> 6: EMA for variance of badge norm 3, of shape (4*C_base, 1).
     *   If the block should downsample X:
     *   -> 7: EMA for mean of downs. badge norm, of shape (4*C_base, 1).
     *   -> 8: EMA for variance of downs. badge norm, of shape (4*C_base, 1).
     *
     * Outputs:
     * - out: Output, of shape (N, 4*C_base*Hout*Wout).
     * - Hout: Output height.
     * - Wout: Output width.
     * - ema_means_vars_upd: List of updated exponential moving averages
     *     for mean and variance of badge normalization layers.
     * - cached_out: Outputs of each layer for computation of backward
     *     pass. Refer to the code for the order of elements.
     * - cached_means_vars: List of cached means and vars returned from
     *     each batch normalization layer. This is required for the
     *     backward pass of the network.
     */
    downsample = strideh > 1 | stridew > 1 | C_in != 4 * C_base
    # default values
    mu_bn = 0.1
    epsilon_bn = 1e-05

    # get all params
    W_conv1 = as.matrix(weights[1])
    gamma_bn1 = as.matrix(weights[2])
    beta_bn1 = as.matrix(weights[3])
    W_conv2 = as.matrix(weights[4])
    gamma_bn2 = as.matrix(weights[5])
    beta_bn2 = as.matrix(weights[6])
    W_conv3 = as.matrix(weights[7])
    gamma_bn3 = as.matrix(weights[8])
    beta_bn3 = as.matrix(weights[9])

    ema_mean_bn1 = as.matrix(ema_means_vars[1])
    ema_var_bn1 = as.matrix(ema_means_vars[2])
    ema_mean_bn2 = as.matrix(ema_means_vars[3])
    ema_var_bn2 = as.matrix(ema_means_vars[4])
    ema_mean_bn3 = as.matrix(ema_means_vars[5])
    ema_var_bn3 = as.matrix(ema_means_vars[6])

    if (downsample) {
        # gather params for donwsampling
        W_conv4 = as.matrix(weights[10])
        gamma_bn4 = as.matrix(weights[11])
        beta_bn4 = as.matrix(weights[12])
        ema_mean_bn4 = as.matrix(ema_means_vars[7])
        ema_var_bn4 = as.matrix(ema_means_vars[8])
    }

    # RESIDUAL PATH
    # First convolutional layer
    [out_conv1, Hout_conv1, Wout_conv1] = conv1x1_forward(X, W_conv1, C_in, C_base, Hin, Win, 1, 1)
    [out_bn1, ema_mean_bn1_upd, ema_var_bn1_upd, c_m_bn1, c_v_bn1] = bn2d::forward(out_conv1, gamma_bn1, beta_bn1,
        C_base, Hout_conv1, Wout_conv1, mode, ema_mean_bn1, ema_var_bn1, mu_bn, epsilon_bn)
    out_re1 = relu::forward(out_bn1)

    # Second convolutional layer
    [out_conv2, Hout_conv2, Wout_conv2] = conv3x3_forward(out_re1, W_conv2, C_base, C_base, Hout_conv1, Wout_conv1,
        strideh, stridew)
    [out_bn2, ema_mean_bn2_upd, ema_var_bn2_upd, c_m_bn2, c_v_bn2] = bn2d::forward(out_conv2, gamma_bn2, beta_bn2,
        C_base, Hout_conv2, Wout_conv2, mode, ema_mean_bn2, ema_var_bn2, mu_bn, epsilon_bn)
    out_re2 = relu::forward(out_bn2)

    # Third convolutional layer
    [out_conv3, Hout_conv3, Wout_conv3] = conv1x1_forward(out_re2, W_conv3, C_base, 4*C_base, Hout_conv2, Wout_conv2, 1,
        1)
    [out_bn3, ema_mean_bn3_upd, ema_var_bn3_upd, c_m_bn3, c_v_bn3] = bn2d::forward(out_conv3, gamma_bn3, beta_bn3,
        4*C_base, Hout_conv3, Wout_conv3, mode, ema_mean_bn3, ema_var_bn3, mu_bn, epsilon_bn)

    # IDENTITY PATH
    if (downsample) {
        # downsample input
        [out_conv4, Hout_conv4, Wout_conv4] = conv1x1_forward(X, W_conv4, C_in, 4*C_base, Hin, Win, strideh, stridew)
        [out_bn4, ema_mean_bn4_upd, ema_var_bn4_upd, c_m_bn4, c_v_bn4] = bn2d::forward(out_conv4, gamma_bn4, beta_bn4,
            4*C_base, Hout_conv4, Wout_conv4, mode, ema_mean_bn4, ema_var_bn4, mu_bn, epsilon_bn)
        identity = out_bn4
    } else {
        identity = X
    }

    out_re3 = relu::forward(out_bn3 + identity)

    ema_means_vars_upd = list(ema_mean_bn1_upd, ema_var_bn1_upd, ema_mean_bn2_upd, ema_var_bn2_upd, ema_mean_bn3_upd,
        ema_var_bn3_upd)
    cached_out = list(X, Hin, Win, out_conv1, Hout_conv1, Wout_conv1, out_bn1, out_re1, out_conv2, Hout_conv2,
        Wout_conv2, out_bn2, out_re2, out_conv3, Hout_conv3, Wout_conv3, out_bn3, out_re3)
    cached_means_vars = list(c_m_bn1, c_v_bn1, c_m_bn2, c_v_bn2, c_m_bn3, c_v_bn3)
    if (downsample) {
        ema_means_vars_upd = append(ema_means_vars_upd, ema_mean_bn3_upd)
        ema_means_vars_upd = append(ema_means_vars_upd, ema_var_bn3_upd)
        cached_out = append(cached_out, out_conv4)
        cached_out = append(cached_out, Hout_conv4)
        cached_out = append(cached_out, Wout_conv4)
        cached_out = append(cached_out, out_bn4)
        cached_means_vars = append(cached_means_vars, c_m_bn4)
        cached_means_vars = append(cached_means_vars, c_v_bn4)
    }

    out = out_re3
    Hout = Hout_conv3
    Wout = Wout_conv3
}

bottleneck_block_backward = function(matrix[double] dOut, list[unknown] cached_out,
                                     list[unknown] weights, int C_in, int C_base,
                                     int strideh, int stridew,
                                     list[unknown] cached_means_vars)
    return (matrix[double] dX, list[unknown] gradients) {
    downsample = strideh > 1 | stridew > 1 | C_in != 4*C_base
    # default values
    mu_bn = 0.1
    epsilon_bn = 1e-05

    # get all params
    W_conv1 = as.matrix(weights[1])
    gamma_bn1 = as.matrix(weights[2])
    beta_bn1 = as.matrix(weights[3])
    W_conv2 = as.matrix(weights[4])
    gamma_bn2 = as.matrix(weights[5])
    beta_bn2 = as.matrix(weights[6])
    W_conv3 = as.matrix(weights[7])
    gamma_bn3 = as.matrix(weights[8])
    beta_bn3 = as.matrix(weights[9])

    # cached outputs
    X = as.matrix(cached_out[1])
    Hin = as.integer(as.scalar(cached_out[2]))
    Win = as.integer(as.scalar(cached_out[3]))
    out_conv1 = as.matrix(cached_out[4])
    Hout_conv1 = as.integer(as.scalar(cached_out[5]))
    Wout_conv1 = as.integer(as.scalar(cached_out[6]))
    out_bn1 = as.matrix(cached_out[7])
    out_re1 = as.matrix(cached_out[8])
    out_conv2 = as.matrix(cached_out[9])
    Hout_conv2 = as.integer(as.scalar(cached_out[10]))
    Wout_conv2 = as.integer(as.scalar(cached_out[11]))
    out_bn2 = as.matrix(cached_out[12])
    out_re2 = as.matrix(cached_out[13])
    out_conv3 = as.matrix(cached_out[14])
    Hout_conv3 = as.integer(as.scalar(cached_out[15]))
    Wout_conv3 = as.integer(as.scalar(cached_out[16]))
    out_bn3 = as.matrix(cached_out[17])
    out_re3 = as.matrix(cached_out[18])

    # cached means and vars
    mean_bn1 = as.matrix(cached_means_vars[1])
    var_bn1 = as.matrix(cached_means_vars[2])
    mean_bn2 = as.matrix(cached_means_vars[3])
    var_bn2 = as.matrix(cached_means_vars[4])
    mean_bn3 = as.matrix(cached_means_vars[5])
    var_bn3 = as.matrix(cached_means_vars[6])

    if (downsample) {
        # gather params for downsampling
        W_conv4 = as.matrix(weights[10])
        gamma_bn4 = as.matrix(weights[11])
        beta_bn4 = as.matrix(weights[12])

        out_conv4 = as.matrix(cached_out[19])
        Hout_conv4 = as.integer(as.scalar(cached_out[20]))
        Wout_conv4 = as.integer(as.scalar(cached_out[21]))
        out_bn4 = as.matrix(cached_out[22])
        identity = out_bn4

        mean_bn4 = as.matrix(cached_means_vars[7])
        var_bn4 = as.matrix(cached_means_vars[8])
    } else {
        identity = X
    }

    dX = relu::backward(dOut, identity + out_bn3)

    # IDENTITY PATH
    if (downsample) {
        [dX_id, dgamma_bn4, dbeta_bn4] = bn2d::backward(dX, mean_bn4, var_bn4, out_conv4, gamma_bn4, 4*C_base,
            Hout_conv4, Wout_conv4, epsilon_bn)
        [dX_id, dW_conv4] = conv1x1_backward(dX_id, Hout_conv4, Wout_conv4, X, W_conv4, C_in, 4*C_base, Hin,
            Win, strideh, stridew)
    } else {
        dX_id = dX
    }

    # RESIDUAL PATH
    # Third convolutional layer
    [dX_res, dgamma_bn3, dbeta_bn3] = bn2d::backward(dX, mean_bn3, var_bn3, out_conv3, gamma_bn3, 4*C_base, Hout_conv3,
        Wout_conv3, epsilon_bn)
    [dX_res, dW_conv3] = conv1x1_backward(dX_res, Hout_conv3, Wout_conv3, out_re2, W_conv3, C_base, 4*C_base,
        Hout_conv2, Wout_conv2, 1, 1)

    # Second convolutional layer
    dX_res = relu::backward(dX_res, out_bn2)
    [dX_res, dgamma_bn2, dbeta_bn2] = bn2d::backward(dX_res, mean_bn2, var_bn2, out_conv2, gamma_bn2, C_base,
        Hout_conv2, Wout_conv2, epsilon_bn)
    [dX_res, dW_conv2] = conv3x3_backward(dX_res, Hout_conv2, Wout_conv2, out_re1, W_conv2, C_base, C_base,
        Hout_conv1, Wout_conv1, strideh, stridew)

    # First convoluational layer
    dX_res = relu::backward(dX_res, out_bn1)
    [dX_res, dgamma_bn1, dbeta_bn1] = bn2d::backward(dX_res, mean_bn1, var_bn1, out_conv1, gamma_bn1, C_base,
        Hout_conv1, Wout_conv1, epsilon_bn)
    [dX_res, dW_conv1] = conv1x1_backward(dX_res, Hout_conv1, Wout_conv1, X, W_conv1, C_in, C_base, Hin, Win,
        1, 1)

    gradients = list(dW_conv1, dgamma_bn1, dbeta_bn1, dW_conv2, dgamma_bn2, dbeta_bn2, dW_conv3, dgamma_bn3, dbeta_bn3)
    if (downsample) {
        gradients = append(gradients, dW_conv4)
        gradients = append(gradients, dgamma_bn4)
        gradients = append(gradients, dbeta_bn4)
    }

    # Combine res path and id path
    dX = dX_res + dX_id
}

reslayer_forward = function(matrix[double] X, int Hin, int Win,
                            string block_type, int blocks, int strideh,
                            int stridew, int C_in, int C_base,
                            list[unknown] blocks_weights, string mode,
                            list[unknown] ema_means_vars)
    return (matrix[double] out, int Hout, int Wout,
            list[unknown] ema_means_vars_upd,
            list[unknown] cached_out,
            list[unknown] cached_means_vars) {
    /*
     * Executes the forward pass for a sequence of residual
     * blocks with the same number of base channels, i.e.
     * a residual layer.
     *
     * Inputs:
     * - X: Inputs, of shape (N, C_in*Hin*Win)
     * - Hin: Input height.
     * - Win: Input width.
     * - block_type: 'basic' or 'bottleneck' depending on
     *     which type of block should be used for the
     *     residual layer.
     * - blocks: Number of residual blocks (bigger than 0).
     * - strideh: Stride height for first conv layer of first block.
     * - stridew: Stride width for first conv layer of first block.
     * - C_in: Number of input channels.
     * - C_base: Number of base channels of res layer.
     * - blocks_weights: List of weights of each block.
     *     -> i: List of weights of block i with the content
     *           defined in the docs of basic_block_forward()
     *           or bottleneck_block_forward() depending on
     *           the block type.
     *     -> length == blocks
     * - mode: 'train' or 'test' to indicate if the model is currently
     *     being trained or tested for badge normalization layers.
     *     See badge_norm2d.dml docs for more info.
     * - ema_means_vars: List of exponential moving averages for mean
     *     and variance for badge normalization layers of each block.
     *     -> i: List of EMAs of block i with the content defined
     *           in the docs of basic_block_forward() or
     *           bottleneck_block_forward() depending on the block type.
     *     -> length == blocks
     * - cached_out: Outputs of each layer for computation of backward
     *     pass. Refer to the code for the order of elements.
     * - cached_means_vars: List of cached means and vars returned from
     *     each batch normalization layer. This is required for the
     *     backward pass of the network.
     */
    # first block with provided stride
    if (block_type == "basic") {
        [out, Hout, Wout, emas1_upd, cached_out_b1, cached_mv_b1] = basic_block_forward(X, as.list(blocks_weights[1]),
            C_in, C_base, Hin, Win, strideh, stridew, mode, as.list(ema_means_vars[1]))
    } else {
        [out, Hout, Wout, emas1_upd, cached_out_b1, cached_mv_b1] = bottleneck_block_forward(X,
            as.list(blocks_weights[1]), C_in, C_base, Hin, Win, strideh, stridew, mode, as.list(ema_means_vars[1]))
    }
    ema_means_vars_upd = list(emas1_upd)
    cached_out = list(cached_out_b1)
    cached_means_vars = list(cached_mv_b1)

    # other blocks
    if (blocks > 1) {
        for (i in 2:blocks) {
            current_weights = as.list(blocks_weights[i])
            current_emas = as.list(ema_means_vars[i])

            if (block_type == "basic") {
                [out, Hout, Wout, current_emas_upd, current_cached_out, current_cached_mv] = basic_block_forward(X=out,
                    weights=current_weights, C_in=C_base, C_base=C_base, Hin=Hout, Win=Wout, strideh=1, stridew=1,
                    mode=mode, ema_means_vars=current_emas)
            } else {
                [out, Hout, Wout, current_emas_upd, current_cached_out, current_cached_mv] = bottleneck_block_forward(X=out,
                    weights=current_weights, C_in=C_base*4, C_base=C_base, Hin=Hout, Win=Wout, strideh=1, stridew=1,
                    mode=mode, ema_means_vars=current_emas)
            }

            ema_means_vars_upd = append(ema_means_vars_upd, current_emas_upd)
            cached_out = append(cached_out, current_cached_out)
            cached_means_vars = append(cached_means_vars, current_cached_mv)
        }
    }
}

reslayer_backward = function(matrix[double] dOut, list[unknown] cached_out,
                             string block_type, int blocks, int strideh,
                             int stridew, int C_in, int C_base,
                             list[unknown] weights,
                             list[unknown] cached_means_vars)
    return (matrix[double] dX, list[unknown] gradients) {
    dX = dOut

    # second to last blocks
    if (blocks > 1) {
        gradients_reverse = list()
        for (i in blocks:2) {
            current_cached_out = as.list(cached_out[i])
            current_weights = as.list(weights[i])
            current_cached_means_vars = as.list(cached_means_vars[i])

            if (block_type == "basic") {
                [dX, current_gradients] = basic_block_backward(dOut=dX, cached_out=current_cached_out,
                    weights=current_weights, C_in=C_base, C_base=C_base, strideh=1, stridew=1,
                    cached_means_vars=current_cached_means_vars)
            } else {
                [dX, current_gradients] = bottleneck_block_backward(dOut=dX, cached_out=current_cached_out,
                    weights=current_weights, C_in=4*C_base, C_base=C_base, strideh=1, stridew=1,
                    cached_means_vars=current_cached_means_vars)
            }

            gradients_reverse = append(gradients_reverse, current_gradients)
        }
    }

    # first block
    if (block_type == "basic") {
        [dX, current_gradients] = basic_block_backward(dOut=dX, cached_out=as.list(cached_out[1]),
            weights=as.list(weights[1]), C_in=C_in, C_base=C_base, strideh=strideh, stridew=stridew,
            cached_means_vars=as.list(cached_means_vars[1]))
    } else {
        [dX, current_gradients] = bottleneck_block_backward(dOut=dX, cached_out=as.list(cached_out[1]),
            weights=as.list(weights[1]), C_in=C_in, C_base=C_base, strideh=strideh, stridew=stridew,
            cached_means_vars=as.list(cached_means_vars[1]))
    }

    # sort gradient lists of blocks in correct order
    gradients = list(current_gradients)
    if (blocks > 1) {
        for (i in blocks:2) {
            gradients = append(gradients, as.list(gradients_reverse[i - 1]))
        }
    }
}

resnet_forward = function(matrix[double] X, int Hin, int Win,
                          string block_type, list[unknown] layer_sizes,
                          list[unknown] model, string mode,
                          list[unknown] ema_means_vars)
    return (matrix[double] out, list[unknown] ema_means_vars_upd,
            list[unknown] cached_out, list[unknown] cached_means_vars) {
    /*
     * Forward pass of the ResNet as introduced in
     * "Deep Residual Learning for Image Recognition" by
     * Kaiming He et. al. and inspired by the PyTorch
     * implementation.
     *
     * Inputs:
     * - X: Inputs, of shape (N, C_in*Hin*Win).
     *     C_in = 3 is expected.
     * - Hin: Input height.
     * - Win: Input width.
     * - block_type: 'basic' or 'bottleneck' depending on
     *     which type of block should be used for the
     *     residual network.
     * - layer_sizes: List of the sizes of each of
     *     the 4 residual layers.
     *     For ResNet18: [2, 2, 2, 2], RN34: [3, 4, 6, 3],
     *     RN50: [3, 4, 6, 3], RN101: [3, 4, 23, 3],
     *     RN152: [3, 8, 36, 3]
     * - model: Weights and bias matrices of the model
     *     with the following order/content:
     *   -> 1: Weights of conv 1 7x7, of shape (64, 3*7*7)
     *   -> 2: Weights of batch norm 1, of shape (64, 1).
     *   -> 3: Bias of batch norm 1, of shape (64, 1).
     *   -> 4: List of weights for first residual layer
     *         with 64 base channels.
     *   -> 5: List of weights for second residual layer
     *         with 128 base channels.
     *   -> 6: List of weights for third residual layer
     *         with 256 base channels.
     *   -> 7: List of weights for fourth residual layer
     *         with 512 base channels.
     *      List of residual layers 1, 2, 3 & 4 have
     *      the content/order:
     *      -> i: List of weights for residual block i.
     *            with i in {1, ..., layer_sizes[layer]}
     *         Each list of weights for a residual block
     *         must follow the same order as defined in
     *         the documentation of basic_block_forward()
     *         or bottleneck_block_forward() depending
     *         on the block type.
     *   -> 8: Weights of fully connected layer, of shape (C_out, classes)
     *         where C_out = 512 for basic block type and C_out = 2048
     *         for bottleneck block type.
     *   -> 9: Bias of fully connected layer, of shape (1, classes)
     * - mode: 'train' or 'test' to indicate if the model is currently
     *     being trained or tested for badge normalization layers.
     *     See badge_norm2d.dml docs for more info.
     * - ema_means_vars: List of exponential moving averages for mean
     *     and variance for badge normalization layers.
     *   -> 1: EMA for mean of badge norm 1, of shape (64, 1).
     *   -> 2: EMA for variance of badge norm 1, of shape (64, 1).
     *   -> 3: List of EMA means and vars for residual layer 1.
     *   -> 4: List of EMA means and vars for residual layer 2.
     *   -> 5: List of EMA means and vars for residual layer 3.
     *   -> 6: List of EMA means and vars for residual layer 4.
     *      Lists for EMAs of layer 1, 2, 3 & 4 must have the
     *      following order:
     *      -> i: List of EMA means and vars for residual block i.
     *            with i in {1, ..., layer_sizes[layer]}
     *         Each list of EMAs for a residual block
     *         must follow the same order as defined in
     *         the documentation of basic_block_forward()
     *         or bottleneck_block_forward().
     * - NOTICE: The lists of the first blocks for layer 2, 3 and 4
     *           must include weights and EMAs for 1 extra conv layer
     *           and a batch norm layer for the downsampling on the
     *           identity path.
     *
     * Outputs:
     * - out: Outputs, of shape (N, classes)
     * - ema_means_vars_upd: List of updated exponential moving averages
     *     for mean and variance of badge normalization layers. It follows
     *     the same exact structure as the input EMAs list.
     * - cached_out: Outputs of each layer for computation of backward
     *     pass. Refer to the code for the order of elements.
     * - cached_means_vars: List of cached means and vars returned from
     *     each batch normalization layer. This is required for the
     *     backward pass of the network.
     */
    # default values
    mu_bn = 0.1
    epsilon_bn = 1e-05

    if (block_type == "basic") {
        block_expansion = 1
    } else {
        block_expansion = 4
    }

    # extract model params
    W_conv1 = as.matrix(model[1])
    gamma_bn1 = as.matrix(model[2]); beta_bn1 = as.matrix(model[3])
    weights_reslayer1 = as.list(model[4])
    weights_reslayer2 = as.list(model[5])
    weights_reslayer3 = as.list(model[6])
    weights_reslayer4 = as.list(model[7])
    W_fc = as.matrix(model[8])
    b_fc = as.matrix(model[9])
    ema_mean_bn1 = as.matrix(ema_means_vars[1]); ema_var_bn1 = as.matrix(ema_means_vars[2])
    emas_reslayer1 = as.list(ema_means_vars[3])
    emas_reslayer2 = as.list(ema_means_vars[4])
    emas_reslayer3 = as.list(ema_means_vars[5])
    emas_reslayer4 = as.list(ema_means_vars[6])

    # Convolutional 7x7 layer
    C = 64
    b_conv1 = matrix(0, rows=C, cols=1)
    [out_conv1, Hout_conv1, Wout_conv1] = conv2d::forward(X=X, W=W_conv1, b=b_conv1, C=3, Hin=Hin, Win=Win, Hf=7, Wf=7,
        strideh=2, stridew=2, padh=3, padw=3)
    # Batch Normalization
    [out_bn1, ema_mean_bn1_upd, ema_var_bn1_upd, cached_m, cached_v] = bn2d::forward(X=out_conv1, gamma=gamma_bn1,
        beta=beta_bn1, C=C, Hin=Hout_conv1, Win=Wout_conv1, mode=mode, ema_mean=ema_mean_bn1, ema_var=ema_var_bn1,
        mu=mu_bn, epsilon=epsilon_bn)
    # ReLU
    out_re1 = relu::forward(X=out_bn1)
    # Max Pooling 3x3
    [out_mp, Hout_mp, Wout_mp] = mp2d::forward(X=out_re1, C=C, Hin=Hout_conv1, Win=Wout_conv1, Hf=3, Wf=3, strideh=2,
        stridew=2, padh=1, padw=1)

    # residual layer 1
    block_count = as.integer(as.scalar(layer_sizes[1]))
    [out, Hout, Wout, emas1_upd, cached_out_l1, cached_mv_l1] = reslayer_forward(X=out_mp, Hin=Hout_mp, Win=Wout_mp,
        block_type=block_type, blocks=block_count, strideh=1, stridew=1, C_in=C, C_base=64,
        blocks_weights=weights_reslayer1, mode=mode, ema_means_vars=emas_reslayer1)
    C = 64 * block_expansion
    # residual layer 2
    block_count = as.integer(as.scalar(layer_sizes[2]))
    [out, Hout, Wout, emas2_upd, cached_out_l2, cached_mv_l2] = reslayer_forward(X=out, Hin=Hout, Win=Wout,
        block_type=block_type, blocks=block_count, strideh=2, stridew=2, C_in=C, C_base=128,
        blocks_weights=weights_reslayer2, mode=mode, ema_means_vars=emas_reslayer2)
    C = 128 * block_expansion
    # residual layer 3
    block_count = as.integer(as.scalar(layer_sizes[3]))
    [out, Hout, Wout, emas3_upd, cached_out_l3, cached_mv_l3] = reslayer_forward(X=out, Hin=Hout, Win=Wout,
        block_type=block_type, blocks=block_count, strideh=2, stridew=2, C_in=C, C_base=256,
        blocks_weights=weights_reslayer3, mode=mode, ema_means_vars=emas_reslayer3)
    C = 256 * block_expansion
    # residual layer 4
    block_count = as.integer(as.scalar(layer_sizes[4]))
    [out_res, Hout_res, Wout_res, emas4_upd, cached_out_l4, cached_mv_l4] = reslayer_forward(X=out, Hin=Hout, Win=Wout,
        block_type=block_type, blocks=block_count, strideh=2, stridew=2, C_in=C, C_base=512,
        blocks_weights=weights_reslayer4, mode=mode, ema_means_vars=emas_reslayer4)
    C = 512 * block_expansion

    # Global Average Pooling
    [out_ap, Hout_ap, Wout_ap] = ap2d::forward(X=out_res, C=C, Hin=Hout_res, Win=Wout_res)
    # Affine
    out_fc = fc::forward(X=out_ap, W=W_fc, b=b_fc)

    ema_means_vars_upd = list(ema_mean_bn1_upd, ema_var_bn1_upd, emas1_upd, emas2_upd, emas3_upd, emas4_upd)
    cached_out = list(X, Hin, Win, out_conv1, Hout_conv1, Wout_conv1, out_bn1, out_re1, out_mp, Hout_mp, Wout_mp,
        cached_out_l1, cached_out_l2, cached_out_l3, cached_out_l4, out_res, Hout_res, Wout_res, out_ap, Hout_ap,
        Wout_ap, out_fc)
    cached_means_vars = list(cached_m, cached_v, cached_mv_l1, cached_mv_l2, cached_mv_l3, cached_mv_l4)

    out = out_fc
}

resnet_backward = function(matrix[double] dOut, list[unknown] cached_out,
                           string block_type, list[unknown] layer_sizes,
                           list[unknown] model, list[unknown] cached_means_vars)
    return (matrix[double] dX, list[unknown] gradients) {
    # default values
    epsilon_bn = 1e-05

    if (block_type == "basic") {
        block_expansion = 1
    } else {
        block_expansion = 4
    }

    # extract model params
    W_conv1 = as.matrix(model[1])
    gamma_bn1 = as.matrix(model[2]); beta_bn1 = as.matrix(model[3])
    weights_reslayer1 = as.list(model[4])
    weights_reslayer2 = as.list(model[5])
    weights_reslayer3 = as.list(model[6])
    weights_reslayer4 = as.list(model[7])
    W_fc = as.matrix(model[8])
    b_fc = as.matrix(model[9])

    # extract cached output
    X = as.matrix(cached_out[1])
    Hin = as.integer(as.scalar(cached_out[2]))
    Win = as.integer(as.scalar(cached_out[3]))
    out_conv1 = as.matrix(cached_out[4])
    Hout_conv1 = as.integer(as.scalar(cached_out[5]))
    Wout_conv1 = as.integer(as.scalar(cached_out[6]))
    out_bn1 = as.matrix(cached_out[7])
    out_re1 = as.matrix(cached_out[8])
    out_mp = as.matrix(cached_out[9])
    Hout_mp = as.integer(as.scalar(cached_out[10]))
    Wout_mp = as.integer(as.scalar(cached_out[11]))
    cached_out_l1 = as.list(cached_out[12])
    cached_out_l2 = as.list(cached_out[13])
    cached_out_l3 = as.list(cached_out[14])
    cached_out_l4 = as.list(cached_out[15])
    out_res = as.matrix(cached_out[16])
    Hout_res = as.integer(as.scalar(cached_out[17]))
    Wout_res = as.integer(as.scalar(cached_out[18]))
    out_ap = as.matrix(cached_out[19])
    Hout_ap = as.integer(as.scalar(cached_out[20]))
    Wout_ap = as.integer(as.scalar(cached_out[21]))

    # extract cached means and vars for bn layers
    cached_m_bn1 = as.matrix(cached_means_vars[1])
    cached_v_bn1 = as.matrix(cached_means_vars[2])
    cached_mv_l1 = as.list(cached_means_vars[3])
    cached_mv_l2 = as.list(cached_means_vars[4])
    cached_mv_l3 = as.list(cached_means_vars[5])
    cached_mv_l4 = as.list(cached_means_vars[6])

    # Affine
    [dX_fc, dW_fc, db_fc] = fc::backward(dout=dOut, X=out_ap, W=W_fc, b=b_fc)
    # Global Average Pooling
    dX_ap = ap2d::backward(dout=dX_fc, X=out_res, C=512*block_expansion, Hin=Hout_res, Win=Wout_res)

    # residual layer 4
    block_count = as.integer(as.scalar(layer_sizes[4]))
    [dX_l4, gradients_l4] = reslayer_backward(dOut=dX_ap, cached_out=cached_out_l4, block_type=block_type,
        blocks=block_count, strideh=2, stridew=2, C_in=256*block_expansion, C_base=512,
        weights=weights_reslayer4, cached_means_vars=cached_mv_l4)
    # residual layer 3
    block_count = as.integer(as.scalar(layer_sizes[3]))
    [dX_l3, gradients_l3] = reslayer_backward(dOut=dX_l4, cached_out=cached_out_l3, block_type=block_type,
        blocks=block_count, strideh=2, stridew=2, C_in=128*block_expansion, C_base=256,
        weights=weights_reslayer3, cached_means_vars=cached_mv_l3)
    # residual layer 2
    block_count = as.integer(as.scalar(layer_sizes[2]))
    [dX_l2, gradients_l2] = reslayer_backward(dOut=dX_l3, cached_out=cached_out_l2, block_type=block_type,
        blocks=block_count, strideh=2, stridew=2, C_in=64*block_expansion, C_base=128, weights=weights_reslayer2,
        cached_means_vars=cached_mv_l2)
    # residual layer 1
    block_count = as.integer(as.scalar(layer_sizes[1]))
    [dX_l1, gradients_l1] = reslayer_backward(dOut=dX_l2, cached_out=cached_out_l1, block_type=block_type,
        blocks=block_count, strideh=1, stridew=1, C_in=64, C_base=64, weights=weights_reslayer1,
        cached_means_vars=cached_mv_l1)

    # Max Pooling 3x3
    dX_mp = mp2d::backward(dout=dX_l1, Hout=Hout_mp, Wout=Wout_mp, X=out_re1, C=64, Hin=Hout_conv1, Win=Wout_conv1,
        Hf=3, Wf=3, strideh=2, stridew=2, padh=1, padw=1)
    # ReLU
    dX_re1 = relu::backward(dout=dX_mp, X=out_bn1)
    # Batch Norm
    [dX_bn1, dgamma_bn1, dbeta_bn1] = bn2d::backward(dout=dX_re1, cache_mean=cached_m_bn1, cache_inv_var=cached_v_bn1,
        X=out_conv1, gamma=gamma_bn1, C=64, Hin=Hout_conv1, Win=Wout_conv1, epsilon=epsilon_bn)
    # Convolutional 7x7
    b_conv1 = matrix(0, rows=64, cols=1)
    [dX_conv1, dW_conv1, db_conv1] = conv2d::backward(dout=dX_bn1, Hout=Hout_conv1, Wout=Wout_conv1, X=X, W=W_conv1,
        b=b_conv1, C=3, Hin=Hin, Win=Win, Hf=7, Wf=7, strideh=2, stridew=2, padh=3, padw=3)

    gradients = list(dW_conv1, dgamma_bn1, dbeta_bn1, gradients_l1, gradients_l2, gradients_l3, gradients_l4, dW_fc,
        db_fc)
    dX = dX_conv1
}

init = function(int classes, string block_type, list[unknown] layer_sizes, int seed)
    return (list[unknown] model, list[unknown] emas) {
    C_in = 3
    # Conv 7x7
    [W_conv1, b_conv1] = conv2d::init(F=64, C=C_in, Hf=7, Wf=7, seed=seed)
    C_in = 64
    # Batch norm
    [gamma_bn1, beta_bn1, ema_m_bn1, ema_v_bn1] = bn2d::init(C_in)

    model = list(W_conv1, gamma_bn1, beta_bn1)
    emas = list(ema_m_bn1, ema_v_bn1)

    # residual layers
    C_bases = list(64, 128, 256, 512)
    strides = list(1, 2, 2, 2)
    for (layer in 1:4) {
        layer_size = as.integer(as.scalar(layer_sizes[layer]))
        C_base = as.integer(as.scalar(C_bases[layer]))
        stride = as.integer(as.scalar(strides[layer]))

        weights_layer = list()
        emas_layer = list()

        if (block_type == "basic") {
            # basic blocks
            for (block in 1:layer_size) {
                [t_W_conv1, t_b_conv1] = conv2d::init(F=C_base, C=C_in, Hf=3, Wf=3, seed=seed)
                [t_gamma_bn1, t_beta_bn1, t_ema_m_bn1, t_ema_v_bn1] = bn2d::init(C_base)
                [t_W_conv2, t_b_conv2] = conv2d::init(F=C_base, C=C_base, Hf=3, Wf=3, seed=seed)
                [t_gamma_bn2, t_beta_bn2, t_ema_m_bn2, t_ema_v_bn2] = bn2d::init(C_base)
                weights_block = list(t_W_conv1, t_gamma_bn1, t_beta_bn1, t_W_conv2, t_gamma_bn2, t_beta_bn2)
                emas_block = list(t_ema_m_bn1, t_ema_v_bn1, t_ema_m_bn2, t_ema_v_bn2)

                if (block == 1 & stride > 1) {
                    # downsample params
                    [t_W_conv3, t_b_conv3] = conv2d::init(F=C_base, C=C_in, Hf=1, Wf=1, seed=seed)
                    [t_gamma_bn3, t_beta_bn3, t_ema_m_bn3, t_ema_v_bn3] = bn2d::init(C_base)
                    weights_block = append(weights_block, t_W_conv3)
                    weights_block = append(weights_block, t_gamma_bn3)
                    weights_block = append(weights_block, t_beta_bn3)
                    emas_block = append(emas_block, t_ema_m_bn3)
                    emas_block = append(emas_block, t_ema_v_bn3)
                }

                C_in = C_base

                weights_layer = append(weights_layer, weights_block)
                emas_layer = append(emas_layer, emas_block)
            }
        } else {
            # bottleneck blocks
            for (block in 1:layer_size) {
                [t_W_conv1, t_b_conv1] = conv2d::init(F=C_base, C=C_in, Hf=1, Wf=1, seed=seed)
                [t_gamma_bn1, t_beta_bn1, t_ema_m_bn1, t_ema_v_bn1] = bn2d::init(C_base)
                [t_W_conv2, t_b_conv2] = conv2d::init(F=C_base, C=C_base, Hf=3, Wf=3, seed=seed)
                [t_gamma_bn2, t_beta_bn2, t_ema_m_bn2, t_ema_v_bn2] = bn2d::init(C_base)
                [t_W_conv3, t_b_conv3] = conv2d::init(F=4*C_base, C=C_base, Hf=1, Wf=1, seed=seed)
                [t_gamma_bn3, t_beta_bn3, t_ema_m_bn3, t_ema_v_bn3] = bn2d::init(4*C_base)
                weights_block = list(t_W_conv1, t_gamma_bn1, t_beta_bn1, t_W_conv2, t_gamma_bn2, t_beta_bn2, t_W_conv3,
                    t_gamma_bn3, t_beta_bn3)
                emas_block = list(t_ema_m_bn1, t_ema_v_bn1, t_ema_m_bn2, t_ema_v_bn2, t_ema_m_bn3, t_ema_v_bn3)

                if (block == 1) {
                    # downsample params
                    [t_W_conv4, t_b_conv4] = conv2d::init(F=4*C_base, C=C_in, Hf=1, Wf=1, seed=seed)
                    [t_gamma_bn4, t_beta_bn4, t_ema_m_bn4, t_ema_v_bn4] = bn2d::init(4*C_base)
                    weights_block = append(weights_block, t_W_conv4)
                    weights_block = append(weights_block, t_gamma_bn4)
                    weights_block = append(weights_block, t_beta_bn4)
                    emas_block = append(emas_block, t_ema_m_bn4)
                    emas_block = append(emas_block, t_ema_v_bn4)
                }

                C_in = 4*C_base

                weights_layer = append(weights_layer, weights_block)
                emas_layer = append(emas_layer, emas_block)
            }
        }

        model = append(model, weights_layer)
        emas = append(emas, emas_layer)
    }

    # affine
    [W_fc, b_fc] = fc::init(D=C_in, M=classes, seed)
    model = append(model, W_fc)
    model = append(model, b_fc)
}
