keras/src/backend/tensorflow/numpy.py/trace
def trace(x, offset=0, axis1=0, axis2=1):
    x = convert_to_tensor(x)
    dtype = standardize_dtype(x.dtype)
    if dtype not in ("int64", "uint32", "uint64"):
        dtype = dtypes.result_type(dtype, "int32")
    x_shape = tf.shape(x)
    x = moveaxis(x, (axis1, axis2), (-2, -1))
    # Mask out the diagonal and reduce.
    x = tf.where(
        eye(x_shape[axis1], x_shape[axis2], k=offset, dtype="bool"),
        x,
        tf.zeros_like(x),
    )
    # The output dtype is set to "int32" if the input dtype is "bool"
    if standardize_dtype(x.dtype) == "bool":
        x = tf.cast(x, "int32")
    return tf.cast(tf.reduce_sum(x, axis=(-2, -1)), dtype)