tensor_stream 0.9.8 → 0.9.9

data/lib/tensor_stream/math_gradients.rb

@@ -40,7 +40,10 @@ module TensorStream
         return nil if grads.empty?
         grads.size > 1 ? ts.add_n(grads) : grads[0]
       else
-        return nil if computed_op.nil?
+
+        if computed_op.nil?
+          return nil
+        end
         _propagate(computed_op, tensor.inputs[0], stop_tensor, nodes_to_compute, stop_gradients)
       end
     end
@@ -50,6 +53,7 @@ module TensorStream
       node.graph.name_scope("#{node.name}_grad") do
         x = node.inputs[0] if node.inputs[0]
         y = node.inputs[1] if node.inputs[1]
+        z = node.inputs[2] if node.inputs[2]
 
         case node.operation
         when :add_n
@@ -135,7 +139,6 @@ module TensorStream
             reduced = ts.cast(reduction_indices, :int32)
             idx = ts.range(0, rank)
             other, = ts.setdiff1d(idx, reduced)
-
             [ts.concat([reduced, other], 0),
              ts.reduce_prod(ts.gather(input_shape, reduced)),
              ts.reduce_prod(ts.gather(input_shape, other))]
@@ -239,13 +242,17 @@ module TensorStream
           y = ts.constant(2.0, dtype: x.dtype)
           ts.multiply(grad, ts.multiply(x, y))
         when :where
-          x_mask = i_op(:where, i_op(:ones_like, x), i_op(:zeros_like, y), pred: node.options[:pred])
-          y_mask = i_op(:where, i_op(:zeros_like, x), i_op(:ones_like, y), pred: node.options[:pred])
-          [x_mask * grad, y_mask * grad]
-        when :cond
-          x_cond = i_op(:cond, i_op(:ones_like, x), i_op(:zeros_like, y), pred: node.options[:pred])
-          y_cond = i_op(:cond, i_op(:zeros_like, x), i_op(:ones_like, x), pred: node.options[:pred])
-          [x_cond * grad, y_cond * grad]
+          x_mask = i_op(:where, x, i_op(:ones_like, y), i_op(:zeros_like, z))
+          y_mask = i_op(:where, x, i_op(:zeros_like, y), i_op(:ones_like, z))
+          [nil, x_mask * grad, y_mask * grad]
+        when :case
+          n_preds = node.inputs.size - 2
+
+          case_grads = Array.new(n_preds) do |index|
+            i_op(:case_grad, index, node.inputs[0], node.inputs[2 + index], grad)
+          end
+
+          [nil, i_op(:case_grad, -1, node.inputs[0], node.inputs[1], grad)] + case_grads
         when :mean
           sum_grad = _sum_grad(x, y, grad)[0]
           input_shape = ts.shape(x)
@@ -283,12 +290,12 @@ module TensorStream
           # hack!! not sure how to fix this yet
           return grad if %i[softmax_cross_entropy_with_logits_v2 sparse_softmax_cross_entropy_with_logits].include?(node.inputs[0].operation)
 
-          if node.inputs[0].shape.known? && node.inputs[1].value
+          if node.inputs[0].shape.known? && node.inputs[1].const_value
             multiplier = node.inputs[0].shape.shape[0]
             filler = ts.zeros_like(grad)
 
             res = Array.new(multiplier) do |index|
-              index == node.inputs[1].value ? grad : filler
+              index == node.inputs[1].const_value ? grad : filler
             end
             [res]
           end
@@ -346,7 +353,7 @@ module TensorStream
       tile_scaling = _safe_shape_div(input_shape, output_shape_kept_dims)
       new_grad = _op(:reshape, grad, output_shape_kept_dims)
 
-      grad = _op(:cond, _op(:fill, input_shape, grad), _op(:tile, new_grad, tile_scaling), pred: _op(:rank, grad).zero?)
+      grad = _op(:case, [_op(:rank, grad).zero?], _op(:tile, new_grad, tile_scaling), _op(:fill, input_shape, grad))
 
       [grad, nil]
     end

data/lib/tensor_stream/monkey_patches/float.rb

@@ -1,3 +1,10 @@
 class Float
   include TensorStream::MonkeyPatch
+
+  def self.placeholder(name: nil, width: 32, shape: nil)
+    raise "invalid width passed #{width}" unless [16, 32, 64].include?(width)
+
+    data_type = :"float#{width}"
+    TensorStream.placeholder(data_type, name: name, shape: shape)
+  end
 end

