tensor_stream 1.0.0 → 1.0.1

data/lib/tensor_stream/ops/add.rb

@@ -0,0 +1,23 @@
+TensorStream::OpMaker.define_operation :add do |op|
+  op.what_it_does "Returns x + y element-wise."
+
+  op.parameter :input_a, "tensor X"
+  op.parameter :input_b, "tensor Y"
+
+  op.apply_data_type_coercion!
+  op.supports_broadcasting!
+
+  op.option :name, "Optional name", :nil
+
+  op.define_gradient do |grad, node, params|
+    x, y = params
+    next [grad, grad] if shapes_fully_specified_and_equal(x, y)
+
+    sx = ts.shape(x, name: "add/shape_x")
+    sy = ts.shape(y, name: "add/shape_y")
+    rx, ry = _broadcast_gradient_args(sx, 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),]
+  end
+end

data/lib/tensor_stream/ops/argmax.rb

@@ -0,0 +1,14 @@
+TensorStream::OpMaker.define_operation :argmax do |op|
+  op.what_it_does "Returns the index with the largest value across axes of a tensor."
+
+  op.parameter :input_a, "tensor X", validate: 'NUMERIC_TYPES'
+  op.parameter :axis, "Describes which axis of the input tensor to reduce across. For vectors, use axis = 0", :nil, validate: 'INTEGER_TYPES'
+
+  op.option :name, "Optional name", :nil
+  op.option :dimension, "Same as axis", :nil
+  op.option :output_type, "Output data type defaults to int32", ":int32"
+
+  op.define_gradient do |grad, node, params|
+    [nil, nil]
+  end
+end

data/lib/tensor_stream/ops/argmin.rb

@@ -0,0 +1,14 @@
+TensorStream::OpMaker.define_operation :argmin do |op|
+  op.what_it_does "Returns the index with the smallest value across axes of a tensor."
+
+  op.parameter :input_a, "tensor X", validate: 'NUMERIC_TYPES'
+  op.parameter :axis, "Describes which axis of the input tensor to reduce across. For vectors, use axis = 0", :nil, validate: 'INTEGER_TYPES'
+
+  op.option :name, "Optional name", :nil
+  op.option :dimension, "Same as axis", :nil
+  op.option :output_type, "Output data type defaults to int32", ":int32"
+
+  op.define_gradient do |grad, node, params|
+    [nil, nil]
+  end
+end

data/lib/tensor_stream/ops/case.rb

@@ -0,0 +1,17 @@
+TensorStream::OpMaker.define_operation :case do |op|
+  op.exclude!
+
+  op.define_gradient do |grad, node, params|
+    n_preds = node.inputs.size - 2
+
+    case_grads = Array.new(n_preds) { |index|
+      i_op(:case_grad, index, node.inputs[0], node.inputs[2 + index], grad)
+    }
+
+    [nil, i_op(:case_grad, -1, node.inputs[0], node.inputs[1], grad)] + case_grads
+  end
+
+  op.define_shape do |tensor|
+    tensor.inputs[2]&.shape&.shape
+  end
+end

data/lib/tensor_stream/ops/cast.rb

@@ -0,0 +1,15 @@
+TensorStream::OpMaker.define_operation :cast do |op|
+  op.exclude!
+
+  op.define_gradient do |grad, node, params|
+    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)
+      next ts.cast(grad, src_type)
+    end
+
+    nil
+  end
+end

data/lib/tensor_stream/ops/ceil.rb

@@ -0,0 +1,15 @@
+TensorStream::OpMaker.define_operation :ceil do |op|
+  op.what_it_does "Returns element-wise smallest integer in not less than x"
+
+  op.parameter :input_a, "tensor X", validate: 'FLOATING_POINT_TYPES'
+
+  op.option :name, "Optional name", :nil
+
+  op.define_gradient do |grad, node, params|
+    nil
+  end
+
+  op.define_shape do |tensor|
+    tensor.inputs[0].shape.shape
+  end
+end

data/lib/tensor_stream/ops/cos.rb

