batchnorm.h

#pragma once

#include <torch/nn/cloneable.h>
#include <torch/nn/pimpl.h>
#include <torch/types.h>

#include <cstdint>

namespace torch {
namespace nn {

struct TORCH_API BatchNormOptions {
  /* implicit */ BatchNormOptions(int64_t features);
  TORCH_ARG(int64_t, features);
  TORCH_ARG(bool, affine) = true;
  TORCH_ARG(bool, stateful) = true;
  TORCH_ARG(double, eps) = 1e-5;
  TORCH_ARG(double, momentum) = 0.1;
};

class TORCH_API BatchNormImpl : public torch::nn::Cloneable<BatchNormImpl> {
 public:
  explicit BatchNormImpl(int64_t features)
      : BatchNormImpl(BatchNormOptions(features)) {}
  explicit BatchNormImpl(BatchNormOptions options);

  void reset() override;

  void pretty_print(std::ostream& stream) const override;

  Tensor forward(const Tensor& input);

  Tensor pure_forward(
      const Tensor& input,
      const Tensor& mean,
      const Tensor& variance);

  BatchNormOptions options;

  Tensor weight;

  Tensor bias;

  Tensor running_mean;

  Tensor running_var;
};

TORCH_MODULE(BatchNorm);

} // namespace nn
} // namespace torch