tensor_stream 0.8.5 → 0.8.6

data/lib/tensor_stream/evaluator/ruby_evaluator.rb

@@ -7,6 +7,7 @@ require 'tensor_stream/evaluator/ruby/nn_ops'
 require 'tensor_stream/evaluator/ruby/array_ops'
 require 'tensor_stream/evaluator/ruby/random_ops'
 require 'tensor_stream/evaluator/ruby/images_ops'
+require 'tensor_stream/evaluator/ruby/check_ops'
 
 module TensorStream
   module Evaluator
@@ -39,6 +40,7 @@ module TensorStream
       include TensorStream::ArrayOps
       include TensorStream::RandomOps
       include TensorStream::ImagesOps
+      include TensorStream::CheckOps
 
       def run(tensor, execution_context)
         return tensor.map { |t| run(t, execution_context) } if tensor.is_a?(Array) && !tensor.empty? && tensor[0].is_a?(Tensor)
@@ -116,10 +118,10 @@ module TensorStream
       end
 
       register_op(:cast) do |context, tensor, inputs|
-        call_op(:cast, inputs[0], context, ->(t, _b) { Tensor.cast_dtype(t, tensor.data_type) })
+        call_op(tensor, inputs[0], context, ->(t, _b) { Tensor.cast_dtype(t, tensor.data_type) })
       end
 
-      register_op(:sign) do |context, _tensor, inputs|
+      register_op(:sign) do |context, tensor, inputs|
         func = lambda { |x, _b|
           if x.zero? || (x.is_a?(Float) && x.nan?)
             0
@@ -132,7 +134,7 @@ module TensorStream
           end
         }
 
-        call_op(:sign, inputs[0], context, func)
+        call_op(tensor, inputs[0], context, func)
       end
 
       register_op(:logical_and) do |context, tensor, inputs|
@@ -217,10 +219,6 @@ module TensorStream
         call_vector_op(tensor, :greater_equal, a, b, context, ->(t, u) { t <= u })
       end
 
-      register_op :shape do |_context, tensor, inputs|
-        shape_eval(inputs[0], tensor.options[:out_type])
-      end
-
       register_op :broadcast_transform do |_context, _tensor, inputs|
         broadcast(inputs[0], inputs[1])
       end
@@ -266,14 +264,13 @@ module TensorStream
           raise TensorStream::InvalidArgumentError, "#{message} Invalid argument" if t.nan? || t.infinite?
           t
         }
-        call_op(:check_numerics, inputs[0], context, f)
+        call_op(tensor, inputs[0], context, f)
       end
 
       def eval_operation(tensor, child_context)
         return @context[tensor.name] if @context.key?(tensor.name)
-        # puts "ruby: #{tensor.name}"
         invoke(tensor, child_context).tap do |result|
-
+          # puts "ruby: #{tensor.name}"
           if tensor.breakpoint
             a = resolve_placeholder(tensor.inputs[0], child_context) if tensor.inputs && tensor.inputs[0]
             b = resolve_placeholder(tensor.inputs[1], child_context) if tensor.inputs && tensor.inputs[1]
@@ -412,8 +409,6 @@ module TensorStream
       def call_op(op, a, child_context, func)
         a = complete_eval(a, child_context)
         process_function_op(a, func)
-      rescue FullEvalNotPossible
-        TensorStream.send(op.to_sym, a)
       end
 
       def call_vector_op(tensor, op, a, b, child_context, func)
@@ -484,7 +479,15 @@ module TensorStream
 
         var = if placeholder.is_a?(Placeholder)
                 @context[placeholder.name.to_sym].tap do |c|
