class_docstr: `Cross Network <https://arxiv.org/abs/1708.05123>`_:
  
  Cross Net is a stack of "crossing" operations on a tensor of shape :math:`(*, N)`
  to the same shape, effectively creating :math:`N` learnable polynomical functions
  over the input tensor.
  
  In this module, the crossing operations are defined based on a full rank matrix
  (NxN), such that the crossing effect can cover all bits on each layer. On each layer
  l, the tensor is transformed into:
  
  .. math ::    x_{l+1} = x_0 * (W_l \cdot x_l + b_l) + x_l
  
  where :math:`W_l` is a square matrix :math:`(NxN)`, :math:`*` means element-wise
  multiplication, :math:`\cdot` means matrix multiplication.
  
  Args:
      in_features (int): the dimension of the input.
      num_layers (int): the number of layers in the module.
  
  Example::
  
      batch_size = 3
      num_layers = 2
      in_features = 10
      input = torch.randn(batch_size, in_features)
      dcn = CrossNet(num_layers=num_layers)
      output = dcn(input)
class_docstr_tok: ["`","Cross","Network","<","https",":","\/\/arxiv.org\/abs\/1708.05123",">","`","_",":","Cross","Net","is","a","stack","of","``","crossing","''","operations","on","a","tensor","of","shape",":","math",":","`","(","*",",","N",")","`","to","the","same","shape",",","effectively","creating",":","math",":","`","N","`","learnable","polynomical","functions","over","the","input","tensor",".","In","this","module",",","the","crossing","operations","are","defined","based","on","a","full","rank","matrix","(","NxN",")",",","such","that","the","crossing","effect","can","cover","all","bits","on","each","layer",".","On","each","layer","l",",","the","tensor","is","transformed","into",":","..","math",":",":","x_","{","l+1","}","=","x_0","*","(","W_l","\\cdot","x_l","+","b_l",")","+","x_l","where",":","math",":","`","W_l","`","is","a","square","matrix",":","math",":","`","(","NxN",")","`",",",":","math",":","`","*","`","means","element-wise","multiplication",",",":","math",":","`","\\cdot","`","means","matrix","multiplication",".","Args",":","in_features","(","int",")",":","the","dimension","of","the","input",".","num_layers","(","int",")",":","the","number","of","layers","in","the","module",".","Example",":",":","batch_size","=","3","num_layers","=","2","in_features","=","10","input","=","torch.randn","(","batch_size",",","in_features",")","dcn","=","CrossNet","(","num_layers=num_layers",")","output","=","dcn","(","input",")"]
docstr: Args:
      input (torch.Tensor): tensor with shape [batch_size, in_features].
  
  Returns:
      torch.Tensor: tensor with shape [batch_size, in_features].
docstr_tok: ["Args",":","input","(","torch.Tensor",")",":","tensor","with","shape","[","batch_size",",","in_features","]",".","Returns",":","torch.Tensor",":","tensor","with","shape","[","batch_size",",","in_features","]","."]
code: def forward(self, input: torch.Tensor) -> torch.Tensor:
      """
      Args:
          input (torch.Tensor): tensor with shape [batch_size, in_features].
  
      Returns:
          torch.Tensor: tensor with shape [batch_size, in_features].
      """
      x_0 = input.unsqueeze(2)  # (B, N, 1)
      x_l = x_0
  
      for layer in range(self._num_layers):
          xl_w = torch.matmul(self.kernels[layer], x_l)  # (B, N, 1)
          x_l = x_0 * (xl_w + self.bias[layer]) + x_l  # (B, N, 1)
  
      return torch.squeeze(x_l, dim=2)
  
file_path: torchrec/torchrec/modules/crossnet.py
imports: from typing import Callable, Optional, Union
  import torch
label_desc: Use image node_id 2 for calling the CrossNet obj's underlying member method code with example usage: obj.forward(input) and returns: torch