import tensorflow as tf

var = tf.Variable([1.0, 2.0, 3.0])
accum = tf.Variable([0.1, 0.2, 0.3], dtype=tf.complex64) 
accum2 = tf.Variable([0, 1.2, -3.3], dtype=tf.complex64) 
lr = 0.01
grad = tf.constant([0.1, 0.2, 0.3])
momentum = 0.9

# crash
tf.raw_ops.ResourceSparseApplyAdagrad(
    var=var.handle,
    accum=accum.handle,
    lr=lr,
    grad=grad,
    indices=tf.constant([1,2,3]),
    use_locking=False,
)


# crash only when gpu is available 
# tf.raw_ops.ResourceSparseApplyAdagradDA(var=var.handle,
#                     gradient_accumulator=accum.handle,
#                     gradient_squared_accumulator=accum2.handle,
#                     grad=grad, indices=tf.constant([1,2,3]), lr=lr,
#                     l1=lr, l2=lr, global_step=1000,
#                     use_locking=False)

# crash
# tf.raw_ops.ResourceSparseApplyAdagradV2(var=var.handle,
#                     accum=accum.handle,
#                     epsilon = 0.9,
#                     grad=grad, indices=tf.constant([1,2,3]), lr=lr,
#                     use_locking=False,update_slots=False)

2024-06-06 01:14:57.455749: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2024-06-06 01:14:57.489757: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-06-06 01:14:58.552846: F tensorflow/core/framework/tensor.cc:844] Check failed: dtype() == expected_dtype (8 vs. 1) float expected, got complex64
Aborted (core dumped)