-                  raise "missing placeholder #{placeholder.name}" if c.nil?
+                  raise TensorStream::ValueError, "missing placeholder #{placeholder.name}" if c.nil?
+                  if placeholder.shape.shape
+                    value_shape = shape_eval(c)
+                    placeholder_shape = placeholder.shape.shape
+                    placeholder_shape.zip(value_shape).each do |p_shape, v_shape|
+                      next if p_shape.nil?
+                      raise TensorStream::ValueError, "placeholder expects #{placeholder_shape}, got #{value_shape}" if p_shape != v_shape
+                    end
+                  end
                 end
               else
                 placeholder

data/lib/tensor_stream/graph.rb

@@ -90,8 +90,10 @@ module TensorStream
       @node_keys << node.name
       @nodes[node.name] = node
       @constants[node.name] = node if node.is_const
+      # puts "adding node"
       node.send(:propagate_outputs)
       node.send(:propagate_consumer, node)
+      # puts "#{node.name}"
       node.value = node.eval if @eager_execution
     end
 

data/lib/tensor_stream/helpers/op_helper.rb

@@ -1,3 +1,4 @@
+
 module TensorStream
   # module that contains helper functions useful for ops
   module OpHelper

data/lib/tensor_stream/math_gradients.rb

@@ -10,7 +10,7 @@ module TensorStream
 
     def self.derivative(tensor, wrt_dx, options = {})
       return i_op(:ones_like, tensor) if tensor.equal?(wrt_dx)
-      return i_op(:zeros_like, tensor) unless wrt_dx.consumers.include?(tensor.name)
+      return i_op(:zeros_like, wrt_dx) unless wrt_dx.consumers.include?(tensor.name)
 
       nodes_to_compute = wrt_dx.consumers.select do |t|
         node = tensor.graph.nodes[t]
@@ -30,12 +30,15 @@ module TensorStream
       computed_op = _compute_derivative(tensor, grad)
 
       if computed_op.is_a?(Array)
-        computed_op.each_with_index.collect do |op_grad, index|
+        grads = computed_op.each_with_index.collect do |op_grad, index|
           next if op_grad.nil?
           next unless nodes_to_compute.include?(tensor.inputs[index].name)
 
           _propagate(op_grad, tensor.inputs[index], stop_tensor, nodes_to_compute, stop_gradients)
-        end.compact.reduce(:+)
+        end.compact
+
+        return nil if grads.empty?
+        grads.size > 1 ? ts.add_n(grads) : grads[0]
       else
         return nil if computed_op.nil?
         _propagate(computed_op, tensor.inputs[0], stop_tensor, nodes_to_compute, stop_gradients)
@@ -260,7 +263,11 @@ module TensorStream
           i_op(:softmax_grad, x, grad)
         when :softmax_cross_entropy_with_logits_v2
           output = node
-          [_broadcast_mul(grad, output[1]), nil]
+          logits = node.inputs[0]
+          [_broadcast_mul(grad, output[1]), -ts.nn.log_softmax(logits)]
+        when :sparse_softmax_cross_entropy_with_logits
+          output = node
+           [_broadcast_mul(grad, output[1]), nil]
         when :floor, :ceil
           # non differentiable
           nil
@@ -273,13 +280,51 @@ module TensorStream
         when :transpose
           return [ts.transpose(grad, ts.invert_permutation(y)), nil]
         when :index
-          grad
+          #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
+            multiplier = node.inputs[0].shape.shape[0]
+            filler = ts.zeros_like(grad)
+
+            res = Array.new(multiplier) { |index|
+              index == node.inputs[1].value ? grad : filler
+            }
+            [res]
+          end
+        when :squeeze
+          _reshape_to_input(node, grad)
+        when :expand_dims
+          [_reshape_to_input(node, grad), nil]
+        when :concat
+          _concat_grad_helper(node, grad, 1, node.inputs.size, 0)
+        when :reshape
+          [ts.reshape(grad, ts.shape(node.inputs[0])), nil]
+        when :stack
+          res = ts.unstack(grad, num: node.inputs.size, axis: node.options[:axis])
+          Array.new(node.inputs.size) { |i| res[i] }
+        when :unstack
+          ts.stack(grad, axis: node.options[:axis])
+        when :cast
+          t = %i[float16 float32 float64]
+          src_type = node.inputs[0].data_type
+          dst_type = grad.data_type
+
+          if t.key?(src_type) && t.key?(dst_type)
+            ts.cast(grad, src_type)
+          else
+            nil
+          end
         else
           raise "no derivative op for #{node.operation}"
         end
       end
     end
 
