torchrec 0.0.1 → 0.0.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: b20b0c478258d789f1b6a6b2f80e3350600b6ce23f6ee1146cfb6819551a1150
-  data.tar.gz: 443383a994cd25d6225f84401e0a47886b848e590551ae0eec2eb31dba9468bd
+  metadata.gz: 573d37473e5a3e8e24b60bfe886ef1d14be87833ecc7c207a349c236480f1d6e
+  data.tar.gz: 9fcc58e9201bb545582feaa50a56727dd074688b4c830e495abd452aa31ad572
 SHA512:
-  metadata.gz: a8e6ce693978cd21505f4b7378efe92a87ee672ece2d012cf3ee69dc9a3a3332cc57a84d1e0edcc7c317385aaae212a8351e54d5ea33ad338f795ec7138c5cdc
-  data.tar.gz: eb105de85d924fd2d3c7ef444a61caecdfb8008a5f25dba7297bb233b04f2662a47e144305301bd60f6cc1e9734bb1551946911047e8f654581b5fb192330cf2
+  metadata.gz: bb80aa64a91fb1af2dad904fc4f39b366f79765ba541334c4bf4109011249b9a6075a6d050b1aa644f42236cf037110181ed09cb74e71bd47882c10402e15101
+  data.tar.gz: 7393ad51f8882915e433f3aa8314f3d743a70b192b174854f161e2479c811ad7bb57f7f1953c0e817211056c201de5cf88d4d15e78c361e52768598155e9bad9

data/CHANGELOG.md

@@ -1,3 +1,8 @@
+## 0.0.2 (2022-03-14)
+
+- Added `JaggedTensor`
+- Added `CrossNet`, `DeepFM`, and `FactorizationMachine` modules
+
 ## 0.0.1 (2022-02-28)
 
 - First release

data/README.md

@@ -35,6 +35,9 @@ TorchRec::Models::DLRM::DenseArch.new(in_features, layer_sizes, device: nil)
 
 ```ruby
 TorchRec::Modules::Activation::SwishLayerNorm.new(input_dims, device: nil)
+TorchRec::Modules::CrossNet::CrossNet.new(in_features, num_layers)
+TorchRec::Modules::DeepFM::DeepFM.new(dense_module)
+TorchRec::Modules::DeepFM::FactorizationMachine.new
 TorchRec::Modules::MLP::MLP.new(in_size, layer_sizes, bias: true, activation: :relu, device: nil)
 TorchRec::Modules::MLP::Perceptron.new(in_size, out_size, bias: true, activation: Torch.method(:relu), device: nil)
 ```

data/lib/torchrec/modules/cross_net/cross_net.rb

@@ -0,0 +1,36 @@
+module TorchRec
+  module Modules
+    module CrossNet
+      class CrossNet < Torch::NN::Module
+        def initialize(in_features, num_layers)
+          super()
+          @num_layers = num_layers
+          @kernels = Torch::NN::ParameterList.new(
+            @num_layers.times.map do |i|
+              Torch::NN::Parameter.new(
+                Torch::NN::Init.xavier_normal!(Torch.empty(in_features, in_features))
+              )
+            end
+          )
+          @bias = Torch::NN::ParameterList.new(
+            @num_layers.times.map do |i|
+              Torch::NN::Parameter.new(Torch::NN::Init.zeros!(Torch.empty(in_features, 1)))
+            end
+          )
+        end
+
+        def forward(input)
+          x_0 = input.unsqueeze(2)  # (B, N, 1)
+          x_l = x_0
+
+          @num_layers.times do |layer|
+            xl_w = Torch.matmul(@kernels[layer], x_l)  # (B, N, 1)
+            x_l = x_0 * (xl_w + @bias[layer]) + x_l  # (B, N, 1)
+          end
+
+          Torch.squeeze(x_l, dim: 2)
+        end
+      end
+    end
+  end
+end

data/lib/torchrec/modules/deepfm/deepfm.rb

@@ -0,0 +1,23 @@
+module TorchRec
+  module Modules
+    module DeepFM
+      class DeepFM < Torch::NN::Module
+        def initialize(dense_module)
+          super()
+          @dense_module = dense_module
+        end
+
+        def forward(embeddings)
+          deepfm_input = flatten_input(embeddings)
+          @dense_module.call(deepfm_input)
+        end
+
+        private
+
+        def flatten_input(inputs)
+          Torch.cat(inputs.map { |input| input.flatten(1) }, dim: 1)
+        end
+      end
+    end
+  end
+end

data/lib/torchrec/modules/deepfm/factorization_machine.rb

@@ -0,0 +1,27 @@
+module TorchRec
+  module Modules
+    module DeepFM
+      class FactorizationMachine < Torch::NN::Module
+        def initialize
+          super()
+        end
+
+        def forward(embeddings)
+          fm_input = flatten_input(embeddings)
+          sum_of_input = Torch.sum(fm_input, dim: 1, keepdim: true)
+          sum_of_square = Torch.sum(fm_input * fm_input, dim: 1, keepdim: true)
+          square_of_sum = sum_of_input * sum_of_input
+          cross_term = square_of_sum - sum_of_square
+          cross_term = Torch.sum(cross_term, dim: 1, keepdim: true) * 0.5  # [B, 1]
+          cross_term
+        end
+
+        private
+
+        def flatten_input(inputs)
+          Torch.cat(inputs.map { |input| input.flatten(1) }, dim: 1)
+        end
+      end
+    end
+  end
+end

data/lib/torchrec/sparse/jagged_tensor.rb

@@ -0,0 +1,33 @@
+module TorchRec
+  module Sparse
+    class JaggedTensor
+      def initialize(values, weights: nil, lengths: nil, offsets: nil)
+        @values = values
+        @weights = weights
+        assert_offsets_or_lengths_is_provided(offsets, lengths)
+        if !offsets.nil?
+          assert_tensor_has_no_elements_or_has_integers(offsets, "offsets")
+        end
+        if !lengths.nil?
+          assert_tensor_has_no_elements_or_has_integers(lengths, "lengths")
+        end
+        @lengths = lengths
+        @offsets = offsets
+      end
+
+      private
+
+      def assert_offsets_or_lengths_is_provided(offsets, lengths)
+        if offsets.nil? && lengths.nil?
+          raise ArgumentError, "Must provide lengths or offsets"
+        end
+      end
+
+      def assert_tensor_has_no_elements_or_has_integers(tensor, tensor_name)
+        if tensor.numel != 0 && ![:int64, :int32, :int16, :int8, :uint8].include?(tensor.dtype)
+          raise ArgumentError, "#{tensor_name} must be of integer type, but got #{tensor.dtype}"
+        end
+      end
+    end
+  end
+end

data/lib/torchrec/version.rb

@@ -1,3 +1,3 @@
 module TorchRec
-  VERSION = "0.0.1"
+  VERSION = "0.0.2"
 end

data/lib/torchrec.rb

@@ -8,10 +8,16 @@ require "torchrec/models/dlrm/dense_arch"
 
 # modules
 require "torchrec/modules/activation/swish_layer_norm"
+require "torchrec/modules/cross_net/cross_net"
+require "torchrec/modules/deepfm/deepfm"
+require "torchrec/modules/deepfm/factorization_machine"
 require "torchrec/modules/mlp/mlp"
 require "torchrec/modules/mlp/perceptron"
 require "torchrec/modules/utils"
 
+# sparse
+require "torchrec/sparse/jagged_tensor"
+
 # other
 require "torchrec/version"
 

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: torchrec
 version: !ruby/object:Gem::Version
-  version: 0.0.1
+  version: 0.0.2
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2022-02-28 00:00:00.000000000 Z
+date: 2022-03-15 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: torch-rb
@@ -16,14 +16,14 @@ dependencies:
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '0'
+        version: 0.9.2
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '0'
+        version: 0.9.2
 description:
 email: andrew@ankane.org
 executables: []
@@ -38,9 +38,13 @@ files:
 - lib/torchrec/models/deepfm/over_arch.rb
 - lib/torchrec/models/dlrm/dense_arch.rb
 - lib/torchrec/modules/activation/swish_layer_norm.rb
+- lib/torchrec/modules/cross_net/cross_net.rb
+- lib/torchrec/modules/deepfm/deepfm.rb
+- lib/torchrec/modules/deepfm/factorization_machine.rb
 - lib/torchrec/modules/mlp/mlp.rb
 - lib/torchrec/modules/mlp/perceptron.rb
 - lib/torchrec/modules/utils.rb
+- lib/torchrec/sparse/jagged_tensor.rb
 - lib/torchrec/version.rb
 homepage: https://github.com/ankane/torchrec-ruby
 licenses:
@@ -61,7 +65,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.3.3
+rubygems_version: 3.3.7
 signing_key:
 specification_version: 4
 summary: Deep learning recommendation systems for Ruby