tensor_stream 1.0.4 → 1.0.9

data/lib/tensor_stream/generated_stub/stub_file.erb

@@ -9,12 +9,12 @@ module TensorStream
 <%end%>    #
     #<% if op.supports_broadcasting? %> This operation supports broadcasting
     #<% end %>
-    # Params:
-<% op.parameters.each do |param| %>    # +<%= param[:name] %>+:: <%= param[:description]%><%if param[:validate]%> (of type <%= param[:validate] %>)<%end%>
+<% op.parameters.each do |param| %>    # @param <%= param[:name] %> <%= param[:description]%><%if param[:validate]%> (of type <%= param[:validate] %>)<%end%>
 <% end %>    #
     # Options:
-<% op.options.each do |k, v| %>    # +:<%= k %>+:: <%= v[:description]%><% if v[:default_value] != :nil %> default (<%= v[:default_value] %>)<%end%>
-<%end%>    def <%= op.operation.to_s %>(<%= (op.expand_params(true) + op.expand_options(true)).join(', ') %>)
+<% op.options.each do |k, v| %>    # @option <%= k %> <%= v[:description]%><% if v[:default_value] != :nil %> default (<%= v[:default_value] %>)<%end%>
+<%end%>    # @return Tensor
+    def <%= op.operation.to_s %>(<%= (op.expand_params(true) + op.expand_options(true)).join(', ') %>)
 <%= op.generate_body %>
     end
 <% op.aliases.each do |a|%>

data/lib/tensor_stream/graph.rb

@@ -30,6 +30,7 @@ module TensorStream
         :"#{GraphKeys::TRAINABLE_VARIABLES}" => [],
       }
       @constants = {}
+      TensorStream::Evaluator.clear_storages(self)
     end
 
     def as_default
@@ -129,7 +130,7 @@ module TensorStream
 
     def add_op(operation, *args)
       options = if args.last.is_a?(Hash)
-        args.pop
+        args.pop || {}
       else
         {}
       end
@@ -180,8 +181,7 @@ module TensorStream
 
     def add_variable!(node, options = {})
       node = add_variable(node, options)
-      op = Graph.get_default_graph.add_op!(:variable_v2, container: node, internal_name: node.name, shape: options[:shape], data_type: options[:data_type])
-      node.name = op.name
+      op = Graph.get_default_graph.add_op!(:variable_v2, var_name: node.name, shape: options[:shape], data_type: options[:data_type])
       op
     end
 

data/lib/tensor_stream/graph_deserializers/yaml_loader.rb

@@ -31,15 +31,13 @@ module TensorStream
         options = {}
 
         new_var = nil
-        if op_def.dig(:attrs, :container)
+        if op_def[:op].to_sym == :variable_v2
           new_var = Variable.new(op_def.dig(:attrs, :data_type))
-          var_shape = op_def.dig(:attrs, :container, :shape)
-          var_options = op_def.dig(:attrs, :container, :options)
-          var_options[:name] = op_def[:name]
 
-          new_var.prepare(var_shape.size, var_shape, TensorStream.get_variable_scope, var_options)
-          options[:container] = new_var
+          var_options = {}
+          var_options[:name] = op_def.dig(:attrs, :var_name)
 
+          new_var.prepare(nil, nil, TensorStream.get_variable_scope, var_options)
           @graph.add_variable(new_var, var_options)
         end
 

data/lib/tensor_stream/helpers/infer_shape.rb

@@ -10,13 +10,7 @@ module TensorStream
     def self.infer_shape(tensor)
       case tensor.operation
       when :assign
-        possible_shape = if tensor.inputs[0]&.shape&.shape
-          tensor.inputs[0].shape.shape
-        else
-          tensor.inputs[1].shape.shape
-        end
-
-        possible_shape
+        tensor.inputs[0]&.shape&.shape
       when :const
         shape_eval(tensor.options[:value])
       when :variable_v2

data/lib/tensor_stream/helpers/tensor_mixins.rb

@@ -5,7 +5,16 @@ module TensorStream
     end
 
     def [](index)
-      _op(:index, self, index)
+      if index.is_a?(Range)
+        last = if index.end.nil?
+                 [TensorStream.shape(self)[0]]
+               else
+                 [index.max + 1]
+               end
+        _op(:strided_slice, self, [index.min], last, [1])
+      else
+        _op(:index, self, index)
+      end
     end
 
     def *(other)

data/lib/tensor_stream/images.rb

@@ -7,6 +7,10 @@ module TensorStream
       _op(:decode_png, contents, channels: channels, data_type: dtype, name: name, new_shape: new_shape)
     end
 