@@ -0,0 +1,10 @@
+TensorStream::OpMaker.define_operation :cos do |op|
+  op.what_it_does "Computes cos of input element-wise."
+
+  op.parameter :input_a, "tensor X", validate: 'FLOATING_POINT_TYPES'
+  op.option :name, "Optional name", :nil
+
+  op.define_gradient do |grad, node, params|
+    -grad * ts.sin(params[0])
+  end
+end

data/lib/tensor_stream/ops/div.rb

@@ -0,0 +1,21 @@
+TensorStream::OpMaker.define_operation :div do |op|
+  op.what_it_does "Returns x / y element-wise."
+
+  op.parameter :input_a, "tensor X"
+  op.parameter :input_b, "tensor Y"
+
+  op.apply_data_type_coercion!
+  op.supports_broadcasting!
+
+  op.option :name, "Optional name", :nil
+
+  op.define_gradient do |grad, node, params|
+    x, y = params
+    sx = i_op(:shape, x)
+    sy = i_op(:shape, y)
+    rx, ry = _broadcast_gradient_args(sx, 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),]
+  end
+end

data/lib/tensor_stream/ops/equal.rb

@@ -0,0 +1,15 @@
+TensorStream::OpMaker.define_operation :equal do |op|
+  op.what_it_does "Returns the truth value of (x == y) element-wise."
+
+  op.parameter :input_a, "tensor X"
+  op.parameter :input_b, "tensor Y"
+
+  op.apply_data_type_coercion!
+  op.supports_broadcasting!
+
+  op.option :name, "Optional name", :nil
+
+  op.define_gradient do |grad, node, params|
+    _min_or_max_grad(node.inputs, grad, ->(a, b) { ts.equal(a, b) })
+  end
+end

data/lib/tensor_stream/ops/expand_dims.rb

@@ -0,0 +1,17 @@
+TensorStream::OpMaker.define_operation :expand_dims do |op|
+  op.what_it_does "Inserts a dimension of 1 into a tensor's shape. "
+  op.what_it_does "Given a tensor input, this operation inserts a dimension of 1 at the dimension index axis of input's shape. The "
+  op.what_it_does "dimension index axis starts at zero; if you specify a negative number for axis it is counted backward from the end."
+
+  op.parameter :input, "A tensor"
+  op.parameter :axis, "Specifies the dimension index at which to expand the shape of input. Must be in the range [-rank(input) - 1, rank(input)]."
+  op.option :name, "Optional name", :nil
+
+  op.define_gradient do |grad, node, params|
+    [_reshape_to_input(node, grad), nil]
+  end
+
+  op.define_shape do |tensor|
+    nil
+  end
+end

data/lib/tensor_stream/ops/fill.rb

@@ -0,0 +1,19 @@
+TensorStream::OpMaker.define_operation :fill do |op|
+  op.what_it_does "This operation creates a tensor of shape dims and fills it with value."
+
+  op.parameter :dims, "tensor shape"
+  op.parameter :value, "scalar value to fill with"
+
+  op.option :name, "Optional name", :nil
+
+  op.define_gradient do |grad, node, params|
+    [nil, TensorStream.reduce_sum(grad)]
+  end
+
+  op.define_shape do |tensor|
+    a_shape = tensor.inputs[0] ? tensor.inputs[0].const_value : tensor.options[:shape]
+    next nil if a_shape.nil?
+
+    a_shape.is_a?(Array) ? a_shape : [a_shape]
+  end
+end

data/lib/tensor_stream/ops/floor.rb

@@ -0,0 +1,15 @@
+TensorStream::OpMaker.define_operation :floor do |op|
+  op.what_it_does "Returns element-wise largest integer not greater than x."
+
+  op.parameter :input_a, "tensor X", validate: 'FLOATING_POINT_TYPES'
+
+  op.option :name, "Optional name", :nil
+
+  op.define_gradient do |grad, node, params|
+    nil
+  end
+
+  op.define_shape do |tensor|
+    tensor.inputs[0].shape.shape
+  end
+end

data/lib/tensor_stream/ops/floor_div.rb

@@ -0,0 +1,15 @@
+TensorStream::OpMaker.define_operation :floor_div do |op|
+  op.what_it_does "Returns element-wise integer divistion."
+
+  op.parameter :input_a, "tensor X"
+  op.parameter :input_b, "tensor Y"
+
+  op.apply_data_type_coercion!
+  op.supports_broadcasting!
+
+  op.option :name, "Optional name", :nil
+
+  op.define_gradient do |grad, node, params|
+    [nil, nil]
+  end
+end

data/lib/tensor_stream/ops/greater.rb

@@ -0,0 +1,11 @@
+TensorStream::OpMaker.define_operation :greater do |op|
+  op.what_it_does "Returns the truth value of (x > y) element-wise."
+
+  op.parameter :input_a, "tensor X"
+  op.parameter :input_b, "tensor Y"
+
+  op.apply_data_type_coercion!
+  op.supports_broadcasting!
+
+  op.option :name, "Optional name", :nil
+end

data/lib/tensor_stream/ops/greater_equal.rb

@@ -0,0 +1,11 @@
+TensorStream::OpMaker.define_operation :greater_equal do |op|
+  op.what_it_does "Returns the truth value of (x >= y) element-wise."
+
+  op.parameter :input_a, "tensor X"
+  op.parameter :input_b, "tensor Y"
+
+  op.apply_data_type_coercion!
+  op.supports_broadcasting!
+
+  op.option :name, "Optional name", :nil
+end

data/lib/tensor_stream/ops/less_equal.rb

@@ -0,0 +1,15 @@
+TensorStream::OpMaker.define_operation :less_equal do |op|
+  op.what_it_does "Returns the truth value of (x <= y) element-wise."
+
+  op.parameter :input_a, "tensor X"
+  op.parameter :input_b, "tensor Y"
+
+  op.apply_data_type_coercion!
+  op.supports_broadcasting!
+
+  op.option :name, "Optional name", :nil
+
+  op.define_gradient do |grad, node, params|
+    _min_or_max_grad(node.inputs, grad, ->(a, b) { ts.greater_equal(a, b) })
+  end
+end

data/lib/tensor_stream/ops/log.rb

@@ -0,0 +1,14 @@
+TensorStream::OpMaker.define_operation :log do |op|
+  op.what_it_does "Computes natural logarithm of x element-wise."
+
+  op.parameter :input, "tensor X"
+  op.option :name, "Optional name", :nil
+
+  op.define_gradient do |grad, node, params|
+    grad * TensorStream.reciprocal(params[0])
+  end
+
+  op.define_shape do |tensor|
+    tensor.inputs[0].shape.shape
+  end
+end

data/lib/tensor_stream/ops/mat_mul.rb

@@ -0,0 +1,60 @@
+TensorStream::OpMaker.define_operation :mat_mul do |op|
+  op.other_names %w(matmul)
+  op.what_it_does "Multiplies matrix a by matrix b, producing a * b. The inputs must, following any transpositions, be tensors of rank 2 ."
+
+  op.parameter :input_a, "tensor X"
+  op.parameter :input_b, "tensor Y"
+
+  op.apply_data_type_coercion!
+  op.supports_broadcasting!
+
+  op.option :transpose_a, "Transpose matrix A first", :false
+  op.option :transpose_b, "Transpose matrix B first", :false
+  op.option :name, "Optional name", :nil
+
+  op.define_gradient do |grad, node, params|
+    x, y = params
+    t_a = node.options[:transpose_a]
+    t_b = node.options[:transpose_b]
+
+    if !t_a && !t_b
+      grad_a = ts.matmul(grad, y, transpose_b: true)
+      grad_b = ts.matmul(x, grad, transpose_a: true)
+    elsif !ta && tb
+      grad_a = ts.matmul(grad, y)
+      grad_b = ts.matmul(grad, x, transpose_a: true)
+    elsif t_a && !t_b
+      grad_a = ts.matmul(y, grad, transpose_b: true)
+      grad_b = ts.matmul(x, grad)
+    elsif t_a && t_b
+      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]
+  end
+
+  op.define_shape do |tensor|
+    next nil if tensor.inputs[0].shape.shape.nil? || tensor.inputs[1].shape.shape.nil?
+    next [] if tensor.inputs[0].shape.shape.empty? || tensor.inputs[1].shape.shape.empty?
+    next 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
+
+    next 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]
+  end
+end

