Hi Suraj!
This question is a little simpler if the original tensor is square, that is,
it has the same number or rows as columns. (Not required, but you
don’t have to fiddle around checking sizes.)
Here’s the most compact approach I could find, illustrated as a
pytorch version 0.3.0 script:
import torch
torch.__version__
# use a version of A that is square
# note, A has a leading size = 1 dimension
A = torch.FloatTensor ([ [ [0,0,0,0], [0,-0.4,0,-0.9], [0,-0.7,0.9,0.2], [0,0,0,0] ] ])
rows = (A.sum (dim = 2) != 0).squeeze()
# delete the rows and corresponding columns
C = A[torch.ger (rows, rows)].view ((1, rows.sum(), rows.sum()))
A
C
# get the non-zero-row indices
torch.arange (len (rows)).long()[rows]
# get the zero-row indices
torch.arange (len (rows)).long()[~rows]
Here is the output:
>>> import torch
>>> torch.__version__
'0.3.0b0+591e73e'
>>>
>>> # use a version of A that is square
... # note, A has a leading size = 1 dimension
...
>>> A = torch.FloatTensor ([ [ [0,0,0,0], [0,-0.4,0,-0.9], [0,-0.7,0.9,0.2], [0,0,0,0] ] ])
>>> rows = (A.sum (dim = 2) != 0).squeeze()
>>>
>>> # delete the rows and corresponding columns
...
>>> C = A[torch.ger (rows, rows)].view ((1, rows.sum(), rows.sum()))
>>> A
(0 ,.,.) =
0.0000 0.0000 0.0000 0.0000
0.0000 -0.4000 0.0000 -0.9000
0.0000 -0.7000 0.9000 0.2000
0.0000 0.0000 0.0000 0.0000
[torch.FloatTensor of size 1x4x4]
>>> C
(0 ,.,.) =
-0.4000 0.0000
-0.7000 0.9000
[torch.FloatTensor of size 1x2x2]
>>>
>>> # get the non-zero-row indices
...
>>> torch.arange (len (rows)).long()[rows]
1
2
[torch.LongTensor of size 2]
>>>
>>> # get the zero-row indices
...
>>> torch.arange (len (rows)).long()[~rows]
0
3
[torch.LongTensor of size 2]
I would do it by indexing into a range, as illustrated in the above script.
Best.
K. Frank