keras/src/backend/tensorflow/numpy.py/average
def average(x, axis=None, weights=None):
    x = convert_to_tensor(x)

    if weights is None:  # Treat all weights as 1
        dtype = dtypes.result_type(x.dtype, float)
        x = tf.cast(x, dtype)
        avg = tf.reduce_mean(x, axis=axis)
    else:
        weights = convert_to_tensor(weights)
        dtype = dtypes.result_type(x.dtype, weights.dtype, float)
        x = tf.cast(x, dtype)
        weights = tf.cast(weights, dtype)

        def _rank_equal_case():
            weights_sum = tf.reduce_sum(weights, axis=axis)
            return tf.reduce_sum(x * weights, axis=axis) / weights_sum

        def _rank_not_equal_case():
            weights_sum = tf.reduce_sum(weights)
            axes = tf.convert_to_tensor([[axis], [0]])
            return tf.tensordot(x, weights, axes) / weights_sum

        if axis is None:
            avg = _rank_equal_case()
        else:
            if len(x.shape) == len(weights.shape):
                avg = _rank_equal_case()
            else:
                avg = _rank_not_equal_case()
    return avg