[2026春季][T2-2-1] surprisely

infini_train/include/common/cuda/kernel_helper.cuh

@@ -109,8 +109,10 @@ template <typename T> __device__ __forceinline__ T Cos(const T &x) {
 }
 
 template <typename T> __device__ __forceinline__ T Tanh(const T &x) {
-    if constexpr (std::is_same_v<T, nv_bfloat16> || std::is_same_v<T, half>) {
-        return htanh(x);
+    if constexpr (std::is_same_v<T, nv_bfloat16>) {
+        return __float2bfloat16(tanhf(__bfloat162float(x)));
+    } else if constexpr (std::is_same_v<T, half>) {
+        return __float2half(tanhf(__half2float(x)));
     } else if constexpr (std::is_same_v<T, float>) {
         return tanhf(x);
     } else {

infini_train/include/generator.h

@@ -0,0 +1,55 @@
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "infini_train/include/device.h"
+
+namespace infini_train {
+
+class GeneratorImpl {
+public:
+    virtual ~GeneratorImpl() = default;
+
+    virtual void ManualSeed(uint64_t seed) = 0;
+    virtual uint64_t Seed() = 0;
+    virtual uint64_t InitialSeed() const = 0;
+    virtual std::vector<uint8_t> GetState() const = 0;
+    virtual void SetState(const std::vector<uint8_t> &state) = 0;
+    virtual Device GetDevice() const = 0;
+    virtual std::pair<uint64_t, uint64_t> ReserveRandomOffset(uint64_t increment) = 0;
+
+    virtual void FillUniform(std::vector<float> &buffer, float from, float to) = 0;
+    virtual void FillNormal(std::vector<float> &buffer, float mean, float std) = 0;
+};
+
+class Generator {
+public:
+    explicit Generator(std::shared_ptr<GeneratorImpl> impl);
+
+    void ManualSeed(uint64_t seed);
+    uint64_t Seed();
+    uint64_t InitialSeed() const;
+    std::vector<uint8_t> GetState() const;
+    void SetState(const std::vector<uint8_t> &state);
+    Device GetDevice() const;
+    std::pair<uint64_t, uint64_t> ReserveRandomOffset(uint64_t increment);
+
+    void FillUniform(std::vector<float> &buffer, float from, float to);
+    void FillNormal(std::vector<float> &buffer, float mean, float std);
+
+private:
+    std::shared_ptr<GeneratorImpl> impl_;
+};
+
+std::shared_ptr<Generator> MakeCPUGenerator(uint64_t seed = 42);
+std::shared_ptr<Generator> MakeCUDAGenerator(int8_t device_index, uint64_t seed = 42);
+std::shared_ptr<Generator> GetDefaultCPUGenerator();
+std::shared_ptr<Generator> GetDefaultCUDAGenerator(int8_t device_index);
+std::shared_ptr<Generator> GetDefaultGenerator(const Device &device);
+void ManualSeed(uint64_t seed);
+void ManualSeedAll(uint64_t seed);
+
+} // namespace infini_train

infini_train/include/nn/functional.h

@@ -4,9 +4,12 @@
 #include <memory>
 #include <vector>
 
+#include "infini_train/include/device.h"
+
 namespace infini_train {
+class Generator;
 class Tensor;
-}
+} // namespace infini_train
 
 namespace infini_train::nn::function {
 
@@ -47,6 +50,12 @@ std::shared_ptr<Tensor> Triu(const std::shared_ptr<Tensor> &input, int64_t diago
 //   A tensor of the given shape filled with the scalar value 1.
 std::shared_ptr<Tensor> Ones(const std::vector<int64_t> size);
 
+std::shared_ptr<Tensor> Rand(const std::vector<int64_t> &size, Device device = Device(),
+                             std::shared_ptr<Generator> generator = nullptr);
+
+std::shared_ptr<Tensor> Randn(const std::vector<int64_t> &size, Device device = Device(),
+                              std::shared_ptr<Generator> generator = nullptr);
+
 // Returns a new tensor with the reciprocal of the elements of input.
 //
 // Args:

infini_train/include/nn/init.h

@@ -1,12 +1,11 @@
 #pragma once
 
 #include <memory>
-#include <optional>
-#include <random>
 #include <utility>
 
 #include "infini_train/include/datatype.h"
 #include "infini_train/include/device.h"
+#include "infini_train/include/generator.h"
 
 namespace infini_train {
 class Tensor;
@@ -15,7 +14,7 @@ class Device;
 
 namespace infini_train::nn::init {
 std::shared_ptr<Tensor> Normal(const std::shared_ptr<Tensor> &tensor, float mean = 0.0, float std = 1.0,
-                               std::optional<std::mt19937> generator = std::nullopt);
+                               std::shared_ptr<Generator> generator = nullptr);
 
 std::pair<int64_t, int64_t> CalculateFanInAndFanOut(const std::shared_ptr<Tensor> &tensor);
 
@@ -42,10 +41,10 @@ enum class NonLinearityType : int8_t {
 std::shared_ptr<Tensor> KaimingUniform(const std::shared_ptr<Tensor> &tensor, float a = 0.0f,
                                        KaimingMode mode = KaimingMode::kFanIn,
                                        NonLinearityType non_linearity = NonLinearityType::kLeakyReLU,
-                                       std::optional<std::mt19937> generator = std::nullopt);
+                                       std::shared_ptr<Generator> generator = nullptr);
 
 std::shared_ptr<Tensor> Uniform(const std::shared_ptr<Tensor> &tensor, float a = 0.0f, float b = 1.0f,
-                                std::optional<std::mt19937> generator = std::nullopt);
+                                std::shared_ptr<Generator> generator = nullptr);
 
 std::shared_ptr<Tensor> Ones(const std::shared_ptr<Tensor> &tensor);
 

infini_train/include/tensor.h

@@ -4,7 +4,6 @@
 #include <iostream>
 #include <memory>
 #include <optional>
-#include <random>
 #include <vector>
 
 #include "Eigen/Dense"
@@ -15,6 +14,7 @@
 #include "infini_train/include/scalar.h"
 
 namespace infini_train {
+class Generator;
 namespace autograd {
 class Function;
 class AccumulateGrad;
@@ -150,8 +150,7 @@ class Tensor : public std::enable_shared_from_this<Tensor> {
     std::shared_ptr<Tensor> Unsqueeze(int64_t dim);
 
     // distribution
-    std::shared_ptr<Tensor> Uniform(float from = 0.0f, float to = 1.0f,
-                                    std::optional<std::mt19937> generator = std::nullopt);
+    std::shared_ptr<Tensor> Uniform(float from = 0.0f, float to = 1.0f, std::shared_ptr<Generator> generator = nullptr);
 
     std::shared_ptr<Tensor> Matmul(const std::shared_ptr<Tensor> &other);
     std::shared_ptr<Tensor> Outer(const std::shared_ptr<Tensor> &other);