data/lib/tensor_stream/monkey_patches/integer.rb

@@ -1,3 +1,10 @@
 class Integer
   include TensorStream::MonkeyPatch
+
+  def self.placeholder(name: nil, width: 32, shape: nil)
+    raise "invalid width passed #{width}" unless [16, 32, 64].include?(width)
+
+    data_type = :"int#{width}"
+    TensorStream.placeholder(data_type, name: name, shape: shape)
+  end
 end

data/lib/tensor_stream/monkey_patches/patch.rb

@@ -19,13 +19,13 @@ module TensorStream
       TensorStream.shape_eval(self)
     end
 
-    def t(name = nil)
-      TensorStream.convert_to_tensor(self, name: name)
+    def t(name = nil, dtype: nil)
+      TensorStream.convert_to_tensor(self, name: name, dtype: dtype)
     end
 
     def +(other)
       if other.is_a?(TensorStream::Tensor)
-        TensorStream.convert_to_tensor(self) + other
+        TensorStream.convert_to_tensor(self, dtype: other.data_type) + other
       else
         _tensor_stream_add_orig(other)
       end
@@ -33,7 +33,7 @@ module TensorStream
 
     def -(other)
       if other.is_a?(TensorStream::Tensor)
-        TensorStream.convert_to_tensor(self) - other
+        TensorStream.convert_to_tensor(self, dtype: other.data_type) - other
       else
         _tensor_stream_sub_orig(other)
       end
@@ -41,7 +41,7 @@ module TensorStream
 
     def *(other)
       if other.is_a?(TensorStream::Tensor)
-        TensorStream.convert_to_tensor(self) * other
+        TensorStream.convert_to_tensor(self, dtype: other.data_type) * other
       else
         _tensor_stream_mul_orig(other)
       end
@@ -49,7 +49,7 @@ module TensorStream
 
     def /(other)
       if other.is_a?(TensorStream::Tensor)
-        TensorStream.convert_to_tensor(self) * other
+        TensorStream.convert_to_tensor(self, dtype: other.data_type) * other
       else
         _tensor_stream_div_orig(other)
       end
@@ -57,7 +57,7 @@ module TensorStream
 
     def %(other)
       if other.is_a?(TensorStream::Tensor)
-        TensorStream.convert_to_tensor(self) % other
+        TensorStream.convert_to_tensor(self, dtype: other.data_type) % other
       else
         _tensor_stream_mod_orig(other)
       end
@@ -65,7 +65,7 @@ module TensorStream
 
     def **(other)
       if other.is_a?(TensorStream::Tensor)
-        TensorStream.convert_to_tensor(self)**other
+        TensorStream.convert_to_tensor(self, dtype: other.data_type)**other
       else
         _tensor_stream_pow_orig(other)
       end

data/lib/tensor_stream/nn/nn_ops.rb

@@ -5,7 +5,7 @@ module TensorStream
 
     class << self
       def softmax(logits, axis: nil, name: nil)
-        _op(:softmax, logits, nil, axis: axis, name: name)
+        _op(:softmax, logits, axis: axis, name: name)
       end
 
       def relu(features, name: nil)

data/lib/tensor_stream/operation.rb

@@ -2,33 +2,18 @@ require 'tensor_stream/helpers/infer_shape'
 module TensorStream
   # TensorStream class that defines an operation
   class Operation < Tensor
-    attr_accessor :name, :operation, :inputs, :rank, :options
-    attr_reader :outputs
+    include OpHelper
 
-    def initialize(operation, *args)
-      options = if args.last.is_a?(Hash)
-                  args.pop
-                else
-                  {}
-                end
-
-      inputs = args
-
-      setup_initial_state(options)
-
-      @operation = operation
-      @rank = options[:rank] || 0
-      @name = [@graph.get_name_scope, options[:name] || set_name].compact.reject(&:empty?).join('/')
-      @internal = options[:internal]
-      @given_name = @name
+    attr_accessor :name, :operation, :inputs, :rank, :device, :consumers, :breakpoint
+    attr_reader :outputs, :options, :is_const, :data_type, :shape
 
+    def initialize(graph, inputs:, options:)
+      @consumers = Set.new
+      @outputs = []
+      @op = self
+      @graph = graph
+      @inputs = inputs
       @options = options
-
-      @inputs = inputs.map { |i| options[:preserve_params_type] ? i : TensorStream.convert_to_tensor(i) }
-      @data_type = set_data_type(options[:data_type])
-      @is_const = infer_const
-      @shape = TensorShape.new(TensorStream::InferShape.infer_shape(self))
-      @graph.add_node(self)
     end
 
     def to_s
@@ -37,25 +22,74 @@ module TensorStream
 
     def to_h
       {
-        op: operation,
-        name: name,
-        operands: hashify_tensor(inputs)
+        op: operation.to_s,
+        name: name.to_s,
+        data_type: @data_type,
+        inputs: @inputs.map(&:name),
+        attrs: serialize_options
       }
     end
 
+    def const_value
+      @options ? @options[:value] : nil
+    end
+
+    def container_buffer
+      @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))
+      @rank = @shape.rank
+      @is_const = infer_const
+      @data_type = set_data_type(@options[:data_type])
+    end
+
     def infer_const
       return false if breakpoint
+
       case operation
       when :random_standard_normal, :random_uniform, :truncated_normal, :glorot_uniform, :print, :check_numerics
         false
+      when :const
+        true
+      when :placeholder
+        false
+      when :variable_v2
+        false
       else
         non_const = @inputs.compact.find { |input| !input.is_const }
         non_const ? false : true
       end
     end
 
+    def set_name
+      @operation.to_s
+    end
+
     def set_data_type(passed_data_type)
       case operation
+      when :where
+        @inputs[1].data_type
+      when :case
+        if @inputs[2]
+          @inputs[2].data_type
+        else
+          @inputs[1].data_type
+        end
+      when :case_grad
+        @inputs[2].data_type
+      when :placeholder, :variable_v2, :const
+        options[:data_type]
       when :fill
         @inputs[1].data_type
       when :greater, :less, :equal, :not_equal, :greater_equal, :less_equal, :logical_and
@@ -70,9 +104,8 @@ module TensorStream
         @inputs[1].data_type
       when :index
         if @inputs[0].is_a?(ControlFlow)
-
           if @inputs[1].is_const
-            @inputs[0].inputs[@inputs[1].value].data_type
+            @inputs[0].inputs[@inputs[1].const_value].data_type
           else
             :unknown
           end
@@ -163,8 +196,6 @@ module TensorStream
               "reshape(#{sub_input},#{sub_input2})"
             when :rank
               "#{sub_input}.rank"
-            when :cond
-              "(#{auto_math(options[:pred], name_only, max_depth - 1, cur_depth)} ? #{sub_input} : #{sub_input2})"
             when :less
               "#{sub_input} < #{sub_input2}"
             when :less_equal
@@ -222,12 +253,41 @@ module TensorStream
 
     private
 
+    def serialize_options
+      excludes = %i[internal_name source]
+
+      @options.reject { |k, v| excludes.include?(k) || v.nil? }.map do |k, v|
+        v = case v.class.to_s
+            when 'TensorStream::TensorShape'
+              v.shape
+            when 'Array'
+              v
+            when 'String', 'Integer', 'Float', 'Symbol', 'FalseClass', "TrueClass"
+              v
+            when 'TensorStream::Variable'
+              { name: v.name, options: v.options, shape: v.shape.shape.dup }
+            else
+              raise "unknown type #{v.class}"
+            end
+        [k.to_sym, v]
+      end.to_h
+    end
+
+    def add_consumer(consumer)
+      @consumers << consumer.name if consumer.name != name
+    end
+
+    def setup_output(consumer)
+      @outputs << consumer.name unless @outputs.include?(consumer.name)
+    end
+
     def propagate_consumer(consumer)
-      super
+      add_consumer(consumer)
       @inputs.compact.each do |input|
         if input.is_a?(Array)
-          input.flatten.compact.select { |t| t.is_a?(Tensor) }.each do |t|
+          input.flatten.compact.map(&:op).select { |t| t.is_a?(Tensor) }.each do |t|
             next if t.consumers.include?(consumer.name)
+
             t.send(:propagate_consumer, consumer)
           end
         elsif input.name != name && !input.consumers.include?(consumer.name)
@@ -239,7 +299,7 @@ module TensorStream
     def propagate_outputs
       @inputs.compact.each do |input|
         if input.is_a?(Array)
-          input.flatten.compact.each do |t|
+          input.flatten.compact.map(&:op).each do |t|
             t.send(:setup_output, self) if t.is_a?(Tensor)
           end
         elsif input.is_a?(Tensor) && (input.name != name)
@@ -248,8 +308,10 @@ module TensorStream
       end
     end
 
-    def set_name
-      "#{@operation}#{graph.get_operation_counter}:#{@rank}"
+    def setup_initial_state(options)
+      @outputs = []
+      @graph = options[:__graph] || TensorStream.get_default_graph
+      @source = format_source(caller_locations)
     end
   end
 end

data/lib/tensor_stream/ops.rb

@@ -58,7 +58,8 @@ module TensorStream
     # +wrt_xs+ : A Tensor or list of tensors to be used for differentiation.
     # +stop_gradients+ :  Optional. A Tensor or list of tensors not to differentiate through
     def gradients(tensor_ys, wrt_xs, name: 'gradients', stop_gradients: nil)
-      gs = wrt_xs.collect do |x|
+      tensor_ys = tensor_ys.op
+      gs = wrt_xs.map(&:op).collect do |x|
         stops = stop_gradients ? stop_gradients.map(&:name).join('_') : ''
         gradient_program_name = "grad_#{tensor_ys.name}_#{x.name}_#{stops}".to_sym
         tensor_graph = tensor_ys.graph
@@ -82,21 +83,21 @@ module TensorStream
     # Outputs random values from a uniform distribution.
     def random_uniform(shape, dtype: :float32, minval: 0, maxval: 1, seed: nil, name: nil)
       options = { dtype: dtype, minval: minval, maxval: maxval, seed: seed, name: name }
-      _op(:random_uniform, shape, nil, options)
+      _op(:random_uniform, shape, options)
     end
 
     ##
     # Outputs random values from a normal distribution.
     def random_normal(shape, dtype: :float32, mean: 0.0, stddev: 1.0, seed: nil, name: nil)
       options = { dtype: dtype, mean: mean, stddev: stddev, seed: seed, name: name }
-      _op(:random_standard_normal, shape, nil, options)
+      _op(:random_standard_normal, shape, options)
     end
 
     ##
     # Outputs random values from a truncated normal distribution.
     def truncated_normal(shape, dtype: :float32, mean: 0.0, stddev: 1.0, seed: nil, name: nil)
       options = { dtype: dtype, mean: mean, stddev: stddev, seed: seed, name: name }
-      _op(:truncated_normal, shape, nil, options)
+      _op(:truncated_normal, shape, options)
     end
 
     ##
@@ -172,14 +173,14 @@ module TensorStream
     # initializer that generates tensors initialized to 0.
     #
     def zeros_initializer(dtype: :float32)
-      TensorStream::Initializer.new(-> { _op(:zeros, nil, nil, data_type: dtype) })
+      TensorStream::Initializer.new(-> { _op(:zeros, data_type: dtype) })
     end
 
     ##
     # initializer that generates tensors initialized to 1.
     #
     def ones_initializer(dtype: :float32)
-      TensorStream::Initializer.new(-> { _op(:ones, nil, nil, data_type: dtype) })
+      TensorStream::Initializer.new(-> { _op(:ones, data_type: dtype) })
     end
 
     ##
@@ -189,13 +190,13 @@ module TensorStream
     # where limit is sqrt(6 / (fan_in + fan_out)) where fan_in is the number
     # of input units in the weight tensor and fan_out is the number of output units in the weight tensor.
     def glorot_uniform_initializer(seed: nil, dtype: nil)
-      TensorStream::Initializer.new(-> { _op(:glorot_uniform, nil, nil, seed: seed, data_type: dtype) })
+      TensorStream::Initializer.new(-> { _op(:glorot_uniform, seed: seed, data_type: dtype) })
     end
 
     ##
     # Initializer that generates tensors with a uniform distribution.
     def random_uniform_initializer(minval: 0, maxval: 1, seed: nil, dtype: nil)
-      TensorStream::Initializer.new(-> { _op(:random_uniform, nil, nil, minval: 0, maxval: 1, seed: seed, data_type: dtype) })
+      TensorStream::Initializer.new(-> { _op(:random_uniform, minval: 0, maxval: 1, seed: seed, data_type: dtype) })
     end
 
     ##
@@ -276,7 +277,7 @@ module TensorStream
     ##
     # Computes the mean of elements across dimensions of a tensor.
     def reduce_mean(input_tensor, axis = nil, keepdims: false, name: nil)
-      _op(:mean, input_tensor, axis, keepdims: keepdims, name: name)
+      reduce(:mean, input_tensor, axis, keepdims: keepdims, name: name)
     end
 
     ##
@@ -288,7 +289,7 @@ module TensorStream
     # If axis has no entries, all dimensions are reduced, and a tensor with a single element
     # is returned.
     def reduce_sum(input_tensor, axis = nil, keepdims: false, name: nil)
-      _op(:sum, input_tensor, axis, keepdims: keepdims, name: name)
+      reduce(:sum, input_tensor, axis, keepdims: keepdims, name: name)
     end
 
     ##
@@ -300,7 +301,22 @@ module TensorStream
     #
     # If axis has no entries, all dimensions are reduced, and a tensor with a single element is returned.
     def reduce_prod(input, axis = nil, keepdims: false, name: nil)
-      _op(:prod, input, axis, keepdims: keepdims, name: name)
+      reduce(:prod, input, axis, keepdims: keepdims, name: name)
+    end
+
+    def reduce(op, input, axis = nil, keepdims: false, name: nil)
+      input = TensorStream.convert_to_tensor(input)
+      axis = if !axis.nil?
+               axis
+             elsif input.shape.scalar?
+               op
+             elsif input.shape.known?
+               (0...input.shape.ndims).to_a
+             else
+               range(0, rank(input))
+             end
+
+      _op(op, input, axis, keepdims: keepdims, name: name)
     end
 
     ##
@@ -395,13 +411,13 @@ module TensorStream
     ##
     # Return true_fn() if the predicate pred is true else false_fn().
     def cond(pred, true_fn, false_fn, name: nil)
-      _op(:cond, true_fn, false_fn, pred: pred, name: name)
+      _op(:case, [pred], false_fn, true_fn, name: name)
     end
 
     ##
     # Return the elements, either from x or y, depending on the condition.
     def where(condition, true_t = nil, false_t = nil, name: nil)
-      _op(:where, true_t, false_t, pred: condition, name: name)
+      _op(:where, condition, true_t, false_t, name: name)
     end
 
     ##
@@ -486,6 +502,7 @@ module TensorStream
     def max(input_a, input_b, name: nil)
       check_allowed_types(input_a, NUMERIC_TYPES)
       check_allowed_types(input_b, NUMERIC_TYPES)
+
       input_a, input_b = check_data_types(input_a, input_b)
       _op(:max, input_a, input_b, name: name)
     end
@@ -652,6 +669,13 @@ module TensorStream
       _op(:tanh, input, name: name)
     end
 
+        ##
+    # Computes sec of input element-wise.
+    def sec(input, name: nil)
+      check_allowed_types(input, FLOATING_POINT_TYPES)
+      _op(:sec, input, name: name)
+    end
+
     ##
     # Computes sqrt of input element-wise.
     def sqrt(input, name: nil)
@@ -819,6 +843,34 @@ module TensorStream
       [result[0], result[1]]
     end
 
+    ##
+    # Create a case operation.
+    #
+    # The pred_fn_pairs parameter is a dict or list of pairs of size N.
+    # Each pair contains a boolean scalar tensor and a proc that creates the tensors to be returned if the boolean evaluates to true.
+    # default is a proc generating a list of tensors. All the proc in pred_fn_pairs as well as default (if provided) should return the
+    # same number and types of tensors.
+    #
+    def case(args = {})
+      args = args.dup
+      default = args.delete(:default)
+      exclusive = args.delete(:exclusive)
+      strict = args.delete(:strict)
+      name = args.delete(:name)
+
+      predicates = []
+      functions = []
+
+      args.each do |k, v|
+        raise "Invalid argment or option #{k}" unless k.is_a?(Tensor)
+
+        predicates << k
+        functions << (v.is_a?(Proc) ? v.call : v)
+      end
+
+      _op(:case, predicates, default, *functions, exclusive: exclusive, strict: strict, name: name)
+    end
+
     def cumprod(x, axis: 0, exclusive: false, reverse: false, name: nil)
       _op(:cumprod, x, axis: axis, exclusive: exclusive, reverse: reverse, name: name)
     end