Added support for :bool dtype to numo method

ankane · Aug 19, 2024 · 9dcf3ce · 9dcf3ce
1 parent b509c33
commit 9dcf3ce
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/lib/torch/tensor.rb b/lib/torch/tensor.rb
@@ -132,9 +132,13 @@ def new
 
     # TODO read directly from memory
     def numo
-      cls = Torch._dtype_to_numo[dtype]
-      raise Error, "Cannot convert #{dtype} to Numo" unless cls
-      cls.from_string(_data_str).reshape(*shape)
+      if dtype == :bool
+        Numo::UInt8.from_string(_data_str).ne(0).reshape(*shape)
+      else
+        cls = Torch._dtype_to_numo[dtype]
+        raise Error, "Cannot convert #{dtype} to Numo" unless cls
+        cls.from_string(_data_str).reshape(*shape)
+      end
     end
 
     def requires_grad=(requires_grad)

diff --git a/test/numo_test.rb b/test/numo_test.rb
@@ -9,6 +9,10 @@ def test_numo
     x = Torch.tensor([[1, 2, 3], [4, 5, 6]])
     assert x.numo.is_a?(Numo::Int64)
     assert_equal x.to_a, x.numo.to_a
+
+    x = Torch.tensor([[true, false], [false, true]])
+    assert x.numo.is_a?(Numo::Bit)
+    assert_equal [[1, 0], [0, 1]], x.numo.to_a
   end
 
   def test_from_numo