init.h

#pragma once

#include <torch/csrc/WindowsTorchApiMacro.h>
#include <torch/types.h>

namespace torch {
namespace nn {
namespace init {

enum class Nonlinearity {
  Linear,
  Conv1D,
  Conv2D,
  Conv3D,
  ConvTranspose1D,
  ConvTranspose2D,
  ConvTranspose3D,
  Sigmoid,
  Tanh,
  ReLU,
  LeakyReLU
};

enum class FanMode { FanIn, FanOut };

TORCH_API double calculate_gain(Nonlinearity nonlinearity, double param = 0.01);

TORCH_API Tensor constant_(Tensor tensor, Scalar value);

TORCH_API Tensor dirac_(Tensor tensor);

TORCH_API Tensor eye_(Tensor matrix);

TORCH_API Tensor normal_(Tensor tensor, double mean = 0, double std = 1);

TORCH_API Tensor ones_(Tensor tensor);

TORCH_API Tensor orthogonal_(Tensor tensor, double gain = 1.0);

TORCH_API Tensor sparse_(Tensor tensor, double sparsity, double std = 0.01);

TORCH_API Tensor uniform_(Tensor tensor, double low = 0, double high = 1);

TORCH_API Tensor kaiming_normal_(
    Tensor tensor,
    double a = 0,
    FanMode mode = FanMode::FanIn,
    Nonlinearity nonlinearity = Nonlinearity::LeakyReLU);

TORCH_API Tensor kaiming_uniform_(
    Tensor tensor,
    double a = 0,
    FanMode mode = FanMode::FanIn,
    Nonlinearity nonlinearity = Nonlinearity::LeakyReLU);

TORCH_API Tensor xavier_normal_(Tensor tensor, double gain = 1.0);

TORCH_API Tensor xavier_uniform_(Tensor tensor, double gain = 1.0);

TORCH_API Tensor zeros_(Tensor tensor);

} // namespace init
} // namespace nn
} // namespace torch