Fix segmentation fault with complex Variable operations and assign_add

Fixes tensorflow#105367

The issue was that complex types (complex64, complex128) were missing from:
1. GPU DenseUpdate functor template instantiations for ADD/SUB operations
2. GPU kernel registrations for AssignAddVariableOp and AssignSubVariableOp

This caused a segmentation fault when using assign_add on complex Variables,
particularly when combined with tf.raw_ops.Conj operations.

Changes:
- Added TF_CALL_COMPLEX_TYPES to dense_update_functor_gpu.cu.cc for ADD/SUB
- Added TF_CALL_COMPLEX_TYPES to GPU kernel registrations in resource_variable_ops.cc
- Added comprehensive test cases for complex variable assign_add operations

Author: CodersAcademy006

tensorflow/core/kernels/dense_update_functor_gpu.cu.cc

@@ -57,6 +57,7 @@ struct DenseUpdate<GPUDevice, T, SUB> {
   template struct functor::DenseUpdate<GPUDevice, T, SUB>;
 TF_CALL_GPU_NUMBER_TYPES(DEFINE_GPU_KERNELS);
 TF_CALL_INTEGRAL_TYPES(DEFINE_GPU_KERNELS);
+TF_CALL_COMPLEX_TYPES(DEFINE_GPU_KERNELS);
 TF_CALL_float8_e5m2(DEFINE_GPU_KERNELS);
 TF_CALL_float8_e4m3fn(DEFINE_GPU_KERNELS);
 #undef DEFINE_GPU_KERNELS

tensorflow/core/kernels/resource_variable_ops.cc

@@ -681,6 +681,7 @@ TF_CALL_NUMBER_TYPES(REGISTER_KERNELS);
 
 TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_KERNELS);
 TF_CALL_INTEGRAL_TYPES_NO_INT32(REGISTER_GPU_KERNELS);
+TF_CALL_COMPLEX_TYPES(REGISTER_GPU_KERNELS);
 #undef REGISTER_GPU_KERNELS
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 

tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py

@@ -1864,5 +1864,64 @@ def testGatherBatchDimsNeg(self):
         )
         self.evaluate(result)
 
+  @test_util.run_in_graph_and_eager_modes
+  @test_util.run_gpu_only
+  def testComplexVariableAssignAddWithConj(self):
+    """Test for issue #105367: Segfault with complex Variable, Conj, and assign_add."""
+    # Test with complex64
+    input_data_64 = constant_op.constant([1 + 2j, 3 + 4j], dtype=dtypes.complex64)
+    var_64 = resource_variable_ops.ResourceVariable(input_data_64, dtype=dtypes.complex64)
+    self.evaluate(var_64.initializer)
+    
+    conj_result_64 = math_ops.conj(input_data_64)
+    assign_add_op_64 = var_64.assign_add(conj_result_64)
+    result_64 = self.evaluate(assign_add_op_64)
+    
+    # Expected: [1+2j, 3+4j] + [1-2j, 3-4j] = [2+0j, 6+0j]
+    expected_64 = np.array([2+0j, 6+0j], dtype=np.complex64)
+    self.assertAllClose(result_64, expected_64)
+    
+    # Test with complex128
+    input_data_128 = constant_op.constant([1 + 2j, 3 + 4j], dtype=dtypes.complex128)
+    var_128 = resource_variable_ops.ResourceVariable(input_data_128, dtype=dtypes.complex128)
+    self.evaluate(var_128.initializer)
+    
+    conj_result_128 = math_ops.conj(input_data_128)
+    assign_add_op_128 = var_128.assign_add(conj_result_128)
+    result_128 = self.evaluate(assign_add_op_128)
+    
+    # Expected: [1+2j, 3+4j] + [1-2j, 3-4j] = [2+0j, 6+0j]
+    expected_128 = np.array([2+0j, 6+0j], dtype=np.complex128)
+    self.assertAllClose(result_128, expected_128)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testComplexVariableAssignAddCPU(self):
+    """Test complex Variable assign_add on CPU."""
+    # Test with complex64
+    input_data_64 = constant_op.constant([1 + 2j, 3 + 4j], dtype=dtypes.complex64)
+    var_64 = resource_variable_ops.ResourceVariable(input_data_64, dtype=dtypes.complex64)
+    self.evaluate(var_64.initializer)
+    
+    delta_64 = constant_op.constant([0.5 - 1j, 1 + 0.5j], dtype=dtypes.complex64)
+    assign_add_op_64 = var_64.assign_add(delta_64)
+    result_64 = self.evaluate(assign_add_op_64)
+    
+    # Expected: [1+2j, 3+4j] + [0.5-1j, 1+0.5j] = [1.5+1j, 4+4.5j]
+    expected_64 = np.array([1.5+1j, 4+4.5j], dtype=np.complex64)
+    self.assertAllClose(result_64, expected_64)
+    
+    # Test with complex128
+    input_data_128 = constant_op.constant([1 + 2j, 3 + 4j], dtype=dtypes.complex128)
+    var_128 = resource_variable_ops.ResourceVariable(input_data_128, dtype=dtypes.complex128)
+    self.evaluate(var_128.initializer)
+    
+    delta_128 = constant_op.constant([0.5 - 1j, 1 + 0.5j], dtype=dtypes.complex128)
+    assign_add_op_128 = var_128.assign_add(delta_128)
+    result_128 = self.evaluate(assign_add_op_128)
+    
+    # Expected: [1+2j, 3+4j] + [0.5-1j, 1+0.5j] = [1.5+1j, 4+4.5j]
+    expected_128 = np.array([1.5+1j, 4+4.5j], dtype=np.complex128)
+    self.assertAllClose(result_128, expected_128)
+
 if __name__ == "__main__":
   test.main()

test_issue_105367.py

@@ -0,0 +1,13 @@
+import tensorflow as tf
+
+print("TensorFlow version:", tf.__version__)
+print("Testing complex Variable operations with tf.raw_ops.Conj and assign_add...")
+
+try:
+    input_data = tf.constant([1 + 2j, 3 + 4j], dtype=tf.complex64)
+    var = tf.Variable(input_data, dtype=tf.complex64)
+    conj_result = tf.raw_ops.Conj(input=input_data)
+    assign_add_op = var.assign_add(conj_result)
+    print("Success! Result:", assign_add_op.numpy())
+except Exception as e:
+    print(f"Error occurred: {type(e).__name__}: {e}")