Function sparse_batch_broadcast

trellis/modules/sparse/basic.py:392–405

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)

def sparse_batch_broadcast(input: SparseTensor, other: torch.Tensor) -> torch.Tensor:
    """
    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.
    """
    coords, feats = input.coords, input.feats
    broadcasted = torch.zeros_like(feats)
    for k in range(input.shape[0]):
        broadcasted[input.layout[k]] = other[k]
    return broadcasted


def sparse_batch_op(input: SparseTensor, other: torch.Tensor, op: callable = torch.add) -> SparseTensor: