tensor_stream 0.9.2 → 0.9.5

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 9601653d86556739c89b591768e9d54d13d6335dd2f953fff7d91f22636e8c7b
-  data.tar.gz: 3db4119c9e752df77cbbf7f8f753050a012b03fd2bde04dd6bbfd01f8903af62
+  metadata.gz: 51cb6686663dece94714073ff12b13ad7e57de1aadbee44506897c69a2d5fd67
+  data.tar.gz: 47152a908b2cc966cba6721da8f0c3170e3198a70f416fd224caee7d35a1fca2
 SHA512:
-  metadata.gz: 3ca3eec8ce6cc7e73a1d5e38ccaaf2f65f7a37bb84e0118766de05ec374557847a88f2f14b1e35b896010d7e68a0234589fcf3de95a9417840e2e77ab4a4ea58
-  data.tar.gz: 6b109a2ee9b59e286a2d06d1c3166fc5bbda9a888f0c26aa07bc4f994753ce5dbd7cb0aebb433a7446ddef653f3096ae1b60dbc4610e0dc18c9aaca8d3175f2f
+  metadata.gz: 7fff67042fd35c651409e04e890feafab25e618d88d4c3f69c97b39e839a943c64a4e05be388a4464e32328a777f6ee56dc5ede0a6c69c977cc94bcb7cd38bad
+  data.tar.gz: a30230994e81062626eaae5ca7a13ece36d2a5ea6210cfe1e3e465b5e5e24adad525547223f27e15991f74728d53d8b02e7cac0710fe16bf3c602402f5b65b5c

data/lib/tensor_stream/evaluator/base_evaluator.rb

@@ -140,6 +140,8 @@ module TensorStream
             @context[:profile][:operations][tensor.name] = { op: tensor.operation,
                                                              step: @context[:profile][:step],
                                                              eval_time: end_time - start_time,
+                                                             shape: tensor.shape ? tensor.shape.shape : nil,
+                                                             data_type: tensor.data_type,
                                                              tensor: tensor }
           end
         end
@@ -166,6 +168,7 @@ module TensorStream
       def global_eval(tensor, input, execution_context, op_options = {})
         return nil unless input
         return input unless input.is_a?(Tensor)
+
         # puts "global eval #{tensor.name}"
         @context[:_cache][:placement][input.name] = @session.assign_evaluator(input) if @context[:_cache][:placement][input.name].nil?
         if !on_same_device?(input) # tensor is on another device or evaluator

data/lib/tensor_stream/evaluator/operation_helpers/array_ops_helper.rb

@@ -310,5 +310,30 @@ module TensorStream
 
       reduce_axis(0, axis, val, keep_dims, func)
     end
+
+    def arr_pad(arr, paddings, data_type = :float32, rank = 0)
+      raise "padding #{paddings[rank]} needs to have to elements [before, after]" if paddings[rank].size != 2
+
+      before = paddings[rank][0]
+      after = paddings[rank][1]
+      pad_value = fp_type?(data_type) ? 0.0 : 0
+      if arr[0].is_a?(Array)
+        next_dim_elem = arr.collect { |a| arr_pad(a, paddings, data_type, rank + 1) }
+        padding = deep_dup_array(next_dim_elem[0], pad_value)
+        Array.new(before) { padding } + next_dim_elem + Array.new(after) { padding }
+      else
+        Array.new(before) { pad_value } + arr + Array.new(after) { pad_value }
+      end
+    end
+
+    def deep_dup_array(arr, value = nil)
+      if arr.is_a?(Array)
+        arr.dup.collect do |a|
+          deep_dup_array(a, value)
+        end
+      else
+        value.nil? ? arr : value
+      end
+    end
   end
 end

data/lib/tensor_stream/evaluator/ruby/array_ops.rb

@@ -132,9 +132,9 @@ module TensorStream
           axis = !tensor.options[:axis].is_a?(Array) ? [tensor.options[:axis]] : tensor.options[:axis]
 
           if !axis.empty?
-            axis.each do |axis|
-              if shape[axis] == 1
-                shape[axis] = nil
+            axis.each do |x|
+              if shape[x] == 1
+                shape[x] = nil
               else
                 raise TensorStream::ValueError, "unable to squeeze dimension that does not have a size of 1"
               end
@@ -201,6 +201,7 @@ module TensorStream
           out = []
           input.each_with_index do |x, index|
             next if remove.include?(x)
+
             out << x
             idx << index
           end
@@ -225,6 +226,7 @@ module TensorStream
             break if start == limit
             break if (start < limit) && (cur_step >= limit)
             break if (start > limit) && (cur_step <= limit)
+
             r << cur_step
             cur_step += delta
           end
@@ -281,34 +283,34 @@ module TensorStream
           get_rank(inputs[0])
         end
 
-        register_op :split  do |context, tensor, inputs|
+        register_op :split  do |_context, tensor, inputs|
           value, num_split, axis = inputs
 
           value_shape = shape_eval(value)
           res = if num_split.is_a?(Array)
-            begin_index = 0
-            num_split.collect do |num|
-              end_index = begin_index + num
-              arr = split_tensor(value, begin_index, end_index, axis)
-              begin_index = end_index
-              arr
-            end
-          else
-            raise TensorStream::ValueError, "#{num_split} does not divide #{value_shape[axis]} evenly" if value_shape[axis] % num_split != 0
-            piece_sizes = value_shape[axis] / num_split
-            Array.new(num_split) do |num|
-              begin_index = num * piece_sizes
-              end_index = begin_index + piece_sizes
-              split_tensor(value, begin_index, end_index, axis)
-            end
-          end
+                  begin_index = 0
+                  num_split.collect do |num|
+                    end_index = begin_index + num
+                    arr = split_tensor(value, begin_index, end_index, axis)
+                    begin_index = end_index
+                    arr
+                  end
+                else
+                  raise TensorStream::ValueError, "#{num_split} does not divide #{value_shape[axis]} evenly" if value_shape[axis] % num_split != 0
+
+                  piece_sizes = value_shape[axis] / num_split
+                  Array.new(num_split) do |num|
+                    begin_index = num * piece_sizes
+                    end_index = begin_index + piece_sizes
+                    split_tensor(value, begin_index, end_index, axis)
+                  end
+                end
           TensorStream::Evaluator::OutputGroup.new(res, res.map { tensor.inputs[0].data_type })
         end
 
         register_op :reshape do |_context, _tensor, inputs|
           arr, new_shape = inputs
           arr = [arr] unless arr.is_a?(Array)
-
           flat_arr = arr.flatten
           if new_shape.size.zero? && flat_arr.size == 1
             flat_arr[0]
@@ -318,9 +320,7 @@ module TensorStream
         end
 
         register_op :pad do |context, tensor, inputs|
-          p = complete_eval(tensor.options[:paddings], context)
-
-          arr_pad(inputs[0], p, tensor.data_type)
+          arr_pad(inputs[0], inputs[1], tensor.data_type)
         end
 
         register_op :tile do |_context, _tensor, inputs|

data/lib/tensor_stream/evaluator/ruby/check_ops.rb

@@ -2,7 +2,15 @@ module TensorStream
   module CheckOps
     def CheckOps.included(klass)
       klass.class_eval do
+        register_op :assert_equal do |context, tensor, inputs|
+          result = call_vector_op(tensor, :equal, inputs[0], inputs[1], context, ->(t, u) { t == u })
 
+          result = result.is_a?(Array) ? result.flatten.uniq : [result]
+          prefix = tensor.options[:message] || ""
+          raise TensorStream::InvalidArgumentError, "#{prefix} #{tensor.inputs[0].name} != #{tensor.inputs[1].name}" if result != [true]
+
+          nil
+        end
       end
     end
   end

data/lib/tensor_stream/evaluator/ruby/images_ops.rb

@@ -14,30 +14,28 @@ module TensorStream
           image.grayscale! if channels == 1
           image_data = image.pixels.collect do |pixel|
             color_values = if channels == 4
-              [ ChunkyPNG::Color.r(pixel),
-                ChunkyPNG::Color.g(pixel),
-                ChunkyPNG::Color.b(pixel),
-                ChunkyPNG::Color.a(pixel) ]
-            elsif channels == 3
-              [ ChunkyPNG::Color.r(pixel),
-                ChunkyPNG::Color.g(pixel),
-                ChunkyPNG::Color.b(pixel)]
-            elsif channels == 1
-              [ ChunkyPNG::Color.r(pixel) ]
-            else
-              raise "Invalid channel value #{channels}"
-            end
-            
-            if fp_type?(tensor.data_type)
-              color_values.map! { |v| v.to_f }
-            end
+                             [ChunkyPNG::Color.r(pixel),
+                              ChunkyPNG::Color.g(pixel),
+                              ChunkyPNG::Color.b(pixel),
+                              ChunkyPNG::Color.a(pixel)]
+                           elsif channels == 3
+                             [ChunkyPNG::Color.r(pixel),
+                              ChunkyPNG::Color.g(pixel),
+                              ChunkyPNG::Color.b(pixel)]
+                           elsif channels == 1
+                             [ChunkyPNG::Color.r(pixel)]
+                           else
+                             raise "Invalid channel value #{channels}"
+                           end
+
+            color_values.map!(&:to_f) if fp_type?(tensor.data_type)
 
             color_values
           end
           TensorShape.reshape(image_data, [image.height, image.width, channels])
         end
 
-        register_op :encode_png do |_context, tensor, inputs|
+        register_op :encode_png do |_context, _tensor, inputs|
           image_data = inputs[0]
           height, width, channels = shape_eval(image_data)
 

data/lib/tensor_stream/evaluator/ruby/math_ops.rb

@@ -140,17 +140,33 @@ module TensorStream
         end
 
         register_op(%i[argmax arg_max]) do |_context, tensor, inputs|
-          axis = tensor.options[:axis] || 0
+          axis = inputs[1] || 0
           rank = get_rank(inputs[0])
           raise TensorStream::InvalidArgumentError, "Expected dimension in the range [#{-rank},#{rank}) but got #{axis}" if axis < -rank || axis >= rank
-          get_op_with_axis(inputs[0], axis, 0, tensor.data_type)
+
+          new_shape = shape_eval(inputs[0])
+          ns = new_shape.each_with_index.collect do |shape, index|
+            next nil if index == axis
+
+            shape
+          end.compact
+
+          Tensor.cast_dtype(TensorShape.reshape(get_op_with_axis(inputs[0], axis, 0, :max), ns), tensor.data_type)
         end
 
         register_op(%i[argmin arg_min]) do |_context, tensor, inputs|
-          axis = tensor.options[:axis] || 0
+          axis = inputs[1] || 0
           rank = get_rank(inputs[0])
           raise TensorStream::InvalidArgumentError, "Expected dimension in the range [#{-rank},#{rank}) but got #{axis}" if axis < -rank || axis >= rank
-          get_op_with_axis(inputs[0], axis, 0, tensor.data_type, ->(a, b) { a < b })
+
+          new_shape = shape_eval(inputs[0])
+          ns = new_shape.each_with_index.collect do |shape, index|
+            next nil if index == axis
+
+            shape
+          end.compact
+
+          Tensor.cast_dtype(TensorShape.reshape(get_op_with_axis(inputs[0], axis, 0, :min), ns), tensor.data_type)
         end
 
         register_op :cumprod do |context, tensor, inputs|

data/lib/tensor_stream/evaluator/ruby/nn_ops.rb

@@ -16,11 +16,11 @@ module TensorStream
           assign = tensor.inputs[0] || tensor
           assign_acc = tensor.inputs[1]
           assign_acc.value = multi_array_op(->(t, u) { t * momentum + u }, momentum_var, grad)
-          if tensor.options[:use_nesterov]
-            assign.value = multi_array_op(->(v, g, acc) { v - (g * learning_rate + acc * momentum * learning_rate) }, target_var, grad, momentum_var)
-          else
-            assign.value = multi_array_op(->(v, acc) { v - acc * learning_rate }, target_var, momentum_var)
-          end
+          assign.value = if tensor.options[:use_nesterov]
+                           multi_array_op(->(v, g, acc) { v - (g * learning_rate + acc * momentum * learning_rate) }, target_var, grad, momentum_var)
+                         else
+                           multi_array_op(->(v, acc) { v - acc * learning_rate }, target_var, momentum_var)
+                         end
           assign.value
         end
 
@@ -207,6 +207,10 @@ module TensorStream
             TensorShape.reshape(arr, input_shape)
           end
         end
+
+        register_op :relu6 do |context, tensor, inputs|
+          call_vector_op(tensor, :relu6, inputs[0], inputs[1], context, ->(t, u) { [[t, 0].max, 6].min })
+        end
       end
     end
   end

data/lib/tensor_stream/evaluator/ruby/random_ops.rb

@@ -25,24 +25,24 @@ module TensorStream
           generate_vector(shape, generator: generator)
         end
 
-        register_op :random_uniform, no_eval: true do |_context, tensor, _inputs|
+        register_op :random_uniform, no_eval: true do |_context, tensor, inputs|
           maxval = tensor.options.fetch(:maxval, 1)
           minval = tensor.options.fetch(:minval, 0)
           seed = tensor.options[:seed]
 
           random = _get_randomizer(tensor, seed)
           generator = -> { random.rand * (maxval - minval) + minval }
-          shape = tensor.options[:shape] || tensor.shape.shape
+          shape = inputs[0] || tensor.shape.shape
           generate_vector(shape, generator: generator)
         end
 
-        register_op :random_standard_normal, no_eval: true do |_context, tensor, _inputs|
+        register_op :random_standard_normal, no_eval: true do |_context, tensor, inputs|
           seed = tensor.options[:seed]
           random = _get_randomizer(tensor, seed)
           r = RandomGaussian.new(tensor.options.fetch(:mean), tensor.options.fetch(:stddev), -> { random.rand })
           random = _get_randomizer(tensor, seed)
           generator = -> { r.rand }
-          shape = tensor.options[:shape] || tensor.shape.shape
+          shape = inputs[0] || tensor.shape.shape
           generate_vector(shape, generator: generator)
         end
       end

data/lib/tensor_stream/evaluator/ruby_evaluator.rb

@@ -289,7 +289,7 @@ module TensorStream
           # assertions to make sure inferred shapes == actual evaluated shapes
           if tensor.shape.known? && (result.is_a?(Array) || result.is_a?(Float) || result.is_a?(Integer))
             if shape_eval(result) != tensor.shape.shape
-              raise "assert error #{tensor.name} #{shape_eval(result)} != #{tensor.shape.shape}"
+              # raise "assert error #{tensor.name} #{shape_eval(result)} != #{tensor.shape.shape}"
             end
           end
 
@@ -313,21 +313,20 @@ module TensorStream
           @context[tensor.name] = result
         end
       rescue EvaluatorExcecutionException => e
-        raise e, "error #{e.message} while evaluating #{tensor.name} : #{tensor.to_math(true, 1)} defined at #{tensor.source}"
+        raise e, "error #{e.message} while evaluating #{tensor.name}  defined at #{tensor.source}"
       rescue TensorStreamError => e
-        raise e, "error #{e.message} while evaluating #{tensor.name} : #{tensor.to_math(true, 1)} defined at #{tensor.source}"
+        raise e, "error #{e.message} while evaluating #{tensor.name}  defined at #{tensor.source}"
       rescue StandardError => e
-
-        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]
+        # 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]
         puts e.message
         puts e.backtrace.join("\n")
         # shape_a = a.shape.shape if a
         # shape_b = b.shape.shape if b
         # dtype_a = a.data_type if a
         # dtype_b = b.data_type if b
-        a = complete_eval(a, child_context)
-        b = complete_eval(b, child_context)
+        # a = complete_eval(a, child_context)
+        # b = complete_eval(b, child_context)
         # puts "name: #{tensor.given_name}"
         # # puts "op: #{tensor.to_math(true, 1)}"
         # puts "A #{shape_a} #{dtype_a}: #{a}" if a
@@ -363,36 +362,17 @@ module TensorStream
 
       private
 
-      def get_op_with_axis(a, target_axis, current_axis, output_type, op = ->(t, u) { t > u })
-        if target_axis == current_axis
-          if a[0].is_a?(Array)
-            (0...a[0].size).each.collect do |column_index|
-              max = nil
-              max_index = 0
-              a.each_with_index do |row, row_index|
-                if max.nil? || op.call(row[column_index], max)
-                  max = row[column_index]
-                  max_index = row_index
-                end
-              end
+      def get_op_with_axis(a, target_axis, current_axis, op)
+        rank = get_rank(a)
+        return a.index(a.send(:"#{op}")) if rank == 1
 
-              Tensor.cast_dtype(max_index, output_type)
-            end
-          else
-            max = nil
-            max_index = 0
-            a.each_with_index do |x, index|
-              if max.nil? || op.call(x, max)
-                max = x
-                max_index = index
-              end
-            end
-            Tensor.cast_dtype(max_index, output_type)
-          end
+        if current_axis == target_axis
+          compare_items = a.collect(&:flatten).transpose
+          compare_items.map { |item| item.index(item.send(:"#{op}")) }
+        elsif a[0].is_a?(Array)
+          a.map { |item| get_op_with_axis(item, target_axis, current_axis + 1, op) }
         else
-          a.collect do |row|
-            get_op_with_axis(row, target_axis, current_axis + 1, output_type, op)
-          end
+          return a.index(a.send(:"#{op}"))
         end
       end
 
@@ -403,31 +383,6 @@ module TensorStream
         reduce(val, axis, keep_dims, func)
       end
 
-      def arr_pad(arr, paddings, data_type = :float32, rank = 0)
-        raise "padding #{paddings[rank]} needs to have to elements [before, after]" if paddings[rank].size != 2
-
-        before = paddings[rank][0]
-        after = paddings[rank][1]
-        pad_value = fp_type?(data_type) ? 0.0 : 0
-        if arr[0].is_a?(Array)
-          next_dim_elem = arr.collect { |a| arr_pad(a, paddings, data_type, rank + 1) }
-          padding = deep_dup_array(next_dim_elem[0], pad_value)
-          Array.new(before) { padding } + next_dim_elem + Array.new(after) { padding }
-        else
-          Array.new(before) { pad_value } + arr + Array.new(after) { pad_value }
-        end
-      end
-
-      def deep_dup_array(arr, value = nil)
-        if arr.is_a?(Array)
-          arr.dup.collect do |a|
-            deep_dup_array(a, value)
-          end
-        else
-          value.nil? ? arr : value
-        end
-      end
-
       def call_op(op, a, child_context, func)
         a = complete_eval(a, child_context)
         process_function_op(a, func)

data/lib/tensor_stream/graph_builder.rb

@@ -14,6 +14,7 @@ module TensorStream
       parsed_tree = protobuf.load_from_string(buffer)
       parsed_tree.each do |node|
         next unless node['type'] == 'node'
+
         # puts "build #{node['name']}"
         options = protobuf.options_evaluator(node)
         options[:name] = node['name']

data/lib/tensor_stream/graph_serializers/graphml.rb

@@ -34,7 +34,7 @@ module TensorStream
       arr_buf << "</node>"
 
       to_graph_ml(tensor, arr_buf, {}, groups)
-      #dump groups
+      # dump groups
       groups.each do |k, g|
         arr_buf << create_group(k, k, g)
       end

data/lib/tensor_stream/graph_serializers/pbtext.rb

@@ -15,6 +15,7 @@ module TensorStream
           @lines << "  op: #{camelize(node.operation.to_s).to_json}"
           node.inputs.each do |input|
             next unless input
+
             @lines << "  input: #{input.name.to_json}"
           end
           # type

data/lib/tensor_stream/helpers/infer_shape.rb

@@ -0,0 +1,182 @@
+require 'tensor_stream/evaluator/operation_helpers/array_ops_helper'
+module TensorStream
+  ##
+  # Convenience class for guessing the shape of a tensor
+  #
+  class InferShape
+    extend TensorStream::ArrayOpsHelper
+    extend TensorStream::OpHelper
+
+    def self.infer_shape(tensor)
+      case tensor.operation
+      when :assign
+        possible_shape = if tensor.inputs[0] && tensor.inputs[0].shape.shape
+                           tensor.inputs[0].shape.shape
+                         else
+                           tensor.inputs[1].shape.shape
+                         end
+
+        possible_shape
+      when :index
+        return nil unless tensor.inputs[0].is_a?(Tensor)
+        return nil unless tensor.inputs[0].const_value
+
+        input_shape = tensor.inputs[0].shape
+        return nil unless input_shape.known?
+
+        s = input_shape.shape.dup
+        s.shift
+        s
+      when :arg_min, :argmax, :argmin
+        return nil unless tensor.inputs[0].shape.known?
+        return nil if tensor.inputs[1] && tensor.inputs[1].value.nil?
+
+        axis = tensor.inputs[1].nil? ? 0 : tensor.inputs[1].value
+        new_shape = tensor.inputs[0].shape.shape
+        new_shape.each_with_index.collect do |shape, index|
+          next nil if index == axis
+
+          shape
+        end.compact
+      when :mean, :prod, :sum, :arg_max
+        return [] if tensor.inputs[1].nil?
+        return nil if tensor.inputs[0].nil?
+        return nil unless tensor.inputs[0].shape.known?
+
+        input_shape = tensor.inputs[0].shape.shape
+        rank = input_shape.size
+
+        axis = tensor.inputs[1].const_value
+        return nil if axis.nil?
+
+        axis = [axis] unless axis.is_a?(Array)
+        axis = axis.map { |a| a < 0 ? rank - a.abs : a }
+
+        input_shape.each_with_index.map do |item, index|
+          if axis.include?(index)
+            next 1 if tensor.options[:keepdims]
+
+            next nil
+          end
+          item
+        end.compact
+      when :reshape
+        new_shape = tensor.inputs[1] && tensor.inputs[1].value ? tensor.inputs[1].value : nil
+        return nil if new_shape.nil?
+        return nil if tensor.inputs[0].shape.nil?
+
+        input_shape = tensor.inputs[0].shape.shape
+        return new_shape if input_shape.nil?
+        return nil if input_shape.include?(nil)
+        TensorShape.fix_inferred_elements(new_shape, input_shape.reduce(:*))
+      when :flow_group
+        []
+      when :zeros, :ones, :fill, :random_standard_normal, :random_uniform
+        a_shape = tensor.inputs[0] ? tensor.inputs[0].const_value : tensor.options[:shape]
+        return nil if a_shape.nil?
+        a_shape.is_a?(Array) ? a_shape : [a_shape]
+      when :zeros_like, :ones_like
+        tensor.inputs[0].shape.shape
+      when :shape
+        tensor.inputs[0].shape.shape ? [tensor.inputs[0].shape.shape.size] : nil
+      when :pad
+        return nil unless tensor.inputs[0].shape.known?
+        return nil unless tensor.inputs[1].value
+
+        size = tensor.inputs[0].shape.shape.reduce(:*) || 1
+        dummy_tensor_for_shape = TensorShape.reshape(Array.new(size), tensor.inputs[0].shape)
+        shape_eval(arr_pad(dummy_tensor_for_shape, tensor.inputs[1].value))
+      when :mat_mul
+        return nil if tensor.inputs[0].shape.shape.nil? || tensor.inputs[1].shape.shape.nil?
+        return [] if tensor.inputs[0].shape.shape.empty? || tensor.inputs[1].shape.shape.empty?
+        return nil if tensor.inputs[0].shape.shape.size != 2 || tensor.inputs[1].shape.shape.size != 2
+
+        shape1, m = if tensor.options[:transpose_a]
+                      [tensor.inputs[0].shape.shape[0], tensor.inputs[0].shape.shape[1]]
+                    else
+                      [tensor.inputs[0].shape.shape[1], tensor.inputs[0].shape.shape[0]]
+                    end
+
+        shape2, n = if tensor.options[:transpose_b]
+                      [tensor.inputs[1].shape.shape[1], tensor.inputs[1].shape.shape[0]]
+                    else
+                      [tensor.inputs[1].shape.shape[0], tensor.inputs[1].shape.shape[1]]
+                    end
+
+        return nil if shape1.nil? || shape2.nil? || shape1 < 0 || shape2 < 0
+
+        raise TensorStream::ValueError, "incompatible shape sizes for matrix multiplication (#{shape1} != #{shape2}) #{tensor.inputs[0].shape.shape} vs #{tensor.inputs[1].shape.shape}" if shape1 != shape2
+
+        [m, n]
+      when :transpose
+        return nil unless shape_full_specified(tensor.inputs[0])
+        return nil if tensor.inputs[1].is_a?(Tensor)
+
+        rank = tensor.inputs[0].shape.shape.size
+        perm = tensor.inputs[1] || (0...rank).to_a.reverse
+        perm.map { |p| tensor.inputs[0].shape.shape[p] }
+      when :stack
+        return nil unless shape_full_specified(tensor.inputs[0])
+
+        axis = tensor.options[:axis] || 0
+        new_shape = [tensor.inputs.size]
+        tensor.inputs[0].shape.shape.inject(new_shape) { |ns, i| ns << i }
+        rank = tensor.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
+      when :concat
+        return nil if tensor.inputs[0].value.nil?
+
+        axis = tensor.inputs[0].value # get axis
+
+        axis_size = 0
+
+        tensor.inputs[1..tensor.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 = tensor.inputs[1].shape.shape.dup
+        new_shape[axis] = axis_size
+        new_shape
+      when :slice, :squeeze
+        nil
+      when :tile
+        nil
+      when :expand_dims
+        nil
+      when :broadcast_gradient_args
+        nil
+      when :no_op
+        nil
+      when :softmax_cross_entropy_with_logits_v2, :sparse_softmax_cross_entropy_with_logits
+        nil
+      when :decode_png, :flow_dynamic_stitch, :dynamic_stitch, :gather
+        nil
+      when :eye
+        return [tensor.inputs[0].const_value, tensor.inputs[1].const_value] if tensor.inputs[0].const_value && tensor.inputs[1].const_value
+
+        nil
+      when :size
+        []
+      when :unstack
+        return nil unless tensor.inputs[0].shape.known?
+
+        new_shape = tensor.inputs[0].shape.shape.dup
+        rank = new_shape.size - 1
+        axis = tensor.options[:axis] || 0
+        axis = rank + axis if axis < 0
+        rotated_shape = Array.new(axis + 1) { new_shape.shift }
+        rotated_shape.rotate!(-1) + new_shape
+      else
+        return nil if tensor.inputs[0].nil?
+        return tensor.inputs[0].shape.shape if tensor.inputs.size == 1
+
+        TensorShape.infer_shape(tensor.inputs[0].shape.shape, tensor.inputs[1].shape.shape) if tensor.inputs.size == 2 && tensor.inputs[0] && tensor.inputs[1]
+      end
+    end
+  end
+end

data/lib/tensor_stream/helpers/op_helper.rb

@@ -56,8 +56,8 @@ module TensorStream
 
     def format_source(trace)
       grad_source = trace.select { |c| c.to_s.include?(File.join('lib', 'tensor_stream', 'math_gradients')) }.first
-      source = trace.reject { |c| c.to_s.include?(File.join('lib', 'tensor_stream')) }.first
-      [grad_source, source].compact.join("\n")
+      # source = trace.reject { |c| c.to_s.include?(File.join('lib', 'tensor_stream')) }.first
+      [grad_source, trace].compact.join("\n")
     end
 
     def shapes_fully_specified_and_equal(x, y)

data/lib/tensor_stream/images.rb

@@ -8,7 +8,7 @@ module TensorStream
     end
 
     def self.encode_png(contents, compression: -1, name: nil)
-      check_allowed_types(contents, [:uint8, :uint16])
+      check_allowed_types(contents, %i[uint8 uint16])
       contents = convert_to_tensor(contents, dtype: :uint16)
       _op(:encode_png, contents, compression: compression, name: name)
     end