tensor_stream 0.8.5 → 0.8.6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 1d6e9e482de719b709bfe554085718cb01a3bbfd089983eef51ace418b1b7d2d
-  data.tar.gz: 287e63d7a7269e7143ef1016677297c44235a4ac6afb03234f72bb2b6774d348
+  metadata.gz: '0168aeff699a7490174f164370f6cef9dbc3413f23352a696c3a468832bd1b24'
+  data.tar.gz: f62926f3b919cbf5001872d6fab23095807d6c5028de5d885256449e1dd52a01
 SHA512:
-  metadata.gz: 3ec2af3376d4cc671eb7bfc549290145d3e14167fe3b472f1c4dd1a05a6ab78ddeddc608d6c1c26a93ddc542ce605341172212e7d5d2f5beeb2542ef3790b03f
-  data.tar.gz: 961e74c0acac0179affca04974f7933fb1162a1a88b3bcdde6c57f679db6d5e23c2e92dd4fff8454bda7ba87dc2eaad7f9424458440348151ea4ac90258a8acc
+  metadata.gz: ce886b38d14603f9c3f5dbfc6d96f89041350cd7cfa40a2cb2178e2845a43e22e98bca8bd689b72a2b0acceb5da8b3988d88632d601c2da701b0979375e84200
+  data.tar.gz: 6dac48e933a46144fb78433e17077b2c9d9fbe68e4297a0958a752276e51401279460c6ac47adc88ac0b3890a0bb535bfb0f889b64b2439c2ea38c68fb9d1665

data/CHANGELOG.md

@@ -4,6 +4,17 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.8.6] - 2018-09-11
+
+### Added
+- [TRAINING] Added RMSPropOptimizer, AdagradOptimizer
+- [NEW OP] shape_n, sparse_softmax_cross_entropy_with_logits, split, unstack
+- Added RNN sample
+  
+### Fixes
+- Fixed gradient computation when passing an array of tensors to a function
+- Added gradients for various other ops
+
 ## [0.8.5] - 2018-09-06
 
 ### Added

data/README.md

@@ -11,10 +11,9 @@ The goal of this gem is to have a high performance machine learning and compute
 ## Features
 
 - Replicates most of the commonly used low-level tensorflow ops (tf.add, tf.constant, tf.placeholder, tf.matmul, tf.sin etc...)
-- Supports auto-differentiation via tf.gradients (mostly)
+- Supports auto-differentiation
 - Provision to use your own opcode evaluator (opencl, sciruby and tensorflow backends planned)
 - Goal is to be as close to TensorFlow in behavior but with some freedom to add ruby specific enhancements (with lots of test cases)
-- eager execution (experimental)
 - (08-08-2018) Load pbtext files from tensorflow (Graph.parse_from_string)
 
 ## Compatibility
@@ -31,7 +30,7 @@ gem 'tensor_stream-opencl'
 and then (without bundler)
 
 ```ruby
-require 'tensor_stream-opencl'
+require 'tensor_stream/opencl'
 ```
 
 OpenCL is basically a requirement for deep learning and image processing tasks as the ruby implementation is too slow even with jit speedups using latest ruby implementations.
@@ -91,8 +90,11 @@ pred = X * W + b
 # Mean squared error
 cost = ((pred - Y) ** 2).reduce(:+) / ( 2 * n_samples)
 
-# optimizer =  TensorStream::Train::MomentumOptimizer.new(0.01, 0.5, use_nesterov: true).minimize(cost)
-# optimizer =  TensorStream::Train::AdamOptimizer.new.minimize(cost)
+# optimizer = TensorStream::Train::MomentumOptimizer.new(learning_rate, momentum, use_nesterov: true).minimize(cost)
+# optimizer = TensorStream::Train::AdamOptimizer.new(learning_rate).minimize(cost)
+# optimizer = TensorStream::Train::AdadeltaOptimizer.new(1.0).minimize(cost)
+# optimizer = TensorStream::Train::AdagradOptimizer.new(0.01).minimize(cost)
+# optimizer = TensorStream::Train::RMSPropOptimizer.new(0.01, centered: true).minimize(cost)
 optimizer = TensorStream::Train::GradientDescentOptimizer.new(learning_rate).minimize(cost)
 
 # Initialize the variables (i.e. assign their default value)
