get_partial_batchnorm_right_val Subroutine

  pure subroutine get_partial_batchnorm_right_val(this, upstream_grad, output)
    !! Get partial derivative wrt params for batchnorm (subroutine version)
    implicit none

    class(array_type), intent(in) :: this
    real(real32), dimension(:,:), intent(in) :: upstream_grad
    real(real32), dimension(:,:), intent(out) :: output

    integer :: c, num_dims, num_elements
    real(real32), allocatable :: x_hat(:,:)
    real(real32) :: mu, var, eps
    integer, dimension(size(this%shape)) :: input_shape

    input_shape = this%left_operand%shape

    select type(this)
    type is (batchnorm_array_type)
       eps = this%epsilon
       num_dims = size(this%shape)
       num_elements = product(this%shape(1:num_dims - 1))

       output = 0._real32

       allocate(x_hat(num_elements, size(upstream_grad,2)))

       do c = 1, input_shape(num_dims)
          mu = this%mean(c)
          var = this%variance(c)

          ! Normalised input
          x_hat(:,:) = ( &
               this%left_operand%val((c-1)*num_elements+1:c*num_elements,:) - mu &
          ) / sqrt(var + eps)

          output(c,1) = &
               sum(upstream_grad((c-1)*num_elements+1:c*num_elements,:) * x_hat)
          output(c + input_shape(num_dims),1) = &
               sum(upstream_grad((c-1)*num_elements+1:c*num_elements,:))

       end do
    end select

  end subroutine get_partial_batchnorm_right_val