[#85]: Fix sigmoid backward

python/bindings.cpp

@@ -130,6 +130,8 @@ void bind_tensor_creation_functions(pybind11::module &m) {
   bind_tensor_creation_overloads(m, "ones", OnesFunctor{});
   bind_tensor_creation_overloads(m, "empty", EmptyFunctor{});
 
+  m.def("empty_like", &smollnet::empty_like, "tensor"_a,
+        "requires_grad"_a = false);
   m.def("full_like", &smollnet::full_like, "tensor"_a, "value"_a,
         "requires_grad"_a = false);
   m.def("zeros_like", &smollnet::zeros_like, "tensor"_a,

src/autograd.cpp

@@ -287,10 +287,11 @@ TanhFunction::backward(const std::vector<Tensor> &grad_outputs) {
   return grad_inputs;
 }
 
-// SigmoidFunction implementation
-SigmoidFunction::SigmoidFunction(const Tensor &input) {
+SigmoidFunction::SigmoidFunction(const Tensor &input,
+                                 const Tensor &sigmoid_output) {
   inputs = {input};
   needs_input_grad = {input.initialized() && input.requires_grad()};
+  sigmoid_output_data_ = sigmoid_output.data();
 }
 
 std::vector<Tensor>
@@ -304,7 +305,7 @@ SigmoidFunction::backward(const std::vector<Tensor> &grad_outputs) {
 
     auto grad_input = create_grad_tensor(inputs[0]);
     launch_sigmoid_grad(grad_input.data(), grad_outputs.front().data(),
-                        inputs[0].data(), grad_input.dtype(),
+                        sigmoid_output_data_, grad_input.dtype(),
                         grad_input.numel());
     grad_inputs[0] = grad_input;
   }

src/autograd.hpp

@@ -89,10 +89,14 @@ struct TanhFunction : Function {
 };
 
 struct SigmoidFunction : Function {
-  explicit SigmoidFunction(const Tensor &input);
+  SigmoidFunction(const Tensor &input, const Tensor &sigmoid_output);
   std::vector<Tensor>
   backward(const std::vector<Tensor> &grad_outputs) override;
   void print() const override { printf("SigmoidFunction\n"); }
+
+private:
+  // We're not storing Tensor here to avoid refernce cycle
+  const void *sigmoid_output_data_ = nullptr;
 };
 
 struct SumFunction : Function {

src/tensor.cpp

@@ -702,8 +702,7 @@ Tensor tanh(const Tensor &t) {
 }
 
 Tensor sigmoid(const Tensor &t) {
-  Tensor new_tensor = empty(t.dims().data(), t.ndims(), t.dtype(), t.device(),
-                            t.requires_grad());
+  Tensor new_tensor = empty_like(t, t.requires_grad());
 
   if (t.device() == Device::CUDA) {
     launch_sigmoid(new_tensor.data(), t.data(), t.dtype(), t.numel());
@@ -714,7 +713,8 @@ Tensor sigmoid(const Tensor &t) {
                    1.0f / (1.0f + std::exp(-x)));
     }
   }
-  SetupAutograd<SigmoidFunction>(new_tensor, t);
+  // We reuse the sigmoid redult for efficiency
+  SetupAutograd<SigmoidFunction>(t, new_tensor, new_tensor);
   return new_tensor;
 }
 
@@ -987,6 +987,11 @@ Tensor rand(const int64_t *dims, size_t rank, DataType t, Device d,
   return Tensor{tensor};
 }
 
+Tensor empty_like(const Tensor &t, bool requires_grad) {
+  return empty(t.dims().data(), t.ndims(), t.dtype(), t.device(),
+               requires_grad);
+}
+
 Tensor zeros_like(const Tensor &t, bool requires_grad) {
   return zeros(t.dims().data(), t.ndims(), t.dtype(), t.device(),
                requires_grad);

src/tensor.hpp

@@ -156,6 +156,8 @@ Tensor ones(const int64_t *dims, size_t rank, DataType t, Device d,
             bool requires_grad = false);
 Tensor rand(const int64_t *dims, size_t rank, DataType t, Device d,
             bool requires_grad = false);
+
+Tensor empty_like(const Tensor &t, bool requires_grad = false);
 Tensor full_like(const Tensor &t, float value, bool requires_grad = false);
 Tensor zeros_like(const Tensor &t, bool requires_grad = false);
 Tensor ones_like(const Tensor &t, bool requires_grad = false);

utils/tests.py

@@ -33,21 +33,27 @@ def test_tensor_creation():
     print(f"2D rand shape: {t2d_rand.dims()}")
 
     print("\nCreating like tensors...")
+    t2d_empty_like = smollnet.empty_like(t2d_rand)
     t2d_full_like = smollnet.full_like(t2d_rand, 3.0)
     t2d_zeros_like = smollnet.zeros_like(t2d_rand)
     t2d_ones_like = smollnet.ones_like(t2d_rand, requires_grad=True)
     t2d_rand_like = smollnet.rand_like(t2d_rand)
 
+    assert t2d_empty_like.dims() == t2d_rand.dims()
     assert t2d_full_like.dims() == t2d_rand.dims()
     assert t2d_zeros_like.dims() == t2d_rand.dims()
     assert t2d_ones_like.dims() == t2d_rand.dims()
     assert t2d_rand_like.dims() == t2d_rand.dims()
+    assert t2d_empty_like.dtype() == t2d_rand.dtype()
     assert t2d_full_like.dtype() == t2d_rand.dtype()
     assert t2d_zeros_like.dtype() == t2d_rand.dtype()
+    assert t2d_empty_like.device() == t2d_rand.device()
     assert t2d_ones_like.device() == t2d_rand.device()
+    assert not t2d_empty_like.requires_grad()
     assert not t2d_zeros_like.requires_grad()
     assert t2d_ones_like.requires_grad()
 
+    print(f"2D empty_like shape: {t2d_empty_like.dims()}")
     print(f"2D full_like shape: {t2d_full_like.dims()}")
     print(f"2D zeros_like shape: {t2d_zeros_like.dims()}")
     print(f"2D ones_like shape: {t2d_ones_like.dims()}")