@@ -212,7 +214,7 @@ gem 'tensor_stream-opencl'
 To use the opencl evaluator instead of the ruby evaluator simply require it (if using rails this should be loaded automatically).
 
 ```ruby
-require 'tensor_stream-opencl'
+require 'tensor_stream/opencl'
 ```
 
 Adding the OpenCL evaluator should expose additional devices available to tensor_stream
@@ -227,7 +229,7 @@ By default TensorStream will determine using the given evaluators the best possi
 placement for each tensor operation
 
 ```ruby
-require 'tensor_stream/evaluator/opencl/opencl_evaluator'
+require 'tensor_stream/opencl'
 
 # set session to use the opencl evaluator
 sess = ts.session

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

@@ -1,12 +1,27 @@
 module TensorStream
   # varoius utility functions for array processing
   module ArrayOpsHelper
+    def split_tensor(input, begin_index, end_index, axis = 0)
+      if axis.zero?
+        input[begin_index...end_index]
+      else
+        input.collect do |item|
+          split_tensor(item, begin_index, end_index, axis - 1)
+        end
+      end
+    end
+
     def slice_tensor(input, start, size)
       return input if size.empty?
       start_index = start.shift
-      dimen_size = start_index + size.shift
+      current_size = size.shift
+      dimen_size = if current_size == -1
+        input.size - 1
+      else
+        start_index + current_size - 1
+      end
 
-      input[start_index...dimen_size].collect do |item|
+      input[start_index..dimen_size].collect do |item|
         if item.is_a?(Array)
           slice_tensor(item, start.dup, size.dup)
         else

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

@@ -7,7 +7,8 @@ module TensorStream
           start = inputs[1]
           size = complete_eval(tensor.options[:size], context)
           raise "start index and size not of the same shape #{start.size} != #{size.size}" if start.size != size.size
-          slice_tensor(input, start, size)
+
+          slice_tensor(input, start.dup, size.dup)
         end
 
         register_op %i[flow_dynamic_stitch dynamic_stitch] do |_context, _tensor, inputs|
@@ -22,8 +23,9 @@ module TensorStream
           gather(params, indexes)
         end
 
-        register_op %i[concat concat_v2] do |_context, tensor, inputs|
-          concat_array(inputs, tensor.options[:axis])
+        register_op %i[concat concat_v2] do |_context, _tensor, inputs|
+          axis = inputs.shift
+          concat_array(inputs, axis)
         end
 
         register_op :stack do |_context, tensor, inputs|
@@ -74,6 +76,55 @@ module TensorStream
           TensorShape.reshape(output_buffer, new_shape)
         end
 
+        register_op :unstack do |_context, tensor, inputs|
+          value = inputs[0]
+
+          axis = tensor.options[:axis] || 0
+          new_shape = shape_eval(inputs[0])
+          rank = new_shape.size - 1
+
+          divisors = new_shape.dup.drop(1).reverse.inject([1]) do |a, s|
+            a << s * a.last
+          end.reverse
+
+          axis = rank + axis if axis < 0
+          rotated_shape = Array.new(axis + 1) { new_shape.shift }
+          new_shape = rotated_shape.rotate!(-1) + new_shape
+          output_buffer = Array.new(new_shape.reduce(:*)) { 0 }
+
+          multipliers = new_shape.dup.drop(1).reverse.inject([1]) do |a, s|
+            a << s * a.last
+          end.reverse
+
+          inputs.each_with_index do |input, index|
+            raw_input = input.is_a?(Array) ? input.flatten : [input]
+            start = index * divisors.first
+
+            raw_input.each_with_index do |x, index2|
+              index_map = []
+              ptr = start + index2
+              divisors.each_with_object(index_map) do |div, a|
+                a << (ptr / div.to_f).floor
+                ptr = ptr % div
+              end
+
+              rotated_index = Array.new(axis + 1) { index_map.shift }
+              index_map = rotated_index.rotate!(-1) + index_map
+
+              ptr2 = 0
+              multipliers.each_with_index do |m, idx|
+                ptr2 += index_map[idx] * m
+              end
+
+              output_buffer[ptr2] = x
+            end
+          end
+
+          res = TensorShape.reshape(output_buffer, new_shape)
+
+          TensorStream::Evaluator::OutputGroup.new(res)
+        end
+
         register_op :squeeze do |_context, tensor, inputs|
           val = inputs[0]
           shape = shape_eval(val)
@@ -81,10 +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 
+                shape[axis] = nil
               else
                 raise TensorStream::ValueError, "unable to squeeze dimension that does not have a size of 1"
               end
@@ -93,7 +143,7 @@ module TensorStream
             shape = shape.map { |s| s == 1 ? nil : s }
           end
 
-          TensorShape.reshape(val.flatten, shape.compact)
+          TensorShape.reshape(val, shape.compact)
         end
 
         register_op :expand_dims do |_context, _tensor, inputs|
@@ -105,7 +155,7 @@ module TensorStream
 
           new_shape = shape.dup.insert(axis, 1).compact
 
-          TensorShape.reshape([val].flatten, new_shape)
+          TensorShape.reshape([val], new_shape)
         end
 
         register_op :fill do |_context, _tensor, inputs|
@@ -207,7 +257,6 @@ module TensorStream
                  else
                    -> { int_type?(tensor.data_type) ? 1 : 1.0 }
                  end
-
           if shape.is_a?(Array) && shape.size.zero?
             func.call
           else
@@ -232,16 +281,38 @@ module TensorStream
           get_rank(inputs[0])
         end
 
+        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
+          TensorStream::Evaluator::OutputGroup.new(res)
+        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]
           else
-            new_shape = TensorShape.fix_inferred_elements(new_shape, flat_arr.size)
             TensorShape.reshape(flat_arr, new_shape)
           end
         end
@@ -276,6 +347,17 @@ module TensorStream
           pred = complete_eval(tensor.options[:pred], context)
           call_3way_vector_op(pred, inputs[0], inputs[1], context, ->(t, u, v) { t ? u : v })
         end
+
+        register_op :shape do |_context, tensor, inputs|
+          shape_eval(inputs[0], tensor.options[:out_type])
+        end
+
+        register_op :shape_n do |_context, _tensor, inputs|
+          shapes = inputs.collect do |input|
+            shape_eval(input)
+          end
+          TensorStream::Evaluator::OutputGroup.new(shapes)
+        end
       end
     end
   end

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

@@ -0,0 +1,9 @@
+module TensorStream
+  module CheckOps
+    def CheckOps.included(klass)
+      klass.class_eval do
+
+      end
+    end
+  end
+end

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

@@ -34,7 +34,7 @@ module TensorStream
 
             color_values
           end
-          TensorShape.reshape(image_data.flatten, [image.height, image.width, channels])
+          TensorShape.reshape(image_data, [image.height, image.width, channels])
         end
 
         register_op :encode_png do |_context, tensor, inputs|

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

@@ -2,24 +2,24 @@ module TensorStream
   module MathOps
     def MathOps.included(klass)
       klass.class_eval do
-        register_op :tanh, no_eval: true do |context, _tensor, inputs|
-          call_op(:tanh, inputs[0], context, ->(t, _b) { Math.tanh(t) })
+        register_op :tanh, no_eval: true do |context, tensor, inputs|
+          call_op(tensor, inputs[0], context, ->(t, _b) { Math.tanh(t) })
         end
 
-        register_op :tan, no_eval: true do |context, _tensor, inputs|
-          call_op(:tan, inputs[0], context, ->(t, _b) { Math.tan(t) })
+        register_op :tan, no_eval: true do |context, tensor, inputs|
+          call_op(tensor, inputs[0], context, ->(t, _b) { Math.tan(t) })
         end
 
-        register_op :atan, no_eval: true do |context, _tensor, inputs|
-          call_op(:atan, inputs[0], context, ->(t, _b) { Math.atan(t) })
+        register_op :atan, no_eval: true do |context, tensor, inputs|
+          call_op(tensor, inputs[0], context, ->(t, _b) { Math.atan(t) })
         end
 
-        register_op :sec, no_eval: true do |context, _tensor, inputs|
-          call_op(:sec, inputs[0], context, ->(t, _b) { Math.sec(t) })
+        register_op :sec, no_eval: true do |context, tensor, inputs|
+          call_op(tensor, inputs[0], context, ->(t, _b) { Math.sec(t) })
         end
 
-        register_op :sin, no_eval: true do |context, _tensor, inputs|
-          call_op(:sin, inputs[0], context, ->(t, _b) { Math.sin(t) })
+        register_op :sin, no_eval: true do |context, tensor, inputs|
+          call_op(tensor, inputs[0], context, ->(t, _b) { Math.sin(t) })
         end
 
         register_op :add, no_eval: true do |context, tensor, inputs|
@@ -79,64 +79,64 @@ module TensorStream
           call_op(:round, inputs[0], context, ->(t, _b) { t.round })
         end
 
-        register_op :abs, no_eval: true do |context, _tensor, inputs|
-          call_op(:abs, inputs[0], context, ->(t, _b) { t.abs })
+        register_op :abs, no_eval: true do |context, tensor, inputs|
+          call_op(tensor, inputs[0], context, ->(t, _b) { t.abs })
         end
 
-        register_op :asin, no_eval: true do |context, _tensor, inputs|
-          call_op(:asin, inputs[0], context, ->(t, _b) { Math.asin(t) })
+        register_op :asin, no_eval: true do |context, tensor, inputs|
+          call_op(tensor, inputs[0], context, ->(t, _b) { Math.asin(t) })
         end
 
-        register_op :acos, no_eval: true do |context, _tensor, inputs|
-          call_op(:acos, inputs[0], context, ->(t, _b) { Math.acos(t) })
+        register_op :acos, no_eval: true do |context, tensor, inputs|
+          call_op(tensor, inputs[0], context, ->(t, _b) { Math.acos(t) })
         end
 
-        register_op :cos, no_eval: true do |context, _tensor, inputs|
-          call_op(:cos, inputs[0], context, ->(t, _b) { Math.cos(t) })
+        register_op :cos, no_eval: true do |context, tensor, inputs|
+          call_op(tensor, inputs[0], context, ->(t, _b) { Math.cos(t) })
         end
 
-        register_op :log1p, no_eval: true do |context, _tensor, inputs|
-          call_op(:log1p, inputs[0], context, ->(t, _b) { Math.log(1 + t) })
+        register_op :log1p, no_eval: true do |context, tensor, inputs|
+          call_op(tensor, inputs[0], context, ->(t, _b) { Math.log(1 + t) })
         end
 
-        register_op :log, no_eval: true do |context, _tensor, inputs|
-          call_op(:log, inputs[0], context, ->(t, _b) { t < 0 ? Float::NAN : Math.log(t) })
+        register_op :log, no_eval: true do |context, tensor, inputs|
+          call_op(tensor, inputs[0], context, ->(t, _b) { t < 0 ? Float::NAN : Math.log(t) })
         end
 
-        register_op :exp, no_eval: true do |context, _tensor, inputs|
-          call_op(:exp, inputs[0], context, ->(t, _b) { Math.exp(t) })
+        register_op :exp, no_eval: true do |context, tensor, inputs|
+          call_op(tensor, inputs[0], context, ->(t, _b) { Math.exp(t) })
         end
 
-        register_op :sigmoid, no_eval: true do |context, _tensor, inputs|
-          call_op(:sigmoid, inputs[0], context, ->(t, _b) { sigmoid(t) })
+        register_op :sigmoid, no_eval: true do |context, tensor, inputs|
+          call_op(tensor, inputs[0], context, ->(t, _b) { sigmoid(t) })
         end
 
-        register_op :sqrt, no_eval: true do |context, _tensor, inputs|
-          call_op(:sqrt, inputs[0], context, ->(t, _b) { Math.sqrt(t) })
+        register_op :sqrt, no_eval: true do |context, tensor, inputs|
+          call_op(tensor, inputs[0], context, ->(t, _b) { Math.sqrt(t) })
         end
 
-        register_op :floor, no_eval: true do |context, _tensor, inputs|
-          call_op(:floor, inputs[0], context, ->(t, _b) { t.floor })
+        register_op :floor, no_eval: true do |context, tensor, inputs|
+          call_op(tensor, inputs[0], context, ->(t, _b) { t.floor })
         end
 
-        register_op :ceil, no_eval: true do |context, _tensor, inputs|
-          call_op(:ceil, inputs[0], context, ->(t, _b) { t.ceil })
+        register_op :ceil, no_eval: true do |context, tensor, inputs|
+          call_op(tensor, inputs[0], context, ->(t, _b) { t.ceil })
         end
 
-        register_op :square, no_eval: true do |context, _tensor, inputs|
-          call_op(:square, inputs[0], context, ->(t, _b) { t * t })
+        register_op :square, no_eval: true do |context, tensor, inputs|
+          call_op(tensor, inputs[0], context, ->(t, _b) { t * t })
         end
 
-        register_op :reciprocal, no_eval: true do |context, _tensor, inputs|
-          call_op(:reciprocal, inputs[0], context, ->(t, _b) { 1 / t })
+        register_op :reciprocal, no_eval: true do |context, tensor, inputs|
+          call_op(tensor, inputs[0], context, ->(t, _b) { 1 / t })
         end
 
         register_op %i[neg negate], no_eval: true do |context, tensor, inputs|
           call_vector_op(tensor, :negate, inputs[0], nil, context, ->(t, _u) { -t })
         end
 
-        register_op :tanh_grad, no_eval: true do |context, _tensor, inputs|
-          call_op(:tanh_grad, inputs[0], context, ->(t, _b) { 1 - Math.tanh(t) * Math.tanh(t) })
+        register_op :tanh_grad, no_eval: true do |context, tensor, inputs|
+          call_op(tensor, inputs[0], context, ->(t, _b) { 1 - Math.tanh(t) * Math.tanh(t) })
         end
 
         register_op(%i[argmax arg_max]) do |_context, tensor, inputs|

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

@@ -15,46 +15,78 @@ module TensorStream
           target_var, momentum_var, learning_rate, grad, momentum = inputs
           assign = tensor.inputs[0] || tensor
           assign_acc = tensor.inputs[1]
-          assign_acc.value = multi_array_op(->(t, u) { t * momentum + u }, momentum_var, grad )
+          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)
+            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
         end
 
-        register_op :apply_adadelta do |context, tensor, inputs|
+        register_op :apply_adadelta do |_context, tensor, inputs|
           target_var, accum, accum_update, lr, rho, epsilon, grad = inputs
           assign = tensor.inputs[0] || tensor
           assign_acc = tensor.inputs[1]
           assign_acc_update = tensor.inputs[2]
-          assign_acc.value = multi_array_op(->(acc_t, grad_t) { acc_t * rho + (grad_t * grad_t) * (1.0 - rho ) }, accum, grad)
+          assign_acc.value = multi_array_op(->(acc_t, grad_t) { acc_t * rho + (grad_t * grad_t) * (1.0 - rho) }, accum, grad)
           update = multi_array_op(->(acc_update_t, acc_t, grad_t) { Math.sqrt(acc_update_t + epsilon) * (1.0 / Math.sqrt(acc_t + epsilon)) * grad_t }, accum_update, assign_acc.value, grad)
           assign.value = multi_array_op(->(v, u) { v - (u * lr) }, target_var, update)
-          assign_acc_update.value = multi_array_op(->(acc_update_t, u) {  acc_update_t * rho + (u * u) * (1.0 - rho) }, accum_update, update)
- 
+          assign_acc_update.value = multi_array_op(->(acc_update_t, u) { acc_update_t * rho + (u * u) * (1.0 - rho) }, accum_update, update)
+
+          assign.value
+        end
+
+        register_op :apply_adagrad do |_context, tensor, inputs|
+          target_var, accum, lr, grad = inputs
+          assign = tensor.inputs[0] || tensor
+          assign.value = multi_array_op(->(v, a, g) { v - (g * lr * (1.0 / Math.sqrt(a))) }, target_var, accum, grad)
           assign.value
         end
 
-        register_op :apply_adam do |context, tensor, inputs|
+        register_op :apply_adam do |_context, tensor, inputs|
           target_var, m, v, beta1_power, beta2_power, lr_t, beta1_t, beta2_t, epsilon_t, grad = inputs
-          alpha = lr_t * Math.sqrt( 1.0 - beta2_power) / (1.0 - beta1_power)
+          alpha = lr_t * Math.sqrt(1.0 - beta2_power) / (1.0 - beta1_power)
           assign = tensor.inputs[0]
           assign_m = tensor.inputs[1]
           assign_v = tensor.inputs[2]
 
           assign_m.value = multi_array_op(->(u_d , g) { u_d + (g - u_d) * (1.0 - beta1_t) }, m, grad)
-          assign_v.value = multi_array_op(->(u_d , v_d) { u_d + (v_d ** 2 - u_d) * (1.0 - beta2_t)},  v, grad)
-          assign.value = multi_array_op( ->(t, m_d , v_d) { t - ((m_d * alpha) / (Math.sqrt(v_d) + epsilon_t)) }, target_var, assign_m.value, assign_v.value)
+          assign_v.value = multi_array_op(->(u_d , v_d) { u_d + (v_d**2 - u_d) * (1.0 - beta2_t)},  v, grad)
+          assign.value = multi_array_op(->(t, m_d , v_d) { t - ((m_d * alpha) / (Math.sqrt(v_d) + epsilon_t)) }, target_var, assign_m.value, assign_v.value)
           assign.value
         end
 
+        register_op :apply_rms_prop do |_context, tensor, inputs|
+          var, ms, mom, lr, rho, momentum, epsilon, grad = inputs
+          assign = tensor.inputs[0]
+          assign_ms = tensor.inputs[1]
+          assign_mom = tensor.inputs[2]
+          assign_ms.value = multi_array_op(->(g, m) { m + (g * g - m) * (1.0 - rho)}, grad, ms)
+          assign_mom.value = multi_array_op(->(mom_t, g, m) { mom_t * momentum + (g * lr) / Math.sqrt(m + epsilon)}, mom, grad, assign_ms.value)
+          assign.value = multi_array_op(->(v, m) { v - m }, var, assign_mom.value)
+        end
+
+        register_op :apply_centered_rms_prop do |_context, tensor, inputs|
+          var, mg, ms, mom, lr, rho, momentum, epsilon, grad = inputs
+          assign = tensor.inputs[0]
+          assign_mg = tensor.inputs[1]
+          assign_ms = tensor.inputs[2]
+          assign_mom = tensor.inputs[3]
+
+          assign_ms.value = multi_array_op(->(g, m) { m + (g * g - m) * (1.0 - rho) }, grad, ms)
+          assign_mg.value = multi_array_op(->(g, mg_t) {  (g - mg_t) * (1.0 - rho) }, grad, mg)
+          denom =  multi_array_op(->(s, mg_t) { (s - mg_t * mg_t) + epsilon }, assign_ms.value, mg)
+          assign_mom.value = multi_array_op(->(mom_t, g, d) { mom_t * momentum + (g * lr) / Math.sqrt(d)}, mom, grad, denom)
+          assign.value = multi_array_op(->(v, m) { v - m }, var, assign_mom.value)
+        end
+
         register_op %i[softmax_cross_entropy_with_logits_v2 softmax_cross_entropy_with_logits] do |_context, tensor, inputs|
           last_dimen_list = last_axis(inputs[0])
           input_shape = shape_eval(inputs[0])
           rank = input_shape.size - 1
           labels = last_axis(inputs[1])
+
           func = lambda { |logits, label|
             c = logits.max
             transformed_logits = logits.map { |l| l - c }
@@ -83,6 +115,49 @@ module TensorStream
           end
         end
 
+        register_op :sparse_softmax_cross_entropy_with_logits do |context, tensor, inputs|
+          last_dimen_list = last_axis(inputs[0])
+          input_shape = shape_eval(inputs[0])
+          rank = input_shape.size - 1
+          labels = last_axis(inputs[1])
+          num_classes = input_shape.last
+
+          labels = labels.map do |l|
+            one_hot = Array.new(num_classes) { 0 }
+            one_hot[l] = 1
+            one_hot
+          end
+
+          func = lambda { |logits, label|
+            c = logits.max
+            transformed_logits = logits.map { |l| l - c }
+            sum = transformed_logits.map { |x| Math.exp(x) }.reduce(:+)
+            losses = transformed_logits.zip(label).map { |x, y| (Math.log(sum) - x) * y }
+            probs = transformed_logits.zip(label).map  { |x, y| (Math.exp(x) / sum) - y }
+            [losses, probs]
+          }
+
+          if input_shape.size == 1
+            loss, prob = func.call(last_dimen_list, labels)
+            loss = reduce(loss, rank, false)
+
+            TensorStream::Evaluator::OutputGroup.new([loss, prob], [tensor.inputs[0].data_type, tensor.inputs[0].data_type])
+          else
+            losses = []
+            backprobs = []
+            arr = last_dimen_list.zip(labels).each do |list, label|
+              loss, prob = func.call(list, label)
+              losses << loss
+              backprobs << prob
+            end
+            reshaped_losses = TensorShape.reshape(losses, input_shape)
+            reshaped_backprops = TensorShape.reshape(backprobs, input_shape)
+            reshaped_losses = reduce(reshaped_losses, rank, false)
+
+            TensorStream::Evaluator::OutputGroup.new([reshaped_losses, reshaped_backprops], [tensor.inputs[0].data_type, tensor.inputs[0].data_type])
+          end
+        end
+
         register_op :log_softmax do |_context, _tensor, inputs|
           input_shape = shape_eval(inputs[0])
           last_dimen_list = last_axis(inputs[0])
@@ -100,7 +175,7 @@ module TensorStream
             arr = last_dimen_list.collect do |list|
               func.call(list)
             end
-            TensorShape.reshape(arr.flatten, input_shape)
+            TensorShape.reshape(arr, input_shape)
           end
         end
 
@@ -129,7 +204,7 @@ module TensorStream
             arr = last_dimen_list.zip(last_grad_list).collect do |list, last_grad|
               func.call(list, last_grad)
             end
-            TensorShape.reshape(arr.flatten, input_shape)
+            TensorShape.reshape(arr, input_shape)
           end
         end
       end