src/modules/nn/activations.cc¶

// Copyright 2026 The Tenzr Project Authors
//
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "tenzr/modules/nn/activations.h"

#include "tenzr/common/error.h"

namespace tenzr::nn {

  Array softmax(const Array& input, int64_t axis) {
    TENZR_CHECK(input.IsDefined(), "Input array to softmax is uninitialized");
    int64_t rank = static_cast<int64_t>(input.shape().Rank());
    TENZR_CHECK(rank > 0, "softmax requires input array with rank >= 1");
    int64_t ax = (axis < 0) ? (axis + rank) : axis;
    TENZR_CHECK(ax >= 0 && ax < rank, "softmax axis out of range");

    auto max_val = input.Max(std::vector<int64_t>{ax}, /*keepdims=*/true);
    auto shifted = input - max_val;
    auto exp_shifted = shifted.Exp();
    auto sum_exp = exp_shifted.Sum(std::vector<int64_t>{ax}, /*keepdims=*/true);
    return exp_shifted / sum_exp;
  }

  Array log_softmax(const Array& input, int64_t axis) {
    TENZR_CHECK(input.IsDefined(), "Input array to log_softmax is uninitialized");
    int64_t rank = static_cast<int64_t>(input.shape().Rank());
    TENZR_CHECK(rank > 0, "log_softmax requires input array with rank >= 1");
    int64_t ax = (axis < 0) ? (axis + rank) : axis;
    TENZR_CHECK(ax >= 0 && ax < rank, "log_softmax axis out of range");

    auto max_val = input.Max(std::vector<int64_t>{ax}, /*keepdims=*/true);
    auto shifted = input - max_val;
    auto exp_shifted = shifted.Exp();
    auto sum_exp = exp_shifted.Sum(std::vector<int64_t>{ax}, /*keepdims=*/true);
    return shifted - sum_exp.Log();
  }

}  // namespace tenzr::nn