Upload folder using huggingface_hub

b8d4c08 verified over 1 year ago

3.19 kB

import torch
import torch.nn as nn
import torch.optim as optim


class NdLinear(nn.Module):
    def __init__(self, input_dims: tuple, hidden_size: tuple, transform_outer=True, bias: bool = True):
        """
        NdLinear: A PyTorch layer for projecting tensors into multi-space representations.
        
        Unlike conventional embedding layers that map into a single vector space, NdLinear 
        transforms tensors across a collection of vector spaces, capturing multivariate structure 
        and topical information that standard deep learning architectures typically lose.
        Args:
            input_dims (tuple): Shape of input tensor (excluding batch dimension).
            hidden_size (tuple): Target hidden dimensions after transformation.
        """
        super(NdLinear, self).__init__()

        if len(input_dims) != len(hidden_size):
            raise Exception("Input shape and hidden shape do not match.")

        self.input_dims = input_dims
        self.hidden_size = hidden_size
        self.num_layers = len(input_dims)  # Must match since dims are equal
        # self.relu = nn.ReLU() # self.relu is not being used.
        self.transform_outer = transform_outer
        self.bias = bias

        # Define transformation layers per dimension
        self.align_layers = nn.ModuleList([
            nn.Linear(input_dims[i], hidden_size[i], bias=self.bias) for i in range(self.num_layers)
        ])


    def forward(self, X):
        """
        Forward pass to project input tensor into a new multi-space representation.
        - Incrementally transposes, flattens, applies linear layers, and restores shape.
        Expected Input Shape: [batch_size, *input_dims]
        Output Shape: [batch_size, *hidden_size]
        Args:
            X (torch.Tensor): Input tensor with shape [batch_size, *input_dims]
        Returns:
            torch.Tensor: Output tensor with shape [batch_size, *hidden_size]
        """
        num_transforms = self.num_layers  # Number of transformations
        
        # Define iteration order
        # transform_indices = range(num_transforms) if transform_outer else reversed(range(num_transforms))

        for i in range(num_transforms):
            if self.transform_outer:
                layer = self.align_layers[i]
                transpose_dim = i + 1
            else:
                layer = self.align_layers[num_transforms - (i+1)]
                transpose_dim = num_transforms - i

            # Transpose the selected dimension to the last position
            X = torch.transpose(X, transpose_dim, num_transforms).contiguous()

            # Store original shape before transformation
            X_size = X.shape[:-1]

            # Flatten everything except the last dimension
            X = X.view(-1, X.shape[-1])

            # Apply transformation
            X = layer(X)
            
            # Reshape back to the original spatial structure (with new embedding dim)
            X = X.view(*X_size, X.shape[-1])

            # Transpose the dimension back to its original position
            X = torch.transpose(X, transpose_dim, num_transforms).contiguous()

        return X

	import torch
	import torch.nn as nn
	import torch.optim as optim


	class NdLinear(nn.Module):
	def __init__(self, input_dims: tuple, hidden_size: tuple, transform_outer=True, bias: bool = True):
	"""
	NdLinear: A PyTorch layer for projecting tensors into multi-space representations.

	Unlike conventional embedding layers that map into a single vector space, NdLinear
	transforms tensors across a collection of vector spaces, capturing multivariate structure
	and topical information that standard deep learning architectures typically lose.

	Args:
	input_dims (tuple): Shape of input tensor (excluding batch dimension).
	hidden_size (tuple): Target hidden dimensions after transformation.
	"""
	super(NdLinear, self).__init__()

	if len(input_dims) != len(hidden_size):
	raise Exception("Input shape and hidden shape do not match.")

	self.input_dims = input_dims
	self.hidden_size = hidden_size
	self.num_layers = len(input_dims) # Must match since dims are equal
	# self.relu = nn.ReLU() # self.relu is not being used.
	self.transform_outer = transform_outer
	self.bias = bias

	# Define transformation layers per dimension
	self.align_layers = nn.ModuleList([
	nn.Linear(input_dims[i], hidden_size[i], bias=self.bias) for i in range(self.num_layers)
	])


	def forward(self, X):
	"""
	Forward pass to project input tensor into a new multi-space representation.
	- Incrementally transposes, flattens, applies linear layers, and restores shape.

	Expected Input Shape: [batch_size, *input_dims]
	Output Shape: [batch_size, *hidden_size]

	Args:
	X (torch.Tensor): Input tensor with shape [batch_size, *input_dims]

	Returns:
	torch.Tensor: Output tensor with shape [batch_size, *hidden_size]
	"""
	num_transforms = self.num_layers # Number of transformations

	# Define iteration order
	# transform_indices = range(num_transforms) if transform_outer else reversed(range(num_transforms))

	for i in range(num_transforms):
	if self.transform_outer:
	layer = self.align_layers[i]
	transpose_dim = i + 1
	else:
	layer = self.align_layers[num_transforms - (i+1)]
	transpose_dim = num_transforms - i

	# Transpose the selected dimension to the last position
	X = torch.transpose(X, transpose_dim, num_transforms).contiguous()

	# Store original shape before transformation
	X_size = X.shape[:-1]

	# Flatten everything except the last dimension
	X = X.view(-1, X.shape[-1])

	# Apply transformation
	X = layer(X)

	# Reshape back to the original spatial structure (with new embedding dim)
	X = X.view(*X_size, X.shape[-1])

	# Transpose the dimension back to its original position
	X = torch.transpose(X, transpose_dim, num_transforms).contiguous()

	return X