// 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.
#ifndef TENZR_MODULES_NN_CONV2D_H_
#define TENZR_MODULES_NN_CONV2D_H_
#include <cstdint>
#include <optional>
#include <string>
#include <utility>
#include "tenzr/frontend/dtype.h"
#include "tenzr/frontend/tensor.h"
#include "tenzr/modules/nn/module.h"
namespace tenzr::nn {
class Conv2D : public Module {
public:
// Dynamic DType constructor (explicit dtype, optional seed)
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 = std::nullopt);
// Compile-time templated constructor using C++20 concepts (defaults to float via concepts::FloatingPoint<T>)
template <typename T = float>
requires concepts::FloatingPoint<T>
Conv2D(int64_t in_channels, int64_t out_channels, std::pair<int64_t, int64_t> kernel_size, std::pair<int64_t, int64_t> stride = {1, 1},
std::pair<int64_t, int64_t> padding = {0, 0}, bool use_bias = true, std::optional<uint64_t> seed = std::nullopt)
: Conv2D(in_channels, out_channels, kernel_size, stride, padding, use_bias, DtypeOf<T>(), seed) {}
// Tag dispatch constructor with explicit type: Conv2D(std::type_identity<double>{}, 3, 16, {3, 3})
template <typename T>
requires concepts::FloatingPoint<T>
explicit Conv2D(std::type_identity<T>, int64_t in_channels, int64_t out_channels, std::pair<int64_t, int64_t> kernel_size,
std::pair<int64_t, int64_t> stride = {1, 1}, std::pair<int64_t, int64_t> padding = {0, 0}, bool use_bias = true,
std::optional<uint64_t> seed = std::nullopt)
: Conv2D(in_channels, out_channels, kernel_size, stride, padding, use_bias, DtypeOf<T>(), seed) {}
// Compile-time factory method: Conv2D::create<double>(3, 16, {3, 3})
template <typename T = float>
requires concepts::FloatingPoint<T>
static Conv2D Create(int64_t in_channels, int64_t out_channels, std::pair<int64_t, int64_t> kernel_size, std::pair<int64_t, int64_t> stride = {1, 1},
std::pair<int64_t, int64_t> padding = {0, 0}, bool use_bias = true, std::optional<uint64_t> seed = std::nullopt) {
return Conv2D(in_channels, out_channels, kernel_size, stride, padding, use_bias, DtypeOf<T>(), seed);
}
Tensor Forward(const Tensor& input) override;
Tensor& Weight() { return weight_; }
const Tensor& Weight() const { return weight_; }
Tensor& Bias() { return bias_; }
const Tensor& Bias() const { return bias_; }
int64_t InChannels() const { return in_channels_; }
int64_t OutChannels() const { return out_channels_; }
std::pair<int64_t, int64_t> KernelSize() const { return kernel_size_; }
std::pair<int64_t, int64_t> Stride() const { return stride_; }
std::pair<int64_t, int64_t> Padding() const { return padding_; }
bool HasBias() const { return use_bias_; }
std::string Name() const override { return "Conv2D"; }
private:
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_;
Tensor weight_;
Tensor bias_;
};
} // namespace tenzr::nn
#endif // TENZR_MODULES_NN_CONV2D_H_