data/lib/tensor_stream/ops/max.rb

@@ -0,0 +1,15 @@
+TensorStream::OpMaker.define_operation :max do |op|
+  op.what_it_does "Returns the max of x and y (i.e. x > y ? x : y) element-wise."
+
+  op.parameter :input_a, "tensor X", nil, validate: 'NUMERIC_TYPES'
+  op.parameter :input_b, "tensor Y", nil, validate: 'NUMERIC_TYPES'
+
+  op.apply_data_type_coercion!
+  op.supports_broadcasting!
+
+  op.option :name, "Optional name", :nil
+
+  op.define_gradient do |grad, node, params|
+    _min_or_max_grad(node.inputs, grad, ->(a, b) { ts.greater_equal(a, b) })
+  end
+end

data/lib/tensor_stream/ops/min.rb

@@ -0,0 +1,15 @@
+TensorStream::OpMaker.define_operation :min do |op|
+  op.what_it_does "Returns the min of x and y (i.e. x < y ? x : y) element-wise."
+
+  op.parameter :input_a, "tensor X", nil, validate: 'NUMERIC_TYPES'
+  op.parameter :input_b, "tensor Y", nil, validate: 'NUMERIC_TYPES'
+
+  op.apply_data_type_coercion!
+  op.supports_broadcasting!
+
+  op.option :name, "Optional name", :nil
+
+  op.define_gradient do |grad, node, params|
+    _min_or_max_grad(node.inputs, grad, ->(a, b) { ts.less_equal(a, b) })
+  end
+end

data/lib/tensor_stream/ops/mod.rb

@@ -0,0 +1,23 @@
+TensorStream::OpMaker.define_operation :mod do |op|
+  op.what_it_does "Returns element-wise remainder of division."
+
+  op.parameter :input_a, "tensor X"
+  op.parameter :input_b, "tensor Y"
+
+  op.apply_data_type_coercion!
+  op.supports_broadcasting!
+
+  op.option :name, "Optional name", :nil
+
+  op.define_gradient do |grad, node, params|
+    x, y = params
+    sx = ts.shape(x)
+    sy = ts.shape(y)
+    rx, ry = _broadcast_gradient_args(sx, 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]
+  end
+end

data/lib/tensor_stream/ops/mul.rb

@@ -0,0 +1,21 @@
+TensorStream::OpMaker.define_operation :mul do |op|
+  op.what_it_does "Returns x * y element-wise."
+
+  op.parameter :input_a, "tensor X"
+  op.parameter :input_b, "tensor Y"
+
+  op.apply_data_type_coercion!
+  op.supports_broadcasting!
+
+  op.option :name, "Optional name", :nil
+
+  op.define_gradient do |grad, node, params|
+    x, y = params
+    sx = ts.shape(x)
+    sy = ts.shape(y)
+    rx, ry = _broadcast_gradient_args(sx, sy)
+
+    [ts.reshape(ts.reduce_sum(ts.mul(grad, y), rx), sx),
+     ts.reshape(ts.reduce_sum(ts.mul(x, grad), ry), sy)]
+  end
+end

data/lib/tensor_stream/ops/negate.rb

@@ -0,0 +1,14 @@
+TensorStream::OpMaker.define_operation :negate do |op|
+  op.what_it_does "Computes numerical negative value element-wise."
+
+  op.parameter :input, "tensor X"
+  op.option :name, "Optional name", :nil
+
+  op.define_gradient do |grad, node, params|
+    -grad
+  end
+
+  op.define_shape do |tensor|
+    tensor.inputs[0].shape.shape
+  end
+end