get_partial_softmax_val Subroutine

  pure subroutine get_partial_softmax_val(this, upstream_grad, output)
    !! Get partial derivative of softmax activation (in-place version)
    implicit none
    class(array_type), intent(in) :: this
    real(real32), dimension(:,:), intent(in) :: upstream_grad
    real(real32), dimension(:,:), intent(out) :: output

    integer :: s, dim

    if(this%indices(1).eq.1)then
       dim = 2
    else
       dim = 1
    end if
    output = this%val * upstream_grad
    if(dim.eq.1)then
       do s = 1, size(this%val, 2)
          output(:, s) = output(:, s) - this%val(:, s) * sum(output(:, s))
       end do
    elseif(dim.eq.2)then
       do s = 1, size(this%val, 1)
          output(s, :) = output(s, :) - this%val(s, :) * sum(output(s, :))
       end do
    end if
  end subroutine get_partial_softmax_val