torch-rb 0.17.0 → 0.17.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: cd28e938ac0d61829d538f4a36e9f3377f06257ab08fd3ec8d8e3912194a101a
-  data.tar.gz: 7bbb0c98e783c4a02e4ca46a1835c591fff3630c84d1a44abf1ef8df75c45de7
+  metadata.gz: f04aeddd9c66d668f78ed38d831ec66907743fd5be2cd076e9e547d00f1f2aa3
+  data.tar.gz: e76ca5199c4801c0eaa760ed3e41ebb6e521774bba084809656b1ed9e4b96483
 SHA512:
-  metadata.gz: 38415fe461d01d11d620d2072012ea6a86ad000e5d25920af3d2d4cd05c3b088479cb60a2bd2b6dc013da51e3f744913b3d734fc949ac2708b825c3c041f8c52
-  data.tar.gz: 2c54b30ccfa91f9b7576f7927377c03a9d82b7ff8aaf561e92ad89b2b00a2e7e3e08c268d606e6609cbaa13d2dc2e5e11c8329bf1b3475af7025a21dc8a1921f
+  metadata.gz: 683246aa91d19daa624a9d3a89fa5e4898f19fbb45a72ff2bcac0ee3025859f0f8ad26ac153e58150aedf10547cf670e2fe131cffda246546ded1f1344637040
+  data.tar.gz: bd0b955b997f09a274d8b59a21eafa8bf45972b6a20883b9d457a29417aa2c6d413b324abadfb6772c5bdc3aa0e633d3e25d6f23dd986fa603c55c69ae4a5a85

data/CHANGELOG.md

@@ -1,3 +1,8 @@
+## 0.17.1 (2024-08-19)
+
+- Added `persistent` option to `register_buffer` method
+- Added `prefix` and `recurse` options to `named_buffers` method
+
 ## 0.17.0 (2024-07-26)
 
 - Updated LibTorch to 2.4.0

data/lib/torch/nn/module.rb

@@ -10,16 +10,23 @@ module Torch
         @parameters = {}
         @buffers = {}
         @modules = {}
+        @non_persistent_buffers_set = Set.new
       end
 
       def forward
         raise NotImplementedError
       end
 
-      def register_buffer(name, tensor)
+      def register_buffer(name, tensor, persistent: true)
         # TODO add checks
         @buffers[name] = tensor
         instance_variable_set("@#{name}", tensor)
+
+        if persistent
+          @non_persistent_buffers_set.delete(name)
+        else
+          @non_persistent_buffers_set << name
+        end
       end
 
       def register_parameter(name, param)
@@ -190,8 +197,18 @@ module Torch
         named_buffers.values
       end
 
-      def named_buffers
-        @buffers || {}
+      # TODO set recurse: true in 0.18.0
+      def named_buffers(prefix: "", recurse: false)
+        buffers = {}
+        if recurse
+          named_children.each do |name, mod|
+            buffers.merge!(mod.named_buffers(prefix: "#{prefix}#{name}.", recurse: recurse))
+          end
+        end
+        (@buffers || {}).each do |k, v|
+          buffers[[prefix, k].join] = v
+        end
+        buffers
       end
 
       def children
@@ -390,7 +407,10 @@ module Torch
           destination[prefix + k] = v
         end
         named_buffers.each do |k, v|
-          destination[prefix + k] = v
+          # TODO exclude v.nil?
+          if !@non_persistent_buffers_set.include?(k)
+            destination[prefix + k] = v
+          end
         end
       end
 

data/lib/torch/tensor.rb

@@ -132,9 +132,13 @@ module Torch
 
     # 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)

data/lib/torch/version.rb

@@ -1,3 +1,3 @@
 module Torch
-  VERSION = "0.17.0"
+  VERSION = "0.17.1"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: torch-rb
 version: !ruby/object:Gem::Version
-  version: 0.17.0
+  version: 0.17.1
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2024-07-26 00:00:00.000000000 Z
+date: 2024-08-19 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rice
@@ -16,14 +16,14 @@ dependencies:
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: 4.1.0
+        version: '4.1'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: 4.1.0
+        version: '4.1'
 description:
 email: andrew@ankane.org
 executables: []