+    def self._reshape_to_input(node, grad)
+      ts.reshape(grad, ts.shape(node.inputs[0]))
+    end
+
     def self._broadcast_gradient_args(input_a, input_b)
       res = _op(:broadcast_gradient_args, input_a, input_b)
       [res[0], res[1]]
@@ -335,5 +380,41 @@ module TensorStream
       arr.each { |a| return true if a.equal?(obj) }
       false
     end
+
+    def self._extract_input_shapes(inputs)
+      sizes = []
+      fully_known = true
+      inputs.each do |x|
+        input_shape = ts.shape(x)
+        unless input_shape.is_const
+          fully_known = false
+          break
+        end
+        sizes << input_shape.value
+      end
+
+      if fully_known
+        sizes
+      else
+        ts.shape_n(inputs)
+      end
+    end
+
+    def self._concat_grad_helper(op, grad, start_value_index, end_value_index, dim_index)
+      # Degenerate concatenation, just return grad.
+      if op.inputs.size == 2
+        return end_value_index <= dim_index ? [grad] + [nil] : [nil] + [grad]
+      end
+      concat_dim = op.inputs[dim_index]
+      input_values = op.inputs[start_value_index..end_value_index]
+      non_neg_concat_dim = concat_dim % ts.rank(input_values[0])
+      sizes = _extract_input_shapes(input_values)
+
+      slicer = ts.slice(ts.stack(sizes, axis: 1), [non_neg_concat_dim, 0], [1, -1])
+      sizes = ts.squeeze(slicer)
+
+      out_grads = ts.split(grad, sizes, axis: non_neg_concat_dim, num: op.inputs.size - 1)
+      end_value_index <= dim_index ? out_grads + [nil] : [nil] + out_grads
+    end
   end
 end

data/lib/tensor_stream/nn/nn_ops.rb

@@ -25,11 +25,58 @@ module TensorStream
         logits = tf.convert_to_tensor(logits, name: 'logits')
         labels = tf.convert_to_tensor(labels, name: 'labels')
         labels = tf.cast(labels, logits.dtype)
+
         output = _op(:softmax_cross_entropy_with_logits_v2, logits, labels)
         output[0]
       end
     end
 
+    def self.sparse_softmax_cross_entropy_with_logits(labels: nil, logits: nil, name: nil)
+      TensorStream.name_scope(name, default: "SparseSoftmaxCrossEntropyWithLogits", values: [logits, labels]) do
+        tf = TensorStream
+        labels = tf.convert_to_tensor(labels)
+        logits = tf.convert_to_tensor(logits)
+        precise_logits = logits.data_type == :float16 ? tf.cast(logits, :float32) : logits
+
+        labels_static_shape = labels.shape
+        labels_shape = tf.shape(labels)
+        static_shapes_fully_defined = labels_static_shape.known? && logits.shape.known?
+
+        raise TensorStream::ValueError, "Logits cannot be scalars - received shape #{logits.shape.shape}." if logits.shape.known? && logits.shape.scalar?
+        raise TensorStream::ValueError, "Rank mismatch: Rank of labels (received #{labels_static_shape.ndims}) " +
+          "should equal rank of logits minus 1 (received #{logits.shape.ndims})." if logits.shape.known? && (labels_static_shape.known? && labels_static_shape.ndims != logits.shape.ndims - 1)
+
+        if logits.shape.ndims == 2
+          cost = _op(:sparse_softmax_cross_entropy_with_logits,
+              precise_logits, labels, name: name)
+          if logits.data_type == :float16
+            return tf.cast(cost[0], :float16)
+          else
+            return cost[0]
+          end
+        end
+
+        shape_checks = []
+        if !static_shapes_fully_defined
+          shape_checks << tf.append(tf.assert_equal(tf.rank(labels), tf.rank(logits) - 1))
+        end
+
+        tf.control_dependencies(shape_checks) do
+          num_classes = tf.shape(logits)[tf.rank(logits) - 1]
+          precise_logits = tf.reshape(precise_logits, [-1, num_classes])
+          labels = array_ops.reshape(labels, [-1])
+          cost = _op(:sparse_softmax_cross_entropy_with_logits, precise_logits, labels, name: name)
+          cost = tf.reshape(cost[0], labels_shape)
+
+          if logits.data_type == :float16
+            tf.cast(cost, :float16)
+          else
+            cost
+          end
+        end
+      end
+    end
+
     # Computes log softmax activations.
     def self.log_softmax(logits, axis: -1, name: nil)
       _op(:log_softmax, logits, axis: axis, name: name)

