Spaces:

Jack04810
/

SPAC-4B-ZeroGPU

Sleeping

SPAC-4B-ZeroGPU / spac /sparse /linear.py

Build and validate native SPAC extensions for ZeroGPU Blackwell

c83d2d6 verified 6 days ago

303 Bytes

"""Apply a linear layer to sparse feature rows."""

import torch.nn as nn
from . import SparseTensor


class SparseLinear(nn.Linear):
    def forward(self, value: SparseTensor) -> SparseTensor:
        features = value.feats.to(self.weight.dtype)
        return value.replace(super().forward(features))

	"""Apply a linear layer to sparse feature rows."""

	import torch.nn as nn
	from . import SparseTensor


	class SparseLinear(nn.Linear):
	def forward(self, value: SparseTensor) -> SparseTensor:
	features = value.feats.to(self.weight.dtype)
	return value.replace(super().forward(features))