+    def self.decode_jpeg(contents, channels: 0, dtype: :uint8, name: nil, new_shape: nil)
+      _op(:decode_jpg, contents, channels: channels, data_type: dtype, name: name, new_shape: new_shape)
+    end
+
     def self.encode_png(contents, compression: -1, name: nil, new_shape: nil, resample_method: nil)
       check_allowed_types(contents, %i[uint8 uint16])
       contents = convert_to_tensor(contents, dtype: :uint16)

data/lib/tensor_stream/math/math_ops.rb

@@ -0,0 +1,22 @@
+module TensorStream
+  # High level math functions
+  class Maths
+    extend TensorStream::OpHelper
+
+    module MathFunctions
+
+      ##
+      # Normalizes along dimension axis using an L2 norm.
+      def l2_normalize(x, axis: nil, epsilon: 1e-12, name: nil)
+        TensorStream.name_scope(name, "l2_normalize", values: [x]) do |name|
+          x = TensorStream.convert_to_tensor(x, name: "x")
+          square_sum = TensorStream.reduce_sum(TensorStream.square(x), axis, keepdims: true)
+          x_inv_norm = TensorStream.rsqrt(TensorStream.maximum(square_sum, epsilon))
+          TensorStream.multiply(x, x_inv_norm, name: name)
+        end
+      end
+    end
+
+    extend MathFunctions
+  end
+end

data/lib/tensor_stream/math_gradients.rb

@@ -136,7 +136,7 @@ module TensorStream
         when :sparse_softmax_cross_entropy_with_logits
           output = node
           [_broadcast_mul(grad, output[1]), nil]
-         when :zeros_like
+        when :zeros_like
           # non differentiable
           nil
         when :transpose
@@ -165,12 +165,26 @@ module TensorStream
           ts.stack(grad, axis: node.options[:axis])
         when :conv2d
           _Conv2DGrad(node, grad)
+        when :flow_dynamic_stitch
+          num_values = node.inputs.size / 2
+          indices_grad = [nil] * num_values
+
+          inputs = (0...num_values).map { |i| _int32(node, node.inputs[i]) }
+
+          values_grad = inputs.map { |inp| TensorStream.gather(grad, inp) }
+          indices_grad + values_grad
+        when :gather
+          [_op(:gather_grad, grad, node.inputs[1], TensorStream.shape(node.inputs[0])), nil]
         else
           TensorStream::OpMaker.gradient_op(self, node, grad)
         end
       end
     end
 
+    def self._int32(node, x)
+      (node.inputs[0].data_type == :int32 ? x : TensorStream.cast(x, :int32))
+    end
+
     def self._reshape_to_input(node, grad)
       ts.reshape(grad, ts.shape(node.inputs[0]))
     end

data/lib/tensor_stream/nn/embedding_lookup.rb

