Is there a function that finds the OR of every row in a boolean matrix in pytorch?

Viewed 623

I have a matrix A of size n x m and all the entries are bools. I want all my calculation to be on the GPU and I have matrix A stored as a tensor with each entry being the pytorch's bool datatype. I want an output of a single vector, b, to be a 1 x m tensor that stores the OR of all rows in A.

What I want:
A matrix =
[a1,1, a1,2, , a1,3, ... , a1,m]
[a2,1, a2,2, , a2,3, ... , a2,m]
...
[an,1, an,2, , an,3, ... , an,m]


b = [b1, b2, , b3, , ..., bm]

s.t. bi = a1,i | a2,i | a3,i | ... | an,i
where | is the OR operator in Pytorch

essentially I want a function that applies row or column wise boolean operation. I know the | .__OR__ can be used and Pytorch's OR function takes two boolean tensor as inputs and I would need to loop over all rows to get my desired b vector.

Since OR is commutative (a|(b|c) = (a|b)|c), I would think pytorch would have some nice function that speeds it up by doing the | operations in parallel or in some divide and conquer method, instead of doing this with a loop. Any ideas or references to speed up the process of applying communicative row/column wise boolean operation with pytorch is welcomed. Best if all operations are done on the GPU.

1 Answers

See torch.any and torch.all.

Both take a dim argument and hence you can compute or/and of rows.

Related