data/lib/tensor_stream/operation.rb

@@ -45,7 +45,7 @@ module TensorStream
     def infer_const
       return false if breakpoint
       case operation
-      when :random_standard_normal, :random_uniform, :glorot_uniform, :print
+      when :random_standard_normal, :random_uniform, :glorot_uniform, :print, :check_numerics
         false
       else
         non_const = @inputs.compact.find { |input| !input.is_const }
@@ -59,10 +59,12 @@ module TensorStream
         @inputs[1].data_type
       when :greater, :less, :equal, :not_equal, :greater_equal, :less_equal, :logical_and
         :boolean
-      when :shape, :rank
+      when :shape, :rank, :shape_n
         options[:out_type] || :int32
       when :random_standard_normal, :random_uniform, :glorot_uniform
         passed_data_type || :float32
+      when :concat
+        @inputs[1].data_type
       when :index
         if @inputs[0].is_a?(ControlFlow)
 
@@ -274,9 +276,28 @@ module TensorStream
         rank = inputs[0].shape.shape.size + 1
         axis = rank + axis if axis < 0
         rotated_shape = Array.new(axis + 1) { new_shape.shift }
-        rotated_shape.rotate! + new_shape
+        return rotated_shape.rotate! + new_shape
+      when :concat
+        return nil if inputs[0].value.nil?
+
+        axis = inputs[0].value # get axis
+
+        axis_size = 0
+
+        inputs[1..inputs.size].each do |input_item|
+          return nil if input_item.shape.shape.nil?
+          return nil if input_item.shape.shape[axis].nil?
+
+          axis_size += input_item.shape.shape[axis]
+        end
+
+        new_shape = inputs[1].shape.shape.dup
+        new_shape[axis] = axis_size
+        return new_shape
+      when :slice, :squeeze
+        return nil
       when :tile
-        nil
+        return nil
       else
         return nil if inputs[0].nil?
         return inputs[0].shape.shape if inputs.size == 1

data/lib/tensor_stream/ops.rb

@@ -1,6 +1,11 @@
 module TensorStream
   # Class that defines all available ops supported by TensorStream
   module Ops
+    class OutputHolder
+      def initialize(op)
+        @op = op
+      end
+    end
     FLOATING_POINT_TYPES = %i[float32 float64 float].freeze
     INTEGER_TYPES = %i[uint8 int32 int int64].freeze
     NUMERIC_TYPES = FLOATING_POINT_TYPES + INTEGER_TYPES
@@ -94,10 +99,29 @@ module TensorStream
     ##
     # This operation returns a 1-D integer tensor representing the shape of input
     def shape(input, name: nil, out_type: :int32)
-      return constant(shape_eval(input, out_type), dtype: out_type, name: name) if input.is_a?(Array)
+      return constant(shape_eval(input, out_type), dtype: out_type, name: name) if input.is_a?(Array) && !input[0].is_a?(Tensor)
       return constant(input.shape.shape, dtype: out_type, name: "Shape/#{input.name}") if shape_full_specified(input)
 
