Function sparse_batch_op

trellis/modules/sparse/basic.py:408–417

Broadcast a 1D tensor to a sparse tensor along the batch dimension then perform an operation. Args: input (torch.Tensor): 1D tensor to broadcast. target (SparseTensor): Sparse tensor to broadcast to. op (callable): Operation to perform after broadcasting. Defaul

(input: SparseTensor, other: torch.Tensor, op: callable = torch.add)

def sparse_batch_op(input: SparseTensor, other: torch.Tensor, op: callable = torch.add) -> SparseTensor:
    """
    Broadcast a 1D tensor to a sparse tensor along the batch dimension then perform an operation.
    
    Args:
        input (torch.Tensor): 1D tensor to broadcast.
        target (SparseTensor): Sparse tensor to broadcast to.
        op (callable): Operation to perform after broadcasting. Defaults to torch.add.
    """
    return input.replace(op(input.feats, sparse_batch_broadcast(input, other)))


def sparse_cat(inputs: List[SparseTensor], dim: int = 0) -> SparseTensor: