tensor_stream 0.7.0 → 0.8.0

data/lib/tensor_stream/exceptions.rb

@@ -3,4 +3,5 @@ module TensorStream
   class KeyError < TensorStreamError; end
   class ValueError < TensorStreamError; end
   class InvalidArgumentError < TensorStreamError; end
+  class NotImplementedError < TensorStreamError; end
 end

data/lib/tensor_stream/graph_builder.rb

@@ -36,9 +36,8 @@ module TensorStream
           TensorStream::Placeholder.new(options[:dtype] || options[:T], nil, shape, options)
         else
           op = underscore(node['op']).to_sym
-          unless TensorStream::Evaluator::RubyEvaluator.ops.keys.include?(op)
-            puts "warning unsupported op #{op}"
-          end
+          puts "warning unsupported op #{op}" unless TensorStream::Evaluator::RubyEvaluator.ops.key?(op)
+
           # map input tensor
           inputs = node['input'].map do |input|
             input[0] = '' if input.start_with?('^')

data/lib/tensor_stream/graph_deserializers/protobuf.rb

@@ -23,15 +23,15 @@ module TensorStream
     end
 
     def parse_value(value_node)
-      if value_node['tensor']
-        evaluate_tensor_node(value_node['tensor'])
-      end
+      return unless value_node['tensor']
+
+      evaluate_tensor_node(value_node['tensor'])
     end
 
     def evaluate_tensor_node(node)
       if !node['shape'].empty? && node['tensor_content']
         content = node['tensor_content']
-        unpacked = eval(%Q{"#{content}"})
+        unpacked = eval(%Q("#{content}"))
 
         if node['dtype'] == 'DT_FLOAT'
           TensorShape.reshape(unpacked.unpack('f*'), node['shape'])
@@ -45,14 +45,14 @@ module TensorStream
       else
 
         val = if node['dtype'] == 'DT_FLOAT'
-          node['float_val'] ? node['float_val'].to_f : []
-        elsif node['dtype'] == 'DT_INT32'
-          node['int_val'] ? node['int_val'].to_i : []
-        elsif node['dtype'] == 'DT_STRING'
-          node['string_val']
-        else
-          raise "unknown dtype #{node['dtype']}"
-        end
+                node['float_val'] ? node['float_val'].to_f : []
+              elsif node['dtype'] == 'DT_INT32'
+                node['int_val'] ? node['int_val'].to_i : []
+              elsif node['dtype'] == 'DT_STRING'
+                node['string_val']
+              else
+                raise "unknown dtype #{node['dtype']}"
+              end
 
         if node['shape'] == [1]
           [val]
@@ -83,7 +83,7 @@ module TensorStream
       return {} if node['attributes'].nil?
 
       node['attributes'].map do |attribute|
-        attr_type, attr_value = attribute['value'].collect { |k, v| [k, v] }.flatten(1)
+        attr_type, attr_value = attribute['value'].flat_map { |k, v| [k, v] }
 
         if attr_type == 'tensor'
           attr_value = evaluate_tensor_node(attr_value)
@@ -103,11 +103,10 @@ module TensorStream
       block = []
       node = {}
       node_attr = {}
-      dim = []
       state = :top
 
       lines.each do |str|
-        case(state)
+        case state
         when :top
           node['type'] = parse_node_name(str)
           state = :node_context
@@ -177,7 +176,7 @@ module TensorStream
             next
           else
             key, value = str.split(':', 2)
-            node_attr['value'] << { key => value}
+            node_attr['value'] << { key => value }
           end
         when :tensor_context
           if str == 'tensor_shape {'
@@ -219,7 +218,7 @@ module TensorStream
             state = :shape_context
             next
           else
-            key, value = str.split(':', 2)
+            _key, value = str.split(':', 2)
             node_attr['value']['shape'] << value.strip.to_i
           end
         when :tensor_shape_dim_context
@@ -227,7 +226,7 @@ module TensorStream
             state = :tensor_shape_context
             next
           else
-            key, value = str.split(':', 2)
+            _key, value = str.split(':', 2)
             node_attr['value']['tensor']['shape'] << value.strip.to_i
           end
         end
