LogSoftmax.cpp

Go to the documentation of this file.

//
// Copyright © 2019, 2024 Arm Ltd. All rights reserved.
// SPDX-License-Identifier: MIT
//
 
#include "LogSoftmax.hpp"
 
#include <armnnUtils/TensorUtils.hpp>
#include <armnn/utility/NumericCast.hpp>
 
#include <cmath>
 
namespace
{
 
inline bool ValidateAxis(int axis, unsigned int numDimensions)
{
    const int sNumDimensions = armnn::numeric_cast<int>(numDimensions);
    return axis < sNumDimensions && axis >= -sNumDimensions;
}
 
} // anonymous namespace
 
namespace armnn
{
 
void LogSoftmax(Decoder<float>& input,
                Encoder<float>& output,
                const TensorInfo& inputInfo,
                const LogSoftmaxDescriptor& descriptor)
{
    const unsigned int numDimensions = inputInfo.GetNumDimensions();
 
    ARMNN_THROW_INVALIDARG_MSG_IF_FALSE(ValidateAxis(descriptor.m_Axis, numDimensions),
                                        "Axis index is not in range [-numDimensions, numDimensions).");
 
    unsigned int uAxis = descriptor.m_Axis < 0  ?
        numDimensions - armnn::numeric_cast<unsigned int>(std::abs(descriptor.m_Axis)) :
        armnn::numeric_cast<unsigned int>(descriptor.m_Axis);
 
    const TensorShape& inputShape = inputInfo.GetShape();
    const unsigned int outerSize  = armnnUtils::GetNumElementsBetween(inputShape, 0, uAxis);
    const unsigned int axisSize   = inputShape[uAxis];
    const unsigned int innerSize  = armnnUtils::GetNumElementsBetween(inputShape,
                                                                      uAxis + 1,
                                                                      inputShape.GetNumDimensions());
 
    for (unsigned int outer = 0; outer < outerSize; ++outer)
    {
        for (unsigned int inner = 0; inner < innerSize; ++inner)
        {
            // Find max
            input[outer * axisSize * innerSize + inner];
            float maxValue = input.Get();
            for (unsigned int i = 1u; i < axisSize; ++i)
            {
                input[(outer * axisSize + i) * innerSize + inner];
                maxValue = std::max(maxValue, input.Get());
            }
 
            // Compute sum
            float sum = 0.0f;
            for (unsigned int i = 0u; i < axisSize; ++i)
            {
                input[(outer * axisSize + i) * innerSize + inner];
                sum += std::exp((input.Get() - maxValue) * descriptor.m_Beta);
            }
 
            // Compute log sum
            const float logSum = std::log(sum);
 
            // Compute result
            for (unsigned int i = 0u; i < axisSize; ++i)
            {
                const unsigned int index = (outer * axisSize + i) * innerSize + inner;
 
                input [index];
                output[index];
 
                output.Set((input.Get() - maxValue) * descriptor.m_Beta - logSum);
            }
        }
    }
}
 
} // namespace armnn