src/modules/nn/dropout.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/dropout.h"

#include <vector>

#include "tenzr/common/error.h"
#include "tenzr/modules/random/random.h"

namespace tenzr::nn {

  Dropout::Dropout(double p, std::optional<uint64_t> seed) : p_(p), gen_(seed.has_value() ? *seed : random::DefaultGenerator().NextUint64()) {
    TENZR_CHECK(p >= 0.0 && p < 1.0, "Dropout probability p must be in [0, 1)");
  }

  void Dropout::SetP(double p) {
    TENZR_CHECK(p >= 0.0 && p < 1.0, "Dropout probability p must be in [0, 1)");
    p_ = p;
  }

  Tensor Dropout::Forward(const Tensor& input) {
    TENZR_CHECK(input.IsDefined(), "Input tensor is uninitialized");

    if (!IsTraining() || p_ == 0.0) {
      return input;
    }

    if (input.HasData()) {
      double keep_prob = 1.0 - p_;
      double scale = 1.0 / keep_prob;

      auto mask = random::Bernoulli(input.shape(), static_cast<float>(keep_prob), gen_);
      auto mask_vec = mask.ToVectorAs<double>();
      auto in_vec = input.ToVectorAs<double>();

      std::vector<double> out_vec(in_vec.size());
      for (size_t i = 0; i < in_vec.size(); ++i) {
        out_vec[i] = in_vec[i] * mask_vec[i] * scale;
      }

      return Tensor::FromVector<double>(out_vec, input.shape()).Astype(input.Dtype());
    }

    return Tensor::CustomOp("dropout", {input});
  }

}  // namespace tenzr::nn