piecewise_array Module Function

  module function piecewise_array(input, gradient, limit) result(output)
    !! Apply piecewise activation function to input array
    implicit none

    ! Arguments
    class(array_type), intent(in), target :: input
    real(real32), intent(in) :: gradient
    real(real32), intent(in) :: limit
    type(array_type), pointer :: output
    type(array_type), pointer :: b_array

    output => input%create_result()
    where(input%val.ge.limit)
       output%val = gradient * (input%val - limit) + limit
    elsewhere(input%val.le.-limit)
       output%val = gradient * (input%val + limit) - limit
    elsewhere
       output%val = input%val
    end where

    output%get_partial_left => get_partial_piecewise
    output%get_partial_left_val => get_partial_piecewise_val
    if(input%requires_grad)then
       output%requires_grad = .true.
       output%is_forward = input%is_forward
       output%operation = 'piecewise'
       output%left_operand => input
       output%owns_left_operand = input%is_temporary
    end if
    allocate(b_array)
    b_array%is_sample_dependent = .false.
    b_array%requires_grad = .false.
    call b_array%allocate(array_shape=[2, 1])
    b_array%val(1,1) = gradient
    b_array%val(2,1) = limit
    output%right_operand => b_array
    output%owns_right_operand = .true.

  end function piecewise_array