@@ -0,0 +1,114 @@
+require 'tensor_stream/utils/py_ports'
+##
+# ruby port of https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/ops/embedding_ops.py
+#
+module TensorStream
+  module EmbeddingLookup
+    include TensorStream::PyPorts
+
+    ##
+    # Looks up `ids` in a list of embedding tensors.
+    def embedding_lookup(params, ids, partition_strategy: "mod", name: nil, validate_indices: true, max_norm: nil)
+      _embedding_lookup_and_transform(params, ids, partition_strategy: partition_strategy, name: name, max_norm: max_norm, transform_fn: nil)
+    end
+
+    ##
+    # Helper function for embedding_lookup and _compute_sampled_logits.
+    def _embedding_lookup_and_transform(params, ids, partition_strategy: "mod", name: nil, max_norm: nil, transform_fn: nil)
+      raise TensorStream::ValueError, "Need at least one param" if params.nil?
+
+      params = [params] unless params.is_a?(Array)
+
+      TensorStream.name_scope(name, "embedding_lookup", values: params + [ids]) do |name|
+        np = params.size
+        ids = TensorStream.convert_to_tensor(ids, name: "ids")
+        if (np == 1) && (transform_fn.nil? || (ids.shape.size == 1))
+          result = nil
+          TensorStream.colocate_with(params[0]) do
+            result = _clip(TensorStream.gather(params[0], ids, name: name), ids, max_norm)
+            result = transform_fn.call(result) if transform_fn
+          end
+
+          return TensorStream.identity(result)
+        else
+          flat_ids = TensorStream.reshape(ids, [-1])
+          original_indices = TensorStream.range(TensorStream.size(flat_ids))
+
+          p_assignments = nil
+          new_ids = nil
+
+          if partition_strategy == "mod"
+            p_assignments = flat_ids % np
+            new_ids = floor_div(flat_ids, np)
+          elsif partition_strategy == "div"
+            raise "not yet supported!"
+          else
+            raise TensorStream::ValueError, "Unrecognized partition strategy: " + partition_strategy
+          end
+
+          p_assignments = TensorStream.cast(p_assignments, :int32)
+          gather_ids = TensorStream.dynamic_partition(new_ids, p_assignments, np)
+          pindices = TensorStream.dynamic_partition(original_indices, p_assignments, np)
+          partitioned_result = []
+          (0...np).each do |p|
+            pids = gather_ids[p]
+            result = nil
+            TensorStream.colocate_with(params[p]) do
+              result = TensorStream.gather(params[p], pids)
+              if transform_fn
+                # If transform_fn is provided, the clip_by_norm precedes
+                # the transform and hence must be co-located. See below
+                # for the counterpart if transform_fn is not proveded.
+                result = transform_fn.call(_clip(result, pids, max_norm))
+              end
+            end
+            partitioned_result << result
+          end
+          ret = TensorStream.dynamic_stitch(pindices, partitioned_result, name: name)
+
+          if transform_fn.nil?
+            element_shape_s = params[0].shape[1..-1]
+            params[1..-1].each { |p| element_shape_s = element_shape_s.merge_with(p.shape[1..-1]) }
+          else
+            element_shape_s = ret.shape[1..-1]
+          end
+
+           # Compute the dynamic element shape.
+          element_shape_d = if element_shape_s.fully_defined?
+                               element_shape_s
+                            elsif transform_fn.nil?
+                              # It's important that we compute params[0].shape on the right device
+                              # to avoid data motion.
+                              TensorStream.colocate_with(params[0]) do
+                                params_shape = TensorStream.shape(params[0])
+                                params_shape[1..-1]
+                              end
+                            else
+                              TensorStream.shape(ret)[1..-1]
+                            end
+          ret = TensorStream.reshape(ret, TensorStream.concat([TensorStream.shape(ids), element_shape_d], 0))
+          ret = _clip(ret, ids, max_norm) unless transform_fn
+          ret
+        end
+      end
+    end
+
+    def _clip(params, ids, max_norm)
+      return params if max_norm.nil?
+
+      ids_rank, ids_static = _rank(ids)
+      params_rank, params_static = _rank(params)
+
+      TensorStream.clip_by_norm(params, max_norm, axes: ids_static && params_static ? (ids_rank...params_rank).to_a : TensorStream.range(ids_rank, params_rank))
+    end
+
+    def _rank(x)
+      rank = TensorStream.convert_to_tensor(x).shape.ndims
+      if rank
+        [rank, false]
+      else
+        [TensorStream.rank(x), false]
+      end
+    end
+  end
+end

data/lib/tensor_stream/nn/nn_ops.rb

@@ -1,7 +1,10 @@
+require 'tensor_stream/nn/embedding_lookup'
 module TensorStream
   # High level machine learning functions
   class NN
     extend TensorStream::OpHelper
+    extend TensorStream::EmbeddingLookup
+    extend TensorStream::Maths::MathFunctions
 
     class << self
       def softmax(logits, axis: nil, name: nil)
@@ -137,6 +140,19 @@ module TensorStream
       def conv2d(input, filter, strides, padding, name: nil)
         _op(:conv2d, input, filter, strides: strides, padding: padding, name: name)
       end
+
+      ##
+      # Adds bias to value.
+      #
+      # This is a narrow version of tf add where the bias is restructed to 1-D only
+      def bias_add(value, bias, data_format: nil, name: nil)
+        value = TensorStream.convert_to_tensor(value, name: "input")
+        bias = TensorStream.convert_to_tensor(bias, dtype: value.dtype, name: "bias")
+
+        raise TensorStreamError, "value must be at least rank 2" if value.shape.known? && value.shape.ndims < 2
+
+        _op(:bias_add, value, bias, data_format: data_format, name: name)
+      end
     end
   end
 

data/lib/tensor_stream/op_maker.rb

@@ -2,7 +2,8 @@ class TensorStream::OpMaker
   attr_reader :operation, :description, :parameters,
               :options, :gradient, :check_types,
               :supports_broadcast, :data_type_coercion,
-              :aliases, :custom, :infer_type_proc, :exclude
+              :aliases, :custom, :custom_post, :infer_type_proc, :exclude,
+              :data_type_block
 
   def initialize(op)
     @operation = op
@@ -15,6 +16,7 @@ class TensorStream::OpMaker
     @description = []
     @aliases = []
     @custom = []
