broadcasted equal operation

norch/autograd/functions.py

@@ -216,6 +216,7 @@ class MaxBackward:
             grad_output = (grad_output * mask) / mask.sum().tensor.contents.data[0]
 
         else:
+            
             if not self.keepdim:
                 # Remove dimensions of size 1 from the gradient tensor.
                 input_shape = [s for i, s in enumerate(input_shape) if i != self.axis]
@@ -225,8 +226,12 @@ class MaxBackward:
             grad_output_shape.insert(self.axis, 1)
             grad_output = gradient.reshape(grad_output_shape)
             grad_output = grad_output + self.input[0].zeros_like()
-
+            
+            print(self.input[0])
             max_value = self.input[0].max()
+            print(max_value)
+            max_values = self.input[0].max(axis=self.axis, keepdim=True)
+            print('\n\n', max_values, '@@@')      
             mask = self.input[0].equal(max_value)
 
             grad_output = (grad_output * mask) / mask.sum().tensor.contents.data[0]

norch/csrc/cpu.cpp

@@ -306,6 +306,45 @@ void equal_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
     }
 }
 
+void equal_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size) {
+    int max_ndim = tensor1->ndim > tensor2->ndim ? tensor1->ndim : tensor2->ndim;
+
+    // Calculate strides for broadcasting
+    int* strides1 = (int*)malloc(max_ndim * sizeof(int));
+    int* strides2 = (int*)malloc(max_ndim * sizeof(int));
+    if (strides1 == NULL || strides2 == NULL) {
+        fprintf(stderr, "Memory allocation failed\n");
+        exit(1);
+    }
+
+    int stride1 = 1, stride2 = 1;
+    for (int i = max_ndim - 1; i >= 0; i--) {
+        int dim1 = i < tensor1->ndim ? tensor1->shape[tensor1->ndim - max_ndim + i] : 1;
+        int dim2 = i < tensor2->ndim ? tensor2->shape[tensor2->ndim - max_ndim + i] : 1;
+        strides1[i] = dim1 == broadcasted_shape[i] ? stride1 : 0;
+        strides2[i] = dim2 == broadcasted_shape[i] ? stride2 : 0;
+        stride1 *= (dim1 == broadcasted_shape[i]) ? dim1 : 1;
+        stride2 *= (dim2 == broadcasted_shape[i]) ? dim2 : 1;
+    }
+
+    // Perform element-wise equal with broadcasting
+    for (int i = 0; i < broadcasted_size; i++) {
+        int index1 = 0, index2 = 0;
+        int linear_index = i;
+        for (int j = max_ndim - 1; j >= 0; j--) {
+            int pos = linear_index % broadcasted_shape[j];
+            linear_index /= broadcasted_shape[j];
+            if (strides1[j] != 0) index1 += pos * strides1[j];
+            if (strides2[j] != 0) index2 += pos * strides2[j];
+        }
+        result_data[i] = (tensor1->data[index1] == tensor2->data[index2]) ? 1.0f : 0.0f;
+    }
+
+    // Free strides
+    free(strides1);
+    free(strides2);
+}
+
 
 void ones_like_tensor_cpu(Tensor* tensor, float* result_data) {
     

norch/csrc/cpu.h

@@ -22,6 +22,7 @@ void tensor_pow_scalar_cpu(Tensor* tensor, float exponent, float* result_data);
 void log_tensor_cpu(Tensor* tensor, float* result_data);
 void scalar_mul_tensor_cpu(Tensor* tensor, float scalar, float* result_data);
 void equal_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
+void equal_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size);
 void ones_like_tensor_cpu(Tensor* tensor, float* result_data);
 void zeros_like_tensor_cpu(Tensor* tensor, float* result_data);
 void transpose_1D_tensor_cpu(Tensor* tensor, float* result_data);

norch/csrc/cuda.cu

@@ -564,6 +564,67 @@ __host__ void equal_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_
     cudaDeviceSynchronize();
 }
 
+__global__ void equal_broadcasted_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int* broadcasted_shape, int* strides1, int*strides2, int max_ndim, int size) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= size) return;
+    
+    int index1 = 0, index2 = 0;
+    int linear_index = i;
+    for (int j = max_ndim - 1; j >= 0; j--) {
+        int pos = linear_index % broadcasted_shape[j];
+        linear_index /= broadcasted_shape[j];
+        if (strides1[j] != 0) index1 += pos * strides1[j];
+        if (strides2[j] != 0) index2 += pos * strides2[j];
+    }
+    result_data[i] = (data1[index1] == data2[index2]) ? 1.0f : 0.0f;
+}
+
+__host__ void equal_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size) {
+    int max_ndim = tensor1->ndim > tensor2->ndim ? tensor1->ndim : tensor2->ndim;
+
+    int* strides1 = (int*)malloc(max_ndim * sizeof(int));
+    int* strides2 = (int*)malloc(max_ndim * sizeof(int));
+    if (strides1 == NULL || strides2 == NULL) {
+        fprintf(stderr, "Memory allocation failed\n");
+        exit(1);
+    }
+
+    int stride1 = 1, stride2 = 1;
+    for (int i = max_ndim - 1; i >= 0; i--) {
+        int dim1 = i < tensor1->ndim ? tensor1->shape[tensor1->ndim - max_ndim + i] : 1;
+        int dim2 = i < tensor2->ndim ? tensor2->shape[tensor2->ndim - max_ndim + i] : 1;
+        strides1[i] = dim1 == broadcasted_shape[i] ? stride1 : 0;
+        strides2[i] = dim2 == broadcasted_shape[i] ? stride2 : 0;
+        stride1 *= (dim1 == broadcasted_shape[i]) ? dim1 : 1;
+        stride2 *= (dim2 == broadcasted_shape[i]) ? dim2 : 1;
+    }
+    
+    int* d_broadcasted_shape;
+    int* d_strides1;
+    int* d_strides2;
+
+    cudaMalloc((void**)&d_broadcasted_shape, max_ndim * sizeof(int));
+    cudaMemcpy(d_broadcasted_shape, broadcasted_shape, max_ndim * sizeof(int), cudaMemcpyHostToDevice);
+
+    cudaMalloc((void**)&d_strides1, max_ndim * sizeof(int));
+    cudaMemcpy(d_strides1, strides1, max_ndim * sizeof(int), cudaMemcpyHostToDevice);
+
+    cudaMalloc((void**)&d_strides2, max_ndim * sizeof(int));
+    cudaMemcpy(d_strides2, strides2, max_ndim * sizeof(int), cudaMemcpyHostToDevice);
+
+    int number_of_blocks = (broadcasted_size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+    equal_broadcasted_tensor_cuda_kernel<<<number_of_blocks, THREADS_PER_BLOCK>>>(tensor1->data, tensor2->data, result_data, d_broadcasted_shape, d_strides1, d_strides2, max_ndim, broadcasted_size);
+
+    cudaError_t error = cudaGetLastError();
+    if (error != cudaSuccess) {
+        printf("CUDA error: %s\n", cudaGetErrorString(error));
+        exit(-1);
+    }
+
+    cudaDeviceSynchronize();
+    cudaFree(d_broadcasted_shape);
+}
+
 
 __global__ void ones_like_tensor_cuda_kernel(float* data, float* result_data, int size) {
     

norch/csrc/cuda.h

@@ -52,6 +52,9 @@
     __global__ void equal_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size);
     __host__ void equal_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data);
 
+    __global__ void equal_broadcasted_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int* broadcasted_shape, int* strides1, int*strides2, int max_ndim, int size);
+    __host__ void equal_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size);
+
     __global__ void ones_like_tensor_cuda_kernel(float* data, float* result_data, int size);
     __host__ void ones_like_tensor_cuda(Tensor* tensor, float* result_data);
 

norch/csrc/tensor.cpp

@@ -1084,6 +1084,54 @@ extern "C" {
         }     
     }
 
+    Tensor* equal_broadcasted_tensor(Tensor* tensor1, Tensor* tensor2) {
+
+        if (strcmp(tensor1->device, tensor2->device) != 0) {
+            fprintf(stderr, "Tensors must be on the same device: %s and %s\n", tensor1->device, tensor2->device);
+            exit(1);
+        }
+
+        int max_ndim = tensor1->ndim > tensor2->ndim ? tensor1->ndim : tensor2->ndim;
+
+        // Determine the broadcasted shape
+        int* broadcasted_shape = (int*)malloc(max_ndim * sizeof(int));
+        if (broadcasted_shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+        for (int i = 0; i < max_ndim; i++) {
+            int dim1 = i < tensor1->ndim ? tensor1->shape[tensor1->ndim - 1 - i] : 1;
+            int dim2 = i < tensor2->ndim ? tensor2->shape[tensor2->ndim - 1 - i] : 1;
+            if (dim1 != dim2 && dim1 != 1 && dim2 != 1) {
+                fprintf(stderr, "Shapes are not compatible for broadcasting\n");
+                exit(1);
+            }
+            broadcasted_shape[max_ndim - 1 - i] = dim1 > dim2 ? dim1 : dim2;
+        }
+
+        int broadcasted_size = 1;
+        for (int i = 0; i < max_ndim; i++) {
+            broadcasted_size *= broadcasted_shape[i];
+        }
+
+        if (strcmp(tensor1->device, "cuda") == 0) {
+            float* result_data;
+            cudaMalloc((void **)&result_data, broadcasted_size * sizeof(float));
+            equal_broadcasted_tensor_cuda(tensor1, tensor2, result_data, broadcasted_shape, broadcasted_size);
+            return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device);
+        } 
+        else {
+            float* result_data = (float*)malloc(broadcasted_size * sizeof(float));
+            if (result_data == NULL) {
+                fprintf(stderr, "Memory allocation failed\n");
+                exit(1);
+            }
+
+            equal_broadcasted_tensor_cpu(tensor1, tensor2, result_data, broadcasted_shape, broadcasted_size);
+            return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device);
+        }
+    }
+
 
     Tensor* ones_like_tensor(Tensor* tensor) {
         char* device = (char*)malloc(strlen(tensor->device) + 1);

norch/csrc/tensor.h

@@ -31,6 +31,7 @@ extern "C" {
     Tensor* scalar_pow_tensor(float base, Tensor* tensor);
     Tensor* log_tensor(Tensor* tensor);
     Tensor* equal_tensor(Tensor* tensor1, Tensor* tensor2);
+    Tensor* equal_broadcasted_tensor(Tensor* tensor1, Tensor* tensor2);
     void to_device(Tensor* tensor, char* device);
     Tensor* ones_like_tensor(Tensor* tensor);
     Tensor* zeros_like_tensor(Tensor* tensor);

norch/tensor.py

@@ -656,25 +656,59 @@ class Tensor:
                 return self.equal(other)
             else:
                 return False
-
-        if self.shape != other.shape:
-            return False
         
-        Tensor._C.equal_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
-        Tensor._C.equal_tensor.restype = ctypes.POINTER(CTensor)
+        broadcasted_shape_add = []
 
-        result_tensor_ptr = Tensor._C.equal_tensor(self.tensor, other.tensor)
+        # Function to determine if broadcasting is needed and get the broadcasted shape
+        def broadcast_shape(shape1, shape2):
+            if shape1 == shape2:
+                return shape1, False
+            
+            max_len = max(len(shape1), len(shape2))
+            shape1 = [1] * (max_len - len(shape1)) + shape1
+            shape2 = [1] * (max_len - len(shape2)) + shape2
 
-        result_data = Tensor()
-        result_data.tensor = result_tensor_ptr
-        result_data.shape = self.shape.copy()
-        result_data.ndim = self.ndim
-        result_data.device = self.device
-        result_data.numel = self.numel
+            
+            for dim1, dim2 in zip(shape1, shape2):
+                if dim1 != dim2 and dim1 != 1 and dim2 != 1:
+                    raise ValueError("Shapes are not compatible for broadcasting")
+                broadcasted_shape_add.append(max(dim1, dim2))
+            return broadcasted_shape_add, True
+
+        broadcasted_shape_add, needs_broadcasting = broadcast_shape(self.shape, other.shape)
         
+        if needs_broadcasting:
+            # Call equal_broadcasted_tensor if broadcasting is needed
+            Tensor._C.equal_broadcasted_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
+            Tensor._C.equal_broadcasted_tensor.restype = ctypes.POINTER(CTensor)
+
+            result_tensor_ptr = Tensor._C.equal_broadcasted_tensor(self.tensor, other.tensor)
+
+            result_data = Tensor()
+            result_data.tensor = result_tensor_ptr
+            result_data.shape = broadcasted_shape_add.copy()
+            result_data.ndim = len(broadcasted_shape_add)
+
+            result_data.device = self.device
+            result_data.numel = 1
+            for s in result_data.shape:
+                result_data.numel *= s
+
+        else:
+            Tensor._C.equal_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
+            Tensor._C.equal_tensor.restype = ctypes.POINTER(CTensor)
+
+            result_tensor_ptr = Tensor._C.equal_tensor(self.tensor, other.tensor)
+
+            result_data = Tensor()
+            result_data.tensor = result_tensor_ptr
+            result_data.shape = self.shape.copy()
+            result_data.ndim = self.ndim
+            result_data.device = self.device
+            result_data.numel = self.numel
+            
         return result_data
 
-
     def log(self):
         Tensor._C.log_tensor.argtypes = [ctypes.POINTER(CTensor)]
         Tensor._C.log_tensor.restype = ctypes.POINTER(CTensor)

tests/test_operations.py

@@ -546,6 +546,21 @@ class TestTensorOperations(unittest.TestCase):
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
 
+    def test_broadcasted_equal(self):
+        """
+        Test broadcasted equal two tensors: tensor1.equal(tensor2)
+        """
+        norch_tensor1 = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device)
+        norch_tensor2 = norch.Tensor([[[10, 10]], [[5, 6]]]).to(self.device)
+        norch_result = norch_tensor1.equal(norch_tensor2)
+        torch_result = utils.to_torch(norch_result).to(self.device)
+
+        torch_tensor1 = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device)
+        torch_tensor2 = torch.tensor([[[10, 10]], [[5, 6]]]).to(self.device)
+        torch_expected = (torch_tensor1 == torch_tensor2).float()
+
+        self.assertTrue(utils.compare_torch(torch_result, torch_expected))
+
 
 
     def test_zeros_like(self):