torch-rb 0.2.2 → 0.2.7

data/lib/torch/inspector.rb

@@ -1,89 +1,264 @@
+# mirrors _tensor_str.py
 module Torch
   module Inspector
-    # TODO make more performance, especially when summarizing
-    # how? only read data that will be displayed
-    def inspect
-      data =
-        if numel == 0
-          "[]"
-        elsif dim == 0
-          item
+    PRINT_OPTS = {
+      precision: 4,
+      threshold: 1000,
+      edgeitems: 3,
+      linewidth: 80,
+      sci_mode: nil
+    }
+
+    class Formatter
+      def initialize(tensor)
+        @floating_dtype = tensor.floating_point?
+        @complex_dtype = tensor.complex?
+        @int_mode = true
+        @sci_mode = false
+        @max_width = 1
+
+        tensor_view = Torch.no_grad { tensor.reshape(-1) }
+
+        if !@floating_dtype
+          tensor_view.each do |value|
+            value_str = value.item.to_s
+            @max_width = [@max_width, value_str.length].max
+          end
         else
-          summarize = numel > 1000
+          nonzero_finite_vals = Torch.masked_select(tensor_view, Torch.isfinite(tensor_view) & tensor_view.ne(0))
+
+          # no valid number, do nothing
+          return if nonzero_finite_vals.numel == 0
+
+          # Convert to double for easy calculation. HalfTensor overflows with 1e8, and there's no div() on CPU.
+          nonzero_finite_abs = nonzero_finite_vals.abs.double
+          nonzero_finite_min = nonzero_finite_abs.min.double
+          nonzero_finite_max = nonzero_finite_abs.max.double
+
+          nonzero_finite_vals.each do |value|
+            if value.item != value.item.ceil
+              @int_mode = false
+              break
+            end
+          end
 
-          if dtype == :bool
-            fmt = "%s"
+          if @int_mode
+            # in int_mode for floats, all numbers are integers, and we append a decimal to nonfinites
+            # to indicate that the tensor is of floating type. add 1 to the len to account for this.
+            if nonzero_finite_max / nonzero_finite_min > 1000.0 || nonzero_finite_max > 1.0e8
+              @sci_mode = true
+              nonzero_finite_vals.each do |value|
+                value_str = "%.#{PRINT_OPTS[:precision]}e" % value.item
+                @max_width = [@max_width, value_str.length].max
+              end
+            else
+              nonzero_finite_vals.each do |value|
+                value_str = "%.0f" % value.item
+                @max_width = [@max_width, value_str.length + 1].max
+              end
+            end
           else
-            values = to_a.flatten
-            abs = values.select { |v| v != 0 }.map(&:abs)
-            max = abs.max || 1
-            min = abs.min || 1
-
-            total = 0
-            if values.any? { |v| v < 0 }
-              total += 1
+            # Check if scientific representation should be used.
+            if nonzero_finite_max / nonzero_finite_min > 1000.0 || nonzero_finite_max > 1.0e8 || nonzero_finite_min < 1.0e-4
+              @sci_mode = true
+              nonzero_finite_vals.each do |value|
+                value_str = "%.#{PRINT_OPTS[:precision]}e" % value.item
+                @max_width = [@max_width, value_str.length].max
+              end
+            else
+              nonzero_finite_vals.each do |value|
+                value_str = "%.#{PRINT_OPTS[:precision]}f" % value.item
+                @max_width = [@max_width, value_str.length].max
+              end
             end
+          end
+        end
 
-            if floating_point?
-              sci = max / min.to_f > 1000 || max > 1e8 || min < 1e-4
+        @sci_mode = PRINT_OPTS[:sci_mode] unless PRINT_OPTS[:sci_mode].nil?
+      end
 
-              all_int = values.all? { |v| v.finite? && v == v.to_i }
-              decimal = all_int ? 1 : 4
+      def width
+        @max_width
+      end
 
-              total += sci ? 10 : decimal + 1 + max.to_i.to_s.size
+      def format(value)
+        value = value.item
 
-              if sci
-                fmt = "%#{total}.4e"
-              else
-                fmt = "%#{total}.#{decimal}f"
-              end
-            else
-              total += max.to_s.size
-              fmt = "%#{total}d"
+        if @floating_dtype
+          if @sci_mode
+            ret = "%#{@max_width}.#{PRINT_OPTS[:precision]}e" % value
+          elsif @int_mode
+            ret = String.new("%.0f" % value)
+            unless value.infinite? || value.nan?
+              ret += "."
             end
+          else
+            ret = "%.#{PRINT_OPTS[:precision]}f" % value
           end
+        elsif @complex_dtype
+          p = PRINT_OPTS[:precision]
+          raise NotImplementedYet
+        else
+          ret = value.to_s
+        end
+        # Ruby throws error when negative, Python doesn't
+        " " * [@max_width - ret.size, 0].max + ret
+      end
+    end
+
+    def inspect
+      Torch.no_grad do
+        str_intern(self)
+      end
+    rescue => e
+      # prevent stack error
+      puts e.backtrace.join("\n")
+      "Error inspecting tensor: #{e.inspect}"
+    end
+
+    private
+
+    # TODO update
+    def str_intern(slf)
+      prefix = "tensor("
+      indent = prefix.length
+      suffixes = []
+
+      has_default_dtype = [:float32, :int64, :bool].include?(slf.dtype)
+
+      if slf.numel == 0 && !slf.sparse?
+        # Explicitly print the shape if it is not (0,), to match NumPy behavior
+        if slf.dim != 1
+          suffixes << "size: #{shape.inspect}"
+        end
 
-          inspect_level(to_a, fmt, dim - 1, 0, summarize)
+        # In an empty tensor, there are no elements to infer if the dtype
+        # should be int64, so it must be shown explicitly.
+        if slf.dtype != :int64
+          suffixes << "dtype: #{slf.dtype.inspect}"
         end
+        tensor_str = "[]"
+      else
+        if !has_default_dtype
+          suffixes << "dtype: #{slf.dtype.inspect}"
+        end
+
+        if slf.layout != :strided
+          tensor_str = tensor_str(slf.to_dense, indent)
+        else
+          tensor_str = tensor_str(slf, indent)
+        end
+      end
 
-      attributes = []
-      if requires_grad
-        attributes << "requires_grad: true"
+      if slf.layout != :strided
+        suffixes << "layout: #{slf.layout.inspect}"
       end
-      if ![:float32, :int64, :bool].include?(dtype)
-        attributes << "dtype: #{dtype.inspect}"
+
+      # TODO show grad_fn
+      if slf.requires_grad?
+        suffixes << "requires_grad: true"
       end
 
-      "tensor(#{data}#{attributes.map { |a| ", #{a}" }.join("")})"
+      add_suffixes(prefix + tensor_str, suffixes, indent, slf.sparse?)
     end
 
-    private
+    def add_suffixes(tensor_str, suffixes, indent, force_newline)
+      tensor_strs = [tensor_str]
+      # rfind in Python returns -1 when not found
+      last_line_len = tensor_str.length - (tensor_str.rindex("\n") || -1) + 1
+      suffixes.each do |suffix|
+        suffix_len = suffix.length
+        if force_newline || last_line_len + suffix_len + 2 > PRINT_OPTS[:linewidth]
+          tensor_strs << ",\n" + " " * indent + suffix
+          last_line_len = indent + suffix_len
+          force_newline = false
+        else
+          tensor_strs.append(", " + suffix)
+          last_line_len += suffix_len + 2
+        end
+      end
+      tensor_strs.append(")")
+      tensor_strs.join("")
+    end
 
-    # TODO DRY code
-    def inspect_level(arr, fmt, total, level, summarize)
-      if level == total
-        cols =
-          if summarize && arr.size > 7
-            arr[0..2].map { |v| fmt % v } +
-            ["..."] +
-            arr[-3..-1].map { |v| fmt % v }
-          else
-            arr.map { |v| fmt % v }
-          end
+    def tensor_str(slf, indent)
+      return "[]" if slf.numel == 0
+
+      summarize = slf.numel > PRINT_OPTS[:threshold]
+
+      if slf.dtype == :float16 || slf.dtype == :bfloat16
+        slf = slf.float
+      end
+      formatter = Formatter.new(summarize ? summarized_data(slf) : slf)
+      tensor_str_with_formatter(slf, indent, formatter, summarize)
+    end
+
+    def summarized_data(slf)
+      edgeitems = PRINT_OPTS[:edgeitems]
 
-        "[#{cols.join(", ")}]"
+      dim = slf.dim
+      if dim == 0
+        slf
+      elsif dim == 1
+        if size(0) > 2 * edgeitems
+          Torch.cat([slf[0...edgeitems], slf[-edgeitems..-1]])
+        else
+          slf
+        end
+      elsif slf.size(0) > 2 * edgeitems
+        start = edgeitems.times.map { |i| slf[i] }
+        finish = (slf.length - edgeitems).upto(slf.length - 1).map { |i| slf[i] }
+        Torch.stack((start + finish).map { |x| summarized_data(x) })
       else
-        rows =
-          if summarize && arr.size > 7
-            arr[0..2].map { |row| inspect_level(row, fmt, total, level + 1, summarize) } +
-            ["..."] +
-            arr[-3..-1].map { |row| inspect_level(row, fmt, total, level + 1, summarize) }
-          else
-            arr.map { |row| inspect_level(row, fmt, total, level + 1, summarize) }
-          end
+        Torch.stack(slf.map { |x| summarized_data(x) })
+      end
+    end
+
+    def tensor_str_with_formatter(slf, indent, formatter, summarize)
+      edgeitems = PRINT_OPTS[:edgeitems]
+
+      dim = slf.dim
 
-        "[#{rows.join(",#{"\n" * (total - level)}#{" " * (level + 8)}")}]"
+      return scalar_str(slf, formatter) if dim == 0
+      return vector_str(slf, indent, formatter, summarize) if dim == 1
+
+      if summarize && slf.size(0) > 2 * edgeitems
+        slices = (
+          [edgeitems.times.map { |i| tensor_str_with_formatter(slf[i], indent + 1, formatter, summarize) }] +
+          ["..."] +
+          [((slf.length - edgeitems)...slf.length).map { |i| tensor_str_with_formatter(slf[i], indent + 1, formatter, summarize) }]
+        )
+      else
+        slices = slf.size(0).times.map { |i| tensor_str_with_formatter(slf[i], indent + 1, formatter, summarize) }
       end
+
+      tensor_str = slices.join("," + "\n" * (dim - 1) + " " * (indent + 1))
+      "[" + tensor_str + "]"
+    end
+
+    def scalar_str(slf, formatter)
+      formatter.format(slf)
+    end
+
+    def vector_str(slf, indent, formatter, summarize)
+      # length includes spaces and comma between elements
+      element_length = formatter.width + 2
+      elements_per_line = [1, ((PRINT_OPTS[:linewidth] - indent) / element_length.to_f).floor.to_i].max
+      char_per_line = element_length * elements_per_line
+
+      if summarize && slf.size(0) > 2 * PRINT_OPTS[:edgeitems]
+        data = (
+          [slf[0...PRINT_OPTS[:edgeitems]].map { |val| formatter.format(val) }] +
+          [" ..."] +
+          [slf[-PRINT_OPTS[:edgeitems]..-1].map { |val| formatter.format(val) }]
+        )
+      else
+        data = slf.map { |val| formatter.format(val) }
+      end
+
+      data_lines = (0...data.length).step(elements_per_line).map { |i| data[i...(i + elements_per_line)] }
+      lines = data_lines.map { |line| line.join(", ") }
+      "[" + lines.join("," + "\n" + " " * (indent + 1)) + "]"
     end
   end
 end

data/lib/torch/native/function.rb

@@ -66,6 +66,7 @@ module Torch
             end
 
             next if t == "Generator?"
+            next if t == "MemoryFormat"
             next if t == "MemoryFormat?"
             args << {name: k, type: t, default: d, pos: pos, has_default: has_default}
           end

data/lib/torch/native/generator.rb

@@ -18,7 +18,7 @@ module Torch
           functions = functions()
 
           # skip functions
-          skip_args = ["bool[3]", "Dimname", "MemoryFormat", "Layout", "Storage", "ConstQuantizerPtr"]
+          skip_args = ["bool[3]", "Dimname", "Layout", "Storage", "ConstQuantizerPtr"]
 
           # remove functions
           functions.reject! do |f|
@@ -31,7 +31,7 @@ module Torch
           todo_functions, functions =
             functions.partition do |f|
               f.args.any? do |a|
-                a[:type].include?("?") && !["Tensor?", "Generator?", "int?", "ScalarType?"].include?(a[:type]) ||
+                a[:type].include?("?") && !["Tensor?", "Generator?", "int?", "ScalarType?", "Tensor?[]"].include?(a[:type]) ||
                 skip_args.any? { |sa| a[:type].include?(sa) } ||
                 # native_functions.yaml is missing size argument for normal
                 # https://pytorch.org/cppdocs/api/function_namespacetorch_1a80253fe5a3ded4716ec929a348adb4b9.html
@@ -112,6 +112,9 @@ void add_%{type}_functions(Module m) {
                   "OptionalScalarType"
                 when "Tensor[]"
                   "TensorList"
+                when "Tensor?[]"
+                  # TODO make optional
+                  "TensorList"
                 when "int"
                   "int64_t"
                 when "float"

data/lib/torch/native/parser.rb

@@ -75,7 +75,7 @@ module Torch
                 v.is_a?(Tensor)
               when "Tensor?"
                 v.nil? || v.is_a?(Tensor)
-              when "Tensor[]"
+              when "Tensor[]", "Tensor?[]"
                 v.is_a?(Array) && v.all? { |v2| v2.is_a?(Tensor) }
               when "int"
                 if k == "reduction"

data/lib/torch/nn/batch_norm.rb

@@ -70,6 +70,11 @@ module Torch
           momentum: exponential_average_factor, eps: @eps
         )
       end
+
+      def extra_inspect
+        s = "%{num_features}, eps: %{eps}, momentum: %{momentum}, affine: %{affine}, track_running_stats: %{track_running_stats}"
+        format(s, **dict)
+      end
     end
   end
 end

data/lib/torch/nn/conv2d.rb

@@ -20,7 +20,14 @@ module Torch
 
       # TODO add more parameters
       def extra_inspect
-        format("%s, %s, kernel_size: %s, stride: %s", @in_channels, @out_channels, @kernel_size, @stride)
+        s = String.new("%{in_channels}, %{out_channels}, kernel_size: %{kernel_size}, stride: %{stride}")
+        s += ", padding: %{padding}" if @padding != [0] * @padding.size
+        s += ", dilation: %{dilation}" if @dilation != [1] * @dilation.size
+        s += ", output_padding: %{output_padding}" if @output_padding != [0] * @output_padding.size
+        s += ", groups: %{groups}" if @groups != 1
+        s += ", bias: false" unless @bias
+        s += ", padding_mode: %{padding_mode}" if @padding_mode != "zeros"
+        format(s, **dict)
       end
     end
   end

data/lib/torch/nn/convnd.rb

@@ -23,7 +23,7 @@ module Torch
         if bias
           @bias = Parameter.new(Tensor.new(out_channels))
         else
-          raise NotImplementedError
+          register_parameter("bias", nil)
         end
         reset_parameters
       end

data/lib/torch/nn/max_poolnd.rb

@@ -12,7 +12,8 @@ module Torch
       end
 
       def extra_inspect
-        format("kernel_size: %s", @kernel_size)
+        s = "kernel_size: %{kernel_size}, stride: %{stride}, padding: %{padding}, dilation: %{dilation}, ceil_mode: %{ceil_mode}"
+        format(s, **dict)
       end
     end
   end

data/lib/torch/nn/module.rb

@@ -145,7 +145,7 @@ module Torch
         params = {}
         if recurse
           named_children.each do |name, mod|
-            params.merge!(mod.named_parameters(prefix: "#{name}.", recurse: recurse))
+            params.merge!(mod.named_parameters(prefix: "#{prefix}#{name}.", recurse: recurse))
           end
         end
         instance_variables.each do |name|
@@ -186,8 +186,22 @@ module Torch
         named_modules.values
       end
 
-      def named_modules
-        {"" => self}.merge(named_children)
+      # TODO return enumerator?
+      def named_modules(memo: nil, prefix: "")
+        ret = {}
+        memo ||= Set.new
+        unless memo.include?(self)
+          memo << self
+          ret[prefix] = self
+          named_children.each do |name, mod|
+            next unless mod.is_a?(Module)
+            submodule_prefix = prefix + (!prefix.empty? ? "." : "") + name
+            mod.named_modules(memo: memo, prefix: submodule_prefix).each do |m|
+              ret[m[0]] = m[1]
+            end
+          end
+        end
+        ret
       end
 
       def train(mode = true)
@@ -224,13 +238,15 @@ module Torch
 
       def inspect
         name = self.class.name.split("::").last
-        if children.empty?
+        if named_children.empty?
           "#{name}(#{extra_inspect})"
         else
           str = String.new
           str << "#{name}(\n"
-          children.each do |name, mod|
-            str << "  (#{name}): #{mod.inspect}\n"
+          named_children.each do |name, mod|
+            mod_str = mod.inspect
+            mod_str = mod_str.lines.join("  ")
+            str << "  (#{name}): #{mod_str}\n"
           end
           str << ")"
         end
@@ -270,8 +286,11 @@ module Torch
         str % vars
       end
 
+      # used for format
+      # remove tensors for performance
+      # so we can skip call to inspect
       def dict
-        instance_variables.map { |k| [k[1..-1].to_sym, instance_variable_get(k)] }.to_h
+        instance_variables.reject { |k| instance_variable_get(k).is_a?(Tensor) }.map { |k| [k[1..-1].to_sym, instance_variable_get(k)] }.to_h
       end
     end
   end