keras/src/backend/tensorflow/numpy.py/_keepdims
def _keepdims(x, y, axis):
    if axis is None:
        shape = [1 for _ in range(len(x.shape))]
    else:
        shape = list(shape_op(x))
        for axis in tree.flatten(axis):
            shape[axis] = 1
    y = tf.reshape(y, shape)
    return y