-      _op(:shape, input, nil, name: name, out_type: out_type)
+      _op(:shape, input, name: name, out_type: out_type)
+    end
+
+    def shape_n(inputs, name: nil, out_type: :int32)
+      shapes_known = true
+      inputs.each do |input|
+        unless input.shape.known?
+          shapes_known = false
+          break
+        end
+      end
+
+      if shapes_known
+        inputs.collect { |input| cons(input.shape.shape dtype: out_type) }
+      else
+        res = _op(:shape_n, *inputs, out_type: out_type, name: name)
+        Array.new(inputs.size) do |index|
+          res[index]
+        end
+      end
     end
 
     ##
@@ -113,15 +137,30 @@ module TensorStream
     ##
     # Returns the rank of a tensor.
     def rank(input, name: nil)
+      input = convert_to_tensor(input)
+      return cons(input.shape.ndims) if input.shape.known?
+
       _op(:rank, input, name: name)
     end
 
+    def constant_initializer(value, dtype: nil, verify_shape: false)
+      TensorStream::Initializer.new(-> { convert_to_tensor(value, dtype: dtype) })
+    end
+
     ##
     # initializer that generates tensors initialized to 0.
-    def zeros_initializer(dtype: nil)
+    #
+    def zeros_initializer(dtype: :float32)
       TensorStream::Initializer.new(-> { _op(:zeros, nil, nil, 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) })
+    end
+
     ##
     # The Glorot uniform initializer, also called Xavier uniform initializer.
     #
@@ -246,7 +285,62 @@ module TensorStream
     ##
     # Concatenates tensors along one dimension.
     def concat(values, axis, name: 'concat')
-      _op(:concat, *values, axis: axis, name: name)
+      if values.is_a?(Array)
+        _op(:concat, axis, *values, name: name)
+      else
+        _op(:concat, axis, values, name: name)
+      end
+    end
+
+    def split(value, num_or_size_splits, axis: 0, num: nil, name: 'split')
+      value = convert_to_tensor(value)
+      num_or_size_splits = convert_to_tensor(num_or_size_splits)
+      axis = convert_to_tensor(axis)
+
+      raise TensorStream::ValueError, "num_or_size_splits must be integer dtype" unless INTEGER_TYPES.include?(num_or_size_splits.data_type)
+
+      res = _op(:split, value, num_or_size_splits, axis, name: name)
+
+      pieces = if value.shape.known? && num_or_size_splits.is_const && num_or_size_splits.value && axis.is_const
+                  if num_or_size_splits.shape.scalar?
+                    raise TensorStream::ValueError, "num_or_size_splits must divide dimension #{value.shape.shape[axis.value]} evenly" unless value.shape.shape[axis.value] % num_or_size_splits.value == 0
+                    div = num_or_size_splits.value
+                    n = value.shape.shape[axis.value] / div
+
+                    Array.new(div) { |i|
+                      new_shape = value.shape.shape.dup
+                      new_shape[axis.value] = n
+                      new_shape
+                    }
+                  elsif num_or_size_splits.shape.ndims == 1
+                    raise TensorStream::ValueError, "Sum of splits do not match total dimen in axis #{value.shape.shape[axis.value]} != #{ num_or_size_splits.value.reduce(:+)}" if value.shape.shape[axis.value] != num_or_size_splits.value.reduce(:+)
+                    num_or_size_splits.value.collect do |v|
+                      new_shape = value.shape.shape.dup
+                      new_shape[axis.value] = v
+                      new_shape
+                    end
+                  else
+                    raise TensorStream::ValueError, "Scalar or 1D Tensor expected for num_or_size_splits"
+                  end
+                else
+                  raise TensorStream::ValueError, "Cannot automatically determine num, please specify num: in options" if num.nil?
+
+                  Array.new(num) { nil }
+                end
+
+      pieces.collect.with_index do |shape, i|
+        op = index(res, i, name: "split/index:#{i}")
+        if shape
+          op.shape = TensorShape.new(shape)
+        end
+        op
+      end
+    end
+
+    ##
+    # select an index in an array or a set of tensor outputs
+    def index(tensor, sel, name: nil)
+      _op(:index, tensor, sel, name: name)
     end
 
     ##
@@ -338,8 +432,11 @@ module TensorStream
     ##
     # Returns element-wise remainder of division.
     def mod(input_a, input_b, name: nil)
+      input_a = convert_to_tensor(input_a)
+      input_b = convert_to_tensor(input_b)
+
       input_a, input_b = check_data_types(input_a, input_b)
-      _op(:mod, input_a, input_b, name: name)
+       _op(:mod, input_a, input_b, name: name)
     end
 
     ##
@@ -393,8 +490,11 @@ module TensorStream
     end
 
     ##
-    # Casts a tensor to a new type.
+    # Casts a tensor to a new type, if needed
     def cast(input, dtype, name: nil)
+      input = convert_to_tensor(input)
+      return input if input.data_type == dtype
+
       _op(:cast, input, nil, data_type: dtype, name: name)
     end
 
@@ -630,14 +730,68 @@ module TensorStream
       _op(:gather, params, indices, validate_indices: validate_indices, name: name, axis: axis)
     end
 
+
+    ##
+    # Stacks a list of rank-R tensors into one rank-(R+1) tensor.
+    #
     def stack(values, axis: 0, name: 'stack')
       _op(:stack, *values, axis: axis, name: name)
     end
 
+    ##
+    # Unpacks the given dimension of a rank-R tensor into rank-(R-1) tensors.
+    #
+    def unstack(value, num: nil, axis: 0, name: 'unstack')
+      res = _op(:unstack, value, num: num, axis: axis, name: name)
+
+      num_vars = if value.shape.known?
+                   new_shape = value.shape.shape.dup
+                   rank = new_shape.size - 1
+                   axis = rank + axis if axis < 0
+                   rotated_shape = Array.new(axis + 1) { new_shape.shift }
+                   new_shape = rotated_shape.rotate!(-1) + new_shape
+                   new_shape[0]
+                 else
+                   raise TensorStream::ValueError, "num is unspecified and cannot be inferred." if num.nil?
+                   num
+                 end
+
+      return res[0] if num_vars == 1
+
+      Array.new(num_vars) do |i|
+        index(res, i, name: "unstack/index:#{i}")
+      end
+    end
+
+    ##
+    # Same as stack
+    def pack(values, axis: 0, name: 'pack')
+      _op(:stack, *values, axis: axis, name: name)
+    end
+
+    ##
+    # Same as unstack
+    #
+    def unpack(value, num: nil, axis: 0, name: 'unpack')
+      unstack(value, num: num, axis: axis, name: name)
+    end
+
+    ##
+    # Removes dimensions of size 1 from the shape of a tensor.
+    #
+    # Given a tensor input, this operation returns a tensor of the same type with all dimensions of size 1 removed.
+    # If you don't want to remove all size 1 dimensions, you can remove specific size 1 dimensions by specifying axis.
     def squeeze(value, axis: [], name: nil)
       _op(:squeeze, value, axis: axis, name: nil)
     end
 
+    ##
+    # Computes the difference between two lists of numbers or strings.
+    # Given a list x and a list y, this operation returns a list out that represents all values
+    # that are in x but not in y. The returned list out is sorted in the same order that the numbers appear
+    # in x (duplicates are preserved). This operation also returns a list idx that represents the position of
+    # each out element in x. In other words:
+    #
     def setdiff1d(x, y, index_dtype: :int32, name: nil)
       result = _op(:setdiff1d, x, y, index_dtype: index_dtype, name: name)
       [result[0], result[1]]