tfg.geometry.convolution.utils.convert_to_block_diag_2d

View source on GitHub

Convert a batch of 2d SparseTensors to a 2d block diagonal SparseTensor.

Note:

In the following, A1 to An are optional batch dimensions.

A SparseTensor with dense shape [A1, ..., An, D1, D2] will be reshaped to one with shape [A1*...*An*D1, A1*...*An*D2].

Padded inputs in dims D1 and D2 are allowed. sizes indicates the un-padded shape for each inner [D1, D2] matrix. The additional (padded) rows and columns will be omitted in the block diagonal output.

If padded (sizes != None), the input should not contain any sparse indices outside the bounds indicated by sizes. Setting validate_indices=True will explicitly filter any invalid sparse indices before block diagonalization.

data A SparseTensor with dense shape [A1, ..., An, D1, D2].
sizes A tensor with shape [A1, ..., An, 2]. Can be None (indicates no padding). If not None, sizes indicates the true sizes (before padding) of the inner dimensions of data.
validate_indices A boolean. Ignored if sizes==None. If True, out-of-bounds indices in data are explicitly ignored, otherwise out-of-bounds indices will cause undefined behavior.
name A name for this op. Defaults to 'utils_convert_to_block_diag_2d'.

A 2d block-diagonal SparseTensor.

TypeError if the input types are invalid.
ValueError if the input dimensions are invalid.