+    @custom_post = []
     @infer_type_proc = lambda { |tensor|
       next nil if tensor.inputs[0].nil?
       next tensor.inputs[0].shape.shape if tensor.inputs.size == 1
@@ -31,6 +33,10 @@ class TensorStream::OpMaker
     @custom << custom_code
   end
 
+  def add_custom_post(custom_code)
+    @custom_post << custom_code
+  end
+
   def self.scan
     op_files = Dir[File.join(File.dirname(__FILE__), "ops", "*.rb")]
     op_files.each { |file|
@@ -58,6 +64,22 @@ class TensorStream::OpMaker
     context_caller.instance_exec(tensor, &@ops[tensor.operation].infer_type_proc)
   end
 
+  def self.infer_data_type(context_caller, tensor, passed_data_type)
+    return passed_data_type if passed_data_type
+
+    if @ops[tensor.operation] && @ops[tensor.operation].data_type_block
+      context_caller.instance_exec(tensor, &@ops[tensor.operation].data_type_block)
+    else
+      if tensor.inputs[0]
+        tensor.inputs[0].data_type
+      elsif tensor.inputs[1]
+        tensor.inputs[1].data_type
+      else
+        :unknown
+      end
+    end
+  end
+
   def self.each_op(&block)
     @ops.values.sort_by { |op| op.operation }.reject(&:exclude).each do |op|
       block.call(op)
@@ -94,7 +116,14 @@ class TensorStream::OpMaker
     custom.each do |c|
       body << c
     end
-    body << "_op(:#{operation}, #{(expand_params(false) + options_call).join(', ')})"
+    if custom_post.empty?
+      body << "_op(:#{operation}, #{(expand_params(false) + options_call).join(', ')})"
+    else
+      body << "result = _op(:#{operation}, #{(expand_params(false) + options_call).join(', ')})"
+    end
+    custom_post.each do |c|
+      body << c
+    end
     body.map { |line| "      #{line}"}.join("\n")
   end
 
@@ -122,6 +151,10 @@ class TensorStream::OpMaker
     @infer_type_proc = block
   end
 
+  def define_data_type(&block)
+    @data_type_block = block
+  end
+
   def expand_params(print_defaults)
     @parameters.map { |param|
       print_defaults && param[:default_value] ? "#{param[:name]} = #{default_with_nil(param[:default_value])}" : "#{param[:name]}"
@@ -163,7 +196,7 @@ class TensorStream::OpMaker
   end
 
   def options_call
-    @options.map { |k, v|
+    @options.reject { |k, v| v.dig(:options, :exclude) }.map { |k, v|
       if v.dig(:options, :alias)
         "#{v.dig(:options, :alias)}: #{k}"
       else

data/lib/tensor_stream/operation.rb

@@ -7,7 +7,7 @@ module TensorStream
     attr_accessor :name, :operation, :inputs, :rank, :device, :consumers, :breakpoint
     attr_reader :outputs, :options, :is_const, :data_type, :shape
 
-    def initialize(graph, inputs:, options:)
+    def initialize(graph, inputs: [], options: {})
       @consumers = Set.new
       @outputs = []
       @op = self
@@ -42,14 +42,6 @@ module TensorStream
       @options[:container] ? @options[:container].buffer : nil
     end
 
-    def container
-      @options[:container].read_value
-    end
-
-    def container=(value)
-      @options[:container].value = value
-    end
-
     def set_input(index, value)
       @inputs[index] = value
       @shape = TensorShape.new(TensorStream::InferShape.infer_shape(self))
@@ -58,6 +50,10 @@ module TensorStream
       @data_type = set_data_type(@options[:data_type])
     end
 
+    def set_option(key, value)
+      @options.merge!(key.to_sym => value)
+    end
+
     def infer_const
       return false if breakpoint
 
@@ -68,7 +64,7 @@ module TensorStream
         true
       when :placeholder
         false
-      when :variable_v2
+      when :variable_v2, :assign, :assign_add, :assign_sub
         false
       else
         non_const = @inputs.compact.find { |input| !input.is_const }
@@ -96,7 +92,7 @@ module TensorStream
         options[:data_type]
       when :fill
         @inputs[1].data_type
-      when :greater, :less, :equal, :not_equal, :greater_equal, :less_equal, :logical_and
+      when :logical_and
         :boolean
       when :shape, :rank, :shape_n
         options[:out_type] || :int32
@@ -119,15 +115,7 @@ module TensorStream
           @inputs[0].data_type
         end
       else
-        return passed_data_type if passed_data_type
-
-        if @inputs[0]
-          @inputs[0].data_type
-        elsif @inputs[1]
-          @inputs[1].data_type
-        else
-          :unknown
-        end
+        OpMaker.infer_data_type(self, self, passed_data_type)
       end
     end