bananagator
bananagator

Reputation: 581

Reduce torch tensor

For a boolean tensor of shape (15,10), I want to perform bitwise_or along axis 0 so that the resulting tensor would be of shape 10. torch.bitwise_or does not support this.

I know it can done in numpy using np.bitwise_or.reduce(x,axis=0). I did not find something similar in torch. How to reduce torch tensor?

Upvotes: 1

Views: 3658

Answers (1)

Edwin Cheong
Edwin Cheong

Reputation: 979

Hi figured out the problem here if you look at the docstring for the reduce function it's essentially just a for loop adding itself from 0

# ufunc docstring
# op.identiy is 0
r = op.identity # op = ufunc 
for i in range(len(A)):
  r = op(r, A[i])
return r

So to solve and fix your problem

import numpy as np
import torch

bool_arr = np.random.randint(0, 2, (15, 10), dtype=np.bool) # create a bool arr
tensor_bool_arr = torch.tensor(bool_arr) # Create torch version
output_np = np.bitwise_or.reduce(bool_arr, axis=0) 
# array([ True,  True,  True,  True,  True,  True,  True,  True,  True,
        True])

# Create a pytorch equivalent of bitwise reduce
r = torch.tensor(0)
for i in range(len(tensor_bool_arr)):
    r = torch.bitwise_or(r, tensor_bool_arr[i])

torch_output = r.type(torch.bool)
# tensor([True, True, True, True, True, True, True, True, True, True])
assert torch_output.shape[0] == np_output.shape[0]

Upvotes: 1

Related Questions