@@ -237,7 +236,7 @@ module TensorStream
     end
 
     def parse_node_name(str)
-      name = str.split(' ')[0]
+      str.split(' ')[0]
     end
 
     def process_value(value)
@@ -253,19 +252,19 @@ module TensorStream
       'n' => "\x0a", 'v' => "\x0b", 'f' => "\x0c",
       'r' => "\x0d", 'e' => "\x1b", "\\\\" => "\x5c",
       "\"" => "\x22", "'" => "\x27"
-    }
+    }.freeze
 
     def unescape(str)
       # Escape all the things
-      str.gsub(/\\(?:([#{UNESCAPES.keys.join}])|u([\da-fA-F]{4}))|\\0?x([\da-fA-F]{2})/) {
+      str.gsub(/\\(?:([#{UNESCAPES.keys.join}])|u([\da-fA-F]{4}))|\\0?x([\da-fA-F]{2})/) do
         if $1
           $1 == '\\' ? '\\' : UNESCAPES[$1]
         elsif $2 # escape \u0000 unicode
-          ["#$2".hex].pack('U*')
+          ["#{$2}".hex].pack('U*')
         elsif $3 # escape \0xff or \xff
           [$3].pack('H2')
         end
-      }
+      end
     end
   end
 end

data/lib/tensor_stream/graph_serializers/graphml.rb

@@ -83,7 +83,7 @@ module TensorStream
       arr_buf << '<y:GroupNode>'
       arr_buf << '<y:Fill color="#CAECFF84" transparent="false"/>'
       arr_buf << '<y:BorderStyle color="#666699" type="dotted" width="1.0"/>'
-      arr_buf << '<y:NodeLabel alignment="right" autoSizePolicy="node_width" backgroundColor="#99CCFF" borderDistance="0.0" fontFamily="Dialog" fontSize="15" fontStyle="plain" hasLineColor="false" height="21.4609375" horizontalTextPosition="center" iconTextGap="4" modelName="internal" modelPosition="t" textColor="#000000" verticalTextPosition="bottom" visible="true" width="67.18603515625" x="-8.593017578125" y="0.0">'+ title + '</y:NodeLabel>'
+      arr_buf << '<y:NodeLabel alignment="right" autoSizePolicy="node_width" backgroundColor="#99CCFF" borderDistance="0.0" fontFamily="Dialog" fontSize="15" fontStyle="plain" hasLineColor="false" height="21.4609375" horizontalTextPosition="center" iconTextGap="4" modelName="internal" modelPosition="t" textColor="#000000" verticalTextPosition="bottom" visible="true" width="67.18603515625" x="-8.593017578125" y="0.0">' + title + '</y:NodeLabel>'
       arr_buf << '<y:Shape type="roundrectangle"/>'
       arr_buf << '</y:GroupNode>'
       arr_buf << '</y:Realizers>'
@@ -146,9 +146,9 @@ module TensorStream
           input_buf << "<node id=\"#{_gml_string(input.name)}\">"
           input_buf << "<data key=\"d0\">#{input.name}</data>"
           input_buf << "<data key=\"d2\">green</data>"
-          if @last_session_context[input.name]
-            input_buf << "<data key=\"d3\">#{_val(tensor)}</data>"
-          end
+
+          input_buf << "<data key=\"d3\">#{_val(tensor)}</data>" if @last_session_context[input.name]
+
           input_buf << "<data key=\"d9\">"
           input_buf << "<y:ShapeNode>"
           input_buf << "  <y:Fill color=\"#33CCCC\" transparent=\"false\"/>"
@@ -164,9 +164,9 @@ module TensorStream
           input_buf << "  <y:NodeLabel alignment=\"center\">#{input.name}</y:NodeLabel>"
           input_buf << "</y:ShapeNode>"
           input_buf << "</data>"
-          if @last_session_context[input.name]
-            input_buf << "<data key=\"d3\">#{_val(tensor)}</data>"
-          end
+
+          input_buf << "<data key=\"d3\">#{_val(tensor)}</data>" if @last_session_context[input.name]
+\
           input_buf << "</node>"
         elsif input.is_a?(Tensor)
           input_buf << "<node id=\"#{_gml_string(input.name)}\">"
@@ -175,12 +175,11 @@ module TensorStream
           input_buf << "<data key=\"d9\">"
           input_buf << "<y:ShapeNode>"
 
-          if input.internal?
-            input_buf << "  <y:Fill color=\"#C0C0C0\" transparent=\"false\"/>"
-          else
-            input_buf << "  <y:Fill color=\"#FFFFFF\" transparent=\"false\"/>"
-          end
-
+          input_buf << if input.internal?
+                         "  <y:Fill color=\"#C0C0C0\" transparent=\"false\"/>"
+                       else
+                         "  <y:Fill color=\"#FFFFFF\" transparent=\"false\"/>"
+                       end
 
           input_buf << "  <y:NodeLabel alignment=\"center\">#{input.name}</y:NodeLabel>"
 
@@ -189,7 +188,7 @@ module TensorStream
           input_buf << "</node>"
         end
 
-        if !add_to_group(groups, input.name, input_buf)
+        unless add_to_group(groups, input.name, input_buf)
           if input.is_a?(Variable)
             add_to_group(groups, "variable/#{input.name}", input_buf)
           else
@@ -205,7 +204,7 @@ module TensorStream
     end
 
     def _gml_string(str)
-      str.gsub('/','-')
+      str.tr('/', '-')
     end
 
     def output_edge(input, tensor, arr_buf, index = 0)
@@ -215,22 +214,20 @@ module TensorStream
 
       arr_buf << "<y:PolyLineEdge>"
       arr_buf << "<y:EdgeLabel >"
-      if !@last_session_context.empty?
-        arr_buf << "<![CDATA[  #{_val(input)}  ]]>"
-      else
-        if input.shape.shape.nil?
-          arr_buf << "<![CDATA[ #{input.data_type.to_s} ? ]]>"
-        else
-          arr_buf << "<![CDATA[ #{input.data_type.to_s} #{input.shape.shape.empty? ? 'scalar' : input.shape.shape.to_json}  ]]>"
-        end
-      end
+      arr_buf << if !@last_session_context.empty?
+                   "<![CDATA[  #{_val(input)}  ]]>"
+                 elsif input.shape.shape.nil?
+                   "<![CDATA[ #{input.data_type} ? ]]>"
+                 else
+                   "<![CDATA[ #{input.data_type} #{input.shape.shape.empty? ? 'scalar' : input.shape.shape.to_json}  ]]>"
+                 end
       arr_buf << "</y:EdgeLabel >"
       arr_buf << "<y:Arrows source=\"none\" target=\"standard\"/>"
-      if index == 0
-        arr_buf << "<y:LineStyle color=\"#FF0000\" type=\"line\" width=\"1.0\"/>"
-      else
-        arr_buf << "<y:LineStyle color=\"#0000FF\" type=\"line\" width=\"1.0\"/>"
-      end
+      arr_buf << if index.zero?
+                  "<y:LineStyle color=\"#FF0000\" type=\"line\" width=\"1.0\"/>"
+                 else
+                   "<y:LineStyle color=\"#0000FF\" type=\"line\" width=\"1.0\"/>"
+                 end
       arr_buf << "</y:PolyLineEdge>"
       arr_buf << "</data>"
       arr_buf << "</edge>"

data/lib/tensor_stream/graph_serializers/pbtext.rb

@@ -1,4 +1,5 @@
 module TensorStream
+  # Parses pbtext files and loads it as a graph
   class Pbtext < TensorStream::Serializer
     include TensorStream::StringHelper
     include TensorStream::OpHelper
@@ -47,11 +48,11 @@ module TensorStream
         @lines << "  attr {"
         @lines << "    key: \"#{k}\""
         @lines << "    value {"
-        if (v.is_a?(TrueClass) || v.is_a?(FalseClass))
-          @lines << "      b: #{v.to_s}"
-        elsif (v.is_a?(Integer))
+        if v.is_a?(TrueClass) || v.is_a?(FalseClass)
+          @lines << "      b: #{v}"
+        elsif v.is_a?(Integer)
           @lines << "      int_val: #{v}"
-        elsif (v.is_a?(Float))
+        elsif v.is_a?(Float)
           @lines << "      float_val: #{v}"
         end
         @lines << "    }"
@@ -60,21 +61,23 @@ module TensorStream
     end
 
     def pack_arr_float(float_arr)
-      float_arr.flatten.pack('f*').bytes.map { |b| b.chr =~ /[^[:print:]]/ ? "\\#{sprintf("%o", b).rjust(3, '0')}" : b.chr  }.join
+      float_arr.flatten.pack('f*').bytes.map { |b| b.chr =~ /[^[:print:]]/ ? "\\#{sprintf("%o", b).rjust(3, '0')}" : b.chr }.join
     end
 
     def pack_arr_int(int_arr)
-      int_arr.flatten.pack('l*').bytes.map { |b| b.chr =~ /[^[:print:]]/ ? "\\#{sprintf("%o", b).rjust(3, '0')}" : b.chr  }.join
+      int_arr.flatten.pack('l*').bytes.map { |b| b.chr =~ /[^[:print:]]/ ? "\\#{sprintf("%o", b).rjust(3, '0')}" : b.chr }.join
     end
 
     def shape_buf(tensor, shape_type = 'tensor_shape')
       arr = []
       arr << "  #{shape_type} {"
-      tensor.shape.shape.each do |dim|
-        arr << "    dim {"
-        arr << "      size: #{dim}"
-        arr << "    }"
-      end if tensor.shape.shape
+      if tensor.shape.shape
+        tensor.shape.shape.each do |dim|
+          arr << "    dim {"
+          arr << "      size: #{dim}"
+          arr << "    }"
+        end
+      end
       arr << "  }"
       arr
     end
@@ -102,14 +105,14 @@ module TensorStream
         end
       else
         val_type = if TensorStream::Ops::INTEGER_TYPES.include?(tensor.data_type)
-          "int_val"
-        elsif TensorStream::Ops::FLOATING_POINT_TYPES.include?(tensor.data_type)
-          "float_val"
-        elsif tensor.data_type == :string
-          "string_val"
-        else
-          "val"
-        end
+                     "int_val"
+                   elsif TensorStream::Ops::FLOATING_POINT_TYPES.include?(tensor.data_type)
+                     "float_val"
+                   elsif tensor.data_type == :string
+                     "string_val"
+                   else
+                     "val"
+                   end
         arr << "  #{val_type}: #{tensor.value.to_json}"
       end
       arr << "}"

data/lib/tensor_stream/helpers/string_helper.rb

@@ -12,11 +12,10 @@ module TensorStream
     end
 
     def underscore(string)
-      string.gsub(/::/, '/').
-      gsub(/([A-Z]+)([A-Z][a-z])/,'\1_\2').
-      gsub(/([a-z\d])([A-Z])/,'\1_\2').
-      tr("-", "_").
-      downcase
+      string.gsub(/::/, '/')
+            .gsub(/([A-Z]+)([A-Z][a-z])/, '\1_\2')
+            .gsub(/([a-z\d])([A-Z])/, '\1_\2')
+            .tr("-", "_").downcase
     end
 
     def symbolize_keys(hash)

data/lib/tensor_stream/math_gradients.rb

@@ -1,9 +1,10 @@
 module TensorStream
   # Class that provides auto-differentiation
+  # Most gradients are ported over from tensorflow's math_grad.py
   class MathGradients
     extend TensorStream::OpHelper
 
-    def self.tf
+    def self.ts
       TensorStream
     end
 
@@ -16,7 +17,7 @@ module TensorStream
         node.consumers.include?(tensor.name) || node.equal?(tensor)
       end.compact + [wrt_dx.name]
 
-      grad = i_op(:fill, tf.shape(tensor), tf.constant(1, dtype: wrt_dx.data_type))
+      grad = i_op(:fill, ts.shape(tensor), ts.constant(1, dtype: wrt_dx.data_type))
 
       _propagate(grad, tensor, wrt_dx, nodes_to_compute, options[:stop_gradients] || []) || i_op(:zeros_like, wrt_dx)
     end
@@ -41,6 +42,7 @@ module TensorStream
       end
     end
 
+    #TODO: refactor and implement registerGradient
     def self._compute_derivative(node, grad)
       node.graph.name_scope("#{node.name}_grad") do
         x = node.inputs[0] if node.inputs[0]
@@ -51,116 +53,161 @@ module TensorStream
           return [grad] * node.inputs.size
         when :add
           return [grad, grad] if shapes_fully_specified_and_equal(x, y)
-          sx = tf.shape(x, name: 'add/shape_x')
-          sy = tf.shape(y, name: 'add/shape_y')
+          sx = ts.shape(x, name: 'add/shape_x')
+          sy = ts.shape(y, name: 'add/shape_y')
           rx, ry = _broadcast_gradient_args(sx, sy)
 
-          [tf.reshape(tf.reduce_sum(grad, rx, name: 'add/reduce_sum_x'), sx),
-           tf.reshape(tf.reduce_sum(grad, ry, name: 'add/reduce_sum_y'), sy)]
+          [ts.reshape(ts.reduce_sum(grad, rx, name: 'add/reduce_sum_x'), sx),
+           ts.reshape(ts.reduce_sum(grad, ry, name: 'add/reduce_sum_y'), sy)]
         when :asin
-          tf.control_dependencies([grad]) do
-            x2 = tf.square(x)
-            one = tf.constant(1, dtype: grad.data_type)
-            den = tf.sqrt(tf.subtract(one, x2))
-            inv = tf.reciprocal(den)
+          ts.control_dependencies([grad]) do
+            x2 = ts.square(x)
+            one = ts.constant(1, dtype: grad.data_type)
+            den = ts.sqrt(ts.subtract(one, x2))
+            inv = ts.reciprocal(den)
             grad * inv
           end
         when :acos
-          tf.control_dependencies([grad]) do
-            x2 = tf.square(x)
-            one = tf.constant(1, dtype: grad.data_type)
-            den = tf.sqrt(tf.subtract(one, x2))
-            inv = tf.reciprocal(den)
+          ts.control_dependencies([grad]) do
+            x2 = ts.square(x)
+            one = ts.constant(1, dtype: grad.data_type)
+            den = ts.sqrt(ts.subtract(one, x2))
+            inv = ts.reciprocal(den)
             -grad * inv
           end
+        when :atan
+          ts.control_dependencies([grad]) do
+            x2 = ts.square(x)
+            one = ts.constant(1, dtype: grad.data_type)
+            inv = ts.reciprocal(ts.add(one, x2))
+            grad * inv
+          end
+        when :fill
+          [nil, ts.reduce_sum(grad)]
         when :sub
           return [grad, -grad] if shapes_fully_specified_and_equal(x, y)
 
-          sx = tf.shape(x, name: 'sub/shape_x')
-          sy = tf.shape(y, name: 'sub/shape_y')
+          sx = ts.shape(x, name: 'sub/shape_x')
+          sy = ts.shape(y, name: 'sub/shape_y')
           rx, ry = _broadcast_gradient_args(sx, sy)
 
-          [tf.reshape(tf.reduce_sum(grad, rx, name: 'add/reduce_sub_x'), sx),
-           -tf.reshape(tf.reduce_sum(grad, ry, name: 'add/reduce_sub_y'), sy)]
+          [ts.reshape(ts.reduce_sum(grad, rx, name: 'add/reduce_sub_x'), sx),
+           -ts.reshape(ts.reduce_sum(grad, ry, name: 'add/reduce_sub_y'), sy)]
         when :mul
-          sx = tf.shape(x)
-          sy = tf.shape(y)
+          sx = ts.shape(x)
+          sy = ts.shape(y)
           rx, ry = _broadcast_gradient_args(sx, sy)
 
-          [tf.reshape(tf.reduce_sum(tf.mul(grad, y), rx), sx),
-           tf.reshape(tf.reduce_sum(tf.mul(x, grad), ry), sy)]
+          [ts.reshape(ts.reduce_sum(ts.mul(grad, y), rx), sx),
+           ts.reshape(ts.reduce_sum(ts.mul(x, grad), ry), sy)]
         when :div
           sx = i_op(:shape, x)
           sy = i_op(:shape, y)
           rx, ry = _broadcast_gradient_args(sx, sy)
 
-          [tf.reshape(tf.reduce_sum(tf.div(grad, y), rx), sx),
-           tf.reshape(tf.reduce_sum(grad * tf.div(tf.div(-x, y), y), ry), sy)]
+          [ts.reshape(ts.reduce_sum(ts.div(grad, y), rx), sx),
+           ts.reshape(ts.reduce_sum(grad * ts.div(ts.div(-x, y), y), ry), sy)]
         when :mod
-          sx = tf.shape(x)
-          sy = tf.shape(y)
+          sx = ts.shape(x)
+          sy = ts.shape(y)
           rx, ry = _broadcast_gradient_args(sx, sy)
-          floor_xy = tf.floor_div(x, y)
-          gx = tf.reshape(tf.reduce_sum(grad, rx), sx)
-          gy = tf.reshape(tf.reduce_sum(grad * tf.negative(floor_xy), ry), sy)
+          floor_xy = ts.floor_div(x, y)
+          gx = ts.reshape(ts.reduce_sum(grad, rx), sx)
+          gy = ts.reshape(ts.reduce_sum(grad * ts.negative(floor_xy), ry), sy)
 
           [gx, gy]
+        when :prod
+          input_shape = ts.shape(x)
+          y = ts.range(0, ts.rank(x)) if y.nil?
+          reduction_indices = ts.reshape(y, [-1])
+
+          output_shape_kept_dims = ts.reduced_shape(input_shape, y)
+          tile_scaling = _safe_shape_div(input_shape, output_shape_kept_dims)
+          grad = ts.reshape(grad, output_shape_kept_dims)
+          grad = ts.tile(grad, tile_scaling)
+
+          perm, reduced_num, other_num = ts.device("/cpu:0") do
+            rank = ts.rank(x)
+            reduction_indices = (reduction_indices + rank) % rank
+            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))]
+          end
+
+          permuted = ts.transpose(x, perm)
+          permuted_shape = ts.shape(permuted)
+
+          reshaped = ts.reshape(permuted, [reduced_num, other_num])
+
+          # Calculate product, leaving out the current entry
+          left = ts.cumprod(reshaped, axis: 0, exclusive: true)
+          right = ts.cumprod(reshaped, axis: 0, exclusive: true, reverse: true)
+          y = ts.reshape(left * right, permuted_shape)
+
+          # Invert the transpose and reshape operations.
+          # Make sure to set the statically known shape information through a reshape.
+          out = grad * ts.transpose(y, ts.invert_permutation(perm))
+          [ts.reshape(out, input_shape, name: 'prod'), nil]
         when :squared_difference
           sx = i_op(:shape, x)
           sy = i_op(:shape, y)
           rx, ry = _broadcast_gradient_args(sx, sy)
 
-          x_grad = tf.mul(2.0, grad) * (x - y)
+          x_grad = ts.mul(2.0, grad) * (x - y)
 
-          [tf.reshape(tf.reduce_sum(x_grad, rx), sx),
-           tf.reshape(-tf.reduce_sum(x_grad, ry), sy)]
+          [ts.reshape(ts.reduce_sum(x_grad, rx), sx),
+           ts.reshape(-ts.reduce_sum(x_grad, ry), sy)]
         when :mat_mul
           t_a = node.options[:transpose_a]
           t_b = node.options[:transpose_b]
 
           if !t_a && !t_b
-            grad_a = tf.matmul(grad, y, transpose_b: true)
-            grad_b = tf.matmul(x, grad, transpose_a: true)
+            grad_a = ts.matmul(grad, y, transpose_b: true)
+            grad_b = ts.matmul(x, grad, transpose_a: true)
           elsif !ta && tb
-            grad_a = tf.matmul(grad, y)
-            grad_b = tf.matmul(grad, x, transpose_a: true)
+            grad_a = ts.matmul(grad, y)
+            grad_b = ts.matmul(grad, x, transpose_a: true)
           elsif t_a && !t_b
-            grad_a = tf.matmul(y, grad, transpose_b: true)
-            grad_b = tf.matmul(x, grad)
+            grad_a = ts.matmul(y, grad, transpose_b: true)
+            grad_b = ts.matmul(x, grad)
           elsif t_a && t_b
-            grad_a = tf.matmul(y, grad, transpose_a: true, transpose_b: true)
-            grad_b = tf.matmul(grad, x, transpose_a: true, transpose_b: true)
+            grad_a = ts.matmul(y, grad, transpose_a: true, transpose_b: true)
+            grad_b = ts.matmul(grad, x, transpose_a: true, transpose_b: true)
           end
 
           [grad_a, grad_b]
         when :sin
-          grad * tf.cos(x)
+          grad * ts.cos(x)
         when :tanh
           grad * i_op(:tanh_grad, x)
         when :pow
           z = node
-          sx = tf.shape(x)
-          sy = tf.shape(y)
+          sx = ts.shape(x)
+          sy = ts.shape(y)
           rx, ry = _broadcast_gradient_args(sx, sy)
-          gx = tf.reduce_sum(grad * y * tf.pow(x, y - 1), rx)
+          gx = ts.reduce_sum(grad * y * ts.pow(x, y - 1), rx)
 
-          log_x = tf.where(x > 0, tf.log(x), tf.zeros_like(x))
-          gy = tf.reduce_sum(grad * z * log_x, ry)
+          log_x = ts.where(x > 0, ts.log(x), ts.zeros_like(x))
+          gy = ts.reduce_sum(grad * z * log_x, ry)
 
           [gx, gy]
         when :abs
-          grad * tf.sign(x)
+          grad * ts.sign(x)
         when :log
-          grad * tf.reciprocal(x)
+          grad * ts.reciprocal(x)
         when :cos
-          -grad * tf.sin(x)
+          -grad * ts.sin(x)
         when :max
-          _min_or_max_grad(node.inputs, grad, ->(x, y) { tf.greater_equal(x, y) } )
+          _min_or_max_grad(node.inputs, grad, ->(a, b) { ts.greater_equal(a, b) })
         when :min
-          _min_or_max_grad(node.inputs, grad, ->(x, y) { tf.less_equal(x, y) } )
+          _min_or_max_grad(node.inputs, grad, ->(a, b) { ts.less_equal(a, b) })
         when :tan
-          secx = tf.reciprocal(tf.cos(x))
-          secx2 = tf.square(secx)
+          secx = ts.reciprocal(ts.cos(x))
+          secx2 = ts.square(secx)
           grad * secx2
         when :negate
           -grad
@@ -169,18 +216,25 @@ module TensorStream
         when :identity, :print
           grad
         when :sign
-          tf.zeros(tf.shape(x), dtype: x.data_type)
+          ts.zeros(ts.shape(x), dtype: x.data_type)
+        when :tile
+          input_shape = ts.shape(x)
+          split_shape = ts.reshape(ts.transpose(ts.stack([y, input_shape])), [-1])
+          axes = ts.range(0, ts.size(split_shape), 2)
+          input_grad = ts.reduce_sum(ts.reshape(grad, split_shape), axes)
+
+          [input_grad, nil]
         when :sum
           _sum_grad(x, y, grad)
         when :reciprocal
-          -grad * (tf.constant(1, dtype: x.dtype) / x**2)
+          -grad * (ts.constant(1, dtype: x.dtype) / x**2)
         when :sqrt
-          tf.constant(1, dtype: x.dtype) / (tf.constant(2, dtype: x.dtype) * tf.sqrt(x)) * grad
+          ts.constant(1, dtype: x.dtype) / (ts.constant(2, dtype: x.dtype) * ts.sqrt(x)) * grad
         when :stop_gradient
-          tf.zeros_like(grad)
+          ts.zeros_like(grad)
         when :square
-          y = tf.constant(2.0, dtype: x.dtype)
-          tf.multiply(grad, tf.multiply(x, y))
+          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])
@@ -191,12 +245,12 @@ module TensorStream
           [x_cond * grad, y_cond * grad]
         when :mean
           sum_grad = _sum_grad(x, y, grad)[0]
-          input_shape = tf.shape(x)
-          output_shape = tf.shape(node)
-          factor = _safe_shape_div(tf.reduce_prod(input_shape), tf.reduce_prod(output_shape))
-          tf.div(sum_grad, tf.cast(factor, sum_grad.data_type))
+          input_shape = ts.shape(x)
+          output_shape = ts.shape(node)
+          factor = _safe_shape_div(ts.reduce_prod(input_shape), ts.reduce_prod(output_shape))
+          [ts.div(sum_grad, ts.cast(factor, sum_grad.data_type)), nil]
         when :log1p
-          grad * tf.reciprocal(i_cons(1, dtype: grad.data_type) + x)
+          grad * ts.reciprocal(i_cons(1, dtype: grad.data_type) + x)
         when :sigmoid
           i_op(:sigmoid_grad, x, grad)
         when :sigmoid_grad
@@ -205,7 +259,8 @@ module TensorStream
         when :softmax
           i_op(:softmax_grad, x, grad)
         when :softmax_cross_entropy_with_logits_v2
-          [i_op(:softmax_cross_entropy_with_logits_v2_grad, x, y, grad), nil]
+          output = node
+          [_broadcast_mul(grad, output[1]), nil]
         when :floor, :ceil
           # non differentiable
           nil
@@ -215,6 +270,10 @@ module TensorStream
         when :argmin, :argmax, :floor_div
           # non differentiable
           [nil, nil]
+        when :transpose
+          return [ts.transpose(grad, ts.invert_permutation(y)), nil]
+        when :index
+          grad
         else
           raise "no derivative op for #{node.operation}"
         end
@@ -231,12 +290,12 @@ module TensorStream
     end
 
     def self._safe_shape_div(arg_x, arg_y)
-      _op(:floor_div, arg_x, tf.maximum(arg_y, 1))
+      _op(:floor_div, arg_x, ts.maximum(arg_y, 1))
     end
 
     def self._sum_grad(arg_x, arg_y, grad)
       input_shape = _op(:shape, arg_x)
-      output_shape_kept_dims = tf.reduced_shape(input_shape, arg_y)
+      output_shape_kept_dims = ts.reduced_shape(input_shape, arg_y)
       tile_scaling = _safe_shape_div(input_shape, output_shape_kept_dims)
       new_grad = _op(:reshape, grad, output_shape_kept_dims)
 
@@ -254,19 +313,24 @@ module TensorStream
       x = inputs[0]
       y = inputs[1]
       gdtype = grad.data_type
-      sx = tf.shape(x)
-      sy = tf.shape(y)
-      gradshape = tf.shape(grad)
-      zeros = tf.zeros(gradshape, dtype: gdtype)
+      sx = ts.shape(x)
+      sy = ts.shape(y)
+      gradshape = ts.shape(grad)
+      zeros = ts.zeros(gradshape, dtype: gdtype)
       xmask = selector_op.call(x, y)
       rx, ry = _broadcast_gradient_args(sx, sy)
-      xgrad = tf.where(xmask, grad, zeros, name: 'x')
-      ygrad = tf.where(xmask, zeros, grad, name: 'y')
-      gx = tf.reshape(tf.reduce_sum(xgrad, rx), sx)
-      gy = tf.reshape(tf.reduce_sum(ygrad, ry), sy)
+      xgrad = ts.where(xmask, grad, zeros, name: 'x')
+      ygrad = ts.where(xmask, zeros, grad, name: 'y')
+      gx = ts.reshape(ts.reduce_sum(xgrad, rx), sx)
+      gy = ts.reshape(ts.reduce_sum(ygrad, ry), sy)
       [gx, gy]
     end
 
+    def self._broadcast_mul(vec, mat)
+      vec = ts.expand_dims(vec, -1)
+      vec * mat
+    end
+
     def self._include?(arr, obj)
       arr.each { |a| return true if a.equal?(obj) }
       false