softmax_array Module Function

  module function softmax_array(input, dim) result(output)
    implicit none
    class(array_type), intent(in), target :: input
    integer, intent(in) :: dim
    type(array_type), pointer :: output

    integer :: i

    output => input%create_result()
    if(dim.eq.1)then
       do i = 1, size(input%val, 1)
          output%val(i, :) = exp(input%val(i, :) - maxval(input%val(i,:)))
          output%val(i, :) = output%val(i, :) / sum(output%val(i, :))
       end do
    elseif(dim.eq.2)then
       do i = 1, size(input%val, 2)
          output%val(:, i) = exp(input%val(:, i) - maxval(input%val(:, i)))
          output%val(:, i) = output%val(:, i) / sum(output%val(:, i))
       end do
    else
       call stop_program("softmax_array: Unsupported dimension")
    end if
    allocate(output%indices(1))
    output%indices(1) = dim

    output%get_partial_left => get_partial_softmax
    output%get_partial_left_val => get_partial_softmax_val
    output%get_partial_left_val_sum => get_partial_softmax_val_sum
    if(input%requires_grad)then
       output%requires_grad = .true.
       output%is_forward = input%is_forward
       output%operation = 'softmax'
       output%left_operand => input
       output%owns_left_operand = input%is_temporary
    end if

  end function softmax_array