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

#include <algorithm>
#include <cmath>
#include <vector>

#include "tenzr/common/error.h"
#include "tenzr/frontend/op_registry.h"
#include "tenzr/modules/random/random.h"
#include "tenzr/plax/plax.h"
#include "tenzr/txpr/custom_vjp.h"

namespace tenzr::nn {

  Conv2D::Conv2D(int64_t in_channels, int64_t out_channels, std::pair<int64_t, int64_t> kernel_size, std::pair<int64_t, int64_t> stride,
                 std::pair<int64_t, int64_t> padding, bool use_bias, DType dtype, std::optional<uint64_t> seed)
      : in_channels_(in_channels), out_channels_(out_channels), kernel_size_(kernel_size), stride_(stride), padding_(padding), use_bias_(use_bias) {
    TENZR_CHECK(in_channels > 0, "in_channels must be positive");
    TENZR_CHECK(out_channels > 0, "out_channels must be positive");
    TENZR_CHECK(kernel_size.first > 0 && kernel_size.second > 0, "kernel dimensions must be positive");
    TENZR_CHECK(stride.first > 0 && stride.second > 0, "stride dimensions must be positive");
    TENZR_CHECK(padding.first >= 0 && padding.second >= 0, "padding dimensions must be non-negative");

    int64_t fan_in = in_channels * kernel_size.first * kernel_size.second;
    if (seed.has_value()) {
      random::Generator gen(*seed);
      weight_ = RegisterParameter("weight", random::KaimingUniform({out_channels_, in_channels_, kernel_size_.first, kernel_size_.second}, fan_in, dtype, gen));
    } else {
      weight_ = RegisterParameter(
          "weight", random::KaimingUniform({out_channels_, in_channels_, kernel_size_.first, kernel_size_.second}, fan_in, dtype, random::DefaultGenerator()));
    }

    if (use_bias_) {
      bias_ = RegisterParameter("bias", Tensor::Zeros({out_channels_}, dtype));
    }
  }

  Tensor Conv2D::Forward(const Tensor& input) {
    TENZR_CHECK(input.IsDefined(), "Input tensor is uninitialized");
    TENZR_CHECK(input.shape().Rank() == 4, "Conv2D input must have rank 4 [B, C, H, W]");
    TENZR_CHECK(input.shape().Dim(1) == in_channels_,
                "Input channels (" + std::to_string(input.shape().Dim(1)) + ") does not match layer channels (" + std::to_string(in_channels_) + ")");

    int64_t b = input.shape().Dim(0);
    int64_t h = input.shape().Dim(2);
    int64_t w = input.shape().Dim(3);

    int64_t out_h = (h + 2 * padding_.first - kernel_size_.first) / stride_.first + 1;
    int64_t out_w = (w + 2 * padding_.second - kernel_size_.second) / stride_.second + 1;
    TENZR_CHECK(out_h > 0 && out_w > 0, "Output spatial dimensions must be positive");

    Shape out_shape({b, out_channels_, out_h, out_w});

    // Traceable / differentiable path: lower to plax::conv2d (im2col + matmul),
    // which composes differentiable primitives so jit/grad work.
    if (txpr::IsGradTracingActive() || input.IsSymbolic() || !input.HasData() || !weight_.HasData()) {
      Tensor out = plax::Conv2d(input, weight_, stride_, padding_);
      if (use_bias_) {
        out = out + bias_.Reshape(Shape({1, out_channels_, 1, 1}));
      }
      return out;
    }

    if (input.HasData() && weight_.HasData()) {
      auto in_vec = input.ToVectorAs<double>();
      auto w_vec = weight_.ToVectorAs<double>();
      std::vector<double> b_vec;
      if (use_bias_ && bias_.HasData()) {
        b_vec = bias_.ToVectorAs<double>();
      }

      std::vector<double> out_vec(static_cast<size_t>(out_shape.NumElements()), 0.0);

      int64_t in_c = in_channels_;
      int64_t out_c = out_channels_;
      int64_t kh = kernel_size_.first;
      int64_t kw = kernel_size_.second;
      int64_t sh = stride_.first;
      int64_t sw = stride_.second;
      int64_t ph = padding_.first;
      int64_t pw = padding_.second;

      for (int64_t bi = 0; bi < b; ++bi) {
        for (int64_t oc = 0; oc < out_c; ++oc) {
          double bias_val = b_vec.empty() ? 0.0 : b_vec[static_cast<size_t>(oc)];
          for (int64_t oh = 0; oh < out_h; ++oh) {
            int64_t ih_start = oh * sh - ph;
            for (int64_t ow = 0; ow < out_w; ++ow) {
              int64_t iw_start = ow * sw - pw;
              double sum = bias_val;

              for (int64_t ic = 0; ic < in_c; ++ic) {
                for (int64_t ki = 0; ki < kh; ++ki) {
                  int64_t ih = ih_start + ki;
                  if (ih < 0 || ih >= h) {
                    continue;
                  }
                  for (int64_t kj = 0; kj < kw; ++kj) {
                    int64_t iw = iw_start + kj;
                    if (iw < 0 || iw >= w) {
                      continue;
                    }

                    size_t in_idx = static_cast<size_t>(((bi * in_c + ic) * h + ih) * w + iw);
                    size_t w_idx = static_cast<size_t>(((oc * in_c + ic) * kh + ki) * kw + kj);
                    sum += in_vec[in_idx] * w_vec[w_idx];
                  }
                }
              }

              size_t out_idx = static_cast<size_t>(((bi * out_c + oc) * out_h + oh) * out_w + ow);
              out_vec[out_idx] = sum;
            }
          }
        }
      }

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

    return Tensor::CustomOp("conv2d", {input, weight_});
  }
}  // namespace tenzr::nn