Added logical operators that return non-ByteTensors
The existing equality operators always return ByteTensors. Often the
return values from an equality operation needs to be used with a tensor
of the original type (e.g. point-wise multiplication with a mask), and
the ByteTensor therefore has to be cast back to this type. This casting
is unneccessary, and for CudaTensors would be costly, as it involves a
round-trip to the host.
This commit adds new forms of the equality operators that directly
writes the result into a tensor of the approriate type. For example
x = torch.DoubleTensor(5)
y = torch.rand(5)
x:lt(y,0.5)
The existing ByteTensor-returning forms are still present.
diff --git a/generic/THTensorMath.c b/generic/THTensorMath.c
index 017a791..c884a09 100644
--- a/generic/THTensorMath.c
+++ b/generic/THTensorMath.c
@@ -997,6 +997,13 @@
TH_TENSOR_APPLY2(unsigned char, r_, real, t, \
if (*t_data OP value) *r__data = 1;); \
} \
+ void THTensor_(NAME##ValueT)(THTensor* r_, THTensor* t, real value) \
+ { \
+ THTensor_(rawResize)(r_, t->nDimension, t->size, NULL); \
+ THTensor_(zero)(r_); \
+ TH_TENSOR_APPLY2(real, r_, real, t, \
+ if (*t_data OP value) *r__data = 1;); \
+ } \
void THTensor_(NAME##Tensor)(THByteTensor *r_, THTensor *ta, THTensor *tb) \
{ \
THByteTensor_rawResize(r_, ta->nDimension, ta->size, NULL); \
@@ -1004,6 +1011,14 @@
TH_TENSOR_APPLY3(unsigned char, r_, real, ta, real, tb, \
if(*ta_data OP *tb_data) *r__data = 1;); \
} \
+ void THTensor_(NAME##TensorT)(THTensor *r_, THTensor *ta, THTensor *tb) \
+ { \
+ THTensor_(rawResize)(r_, ta->nDimension, ta->size, NULL); \
+ THTensor_(zero)(r_); \
+ TH_TENSOR_APPLY3(real, r_, real, ta, real, tb, \
+ if(*ta_data OP *tb_data) *r__data = 1;); \
+ } \
+
TENSOR_IMPLEMENT_LOGICAL(lt,<)
TENSOR_IMPLEMENT_LOGICAL(gt,>)
