Average and standard deviation pooling that ignores zero values, in tensorflow

Question

I am trying to apply average, and standard deviation pooling, to each dimension across each vector in a batch, while ignoring any zero values.

Below is a minimum viable reproducible example of my code:

import tensorflow as tf

av_pool = tf.keras.layers.Lambda(
    lambda z: tf.math.reduce_mean(z, axis=1)
)

sd_pool = tf.keras.layers.Lambda(
    lambda z: tf.math.reduce_std(z, axis=1, keepdims=False)
)

# Batch size = 2, sequence length = 4, vector size = 3
a = tf.convert_to_tensor(
    [[[1., 1., 1.], [1., 1., 1.], [1., 0., 0.], [0., 0., 0.]], 
     [[1., 1., 1.], [0., 0., 0.], [0., 0., 0.], [0., 0., 0.]]],
    dtype=tf.float32
)

mask = tf.not_equal(a, tf.zeros_like(a))        

non_zero = tf.ragged.boolean_mask(a, mask)  

# average pooling works
averages = av_pool(non_zero)

# however standard deviation pooling does not
standard_deviations = sd_pool(non_zero)

# producing the following error:
# ValueError: keepdims=True is not supported for RaggedTensors.

Please may somebody help me understand what is going wrong? Why does the average pooling work while the standard deviation pooling does not?

Thank you.