Error when using Tensorflow bucket_by_sequence_length() and tf.py_function() together - tensorflow

Question

I have the following sample code:
import glob
import random
import tensorflow as tf
import cv2
def random_blur(image):
# do stuff which can't be done with a tf.image function...
random_x_blur = random.randint(1, 3)
random_y_blur = random.randint(1, 3)
return cv2.blur(image, (random_x_blur, random_y_blur))
def transform(image):
image = tf.image.random_jpeg_quality(image, 75, 100)
image = tf.image.random_brightness(image, 0.5)
image = tf.image.random_contrast(image, 0.2, 0.5)
# here's the problem...
# image = tf.py_function(func=random_blur, inp=[image], Tout=tf.uint8)
return image
def process_path(file_path):
image = tf.io.read_file(file_path)
image = tf.image.decode_png(image, channels=1)
return transform(image), image
train_directory = 'data/small/'
train_files = glob.glob(train_directory + '*.png')
ds_train = tf.data.Dataset.from_tensor_slices(train_files)
boundaries = [100, 200, 300, 400]
batch_sizes = [16, 16, 16, 16, 16]
ds_train = ds_train.map(process_path, 4)
ds_train = ds_train.bucket_by_sequence_length(element_length_func=lambda x, y: tf.shape(x)[1],
bucket_boundaries=boundaries,
bucket_batch_sizes=batch_sizes)
I'm trying to create a Tensorflow Dataset from variable-width 60px-high images in a directory, using the bucket_by_sequence_length() function to ensure the images in each minibatch have the same dimensions. This all works fine until I uncomment the line beneath "here's the problem" in the code above. When you uncomment that and run it, it produces the following error:
Traceback (most recent call last):
File "test.py", line 34, in <module>
ds_train = ds_train.bucket_by_sequence_length(element_length_func=lambda x, y: tf.shape(x)[1],
File "/Users/garnet/Library/Python/3.8/lib/python/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 3120, in bucket_by_sequence_length
return self.group_by_window(
File "/Users/garnet/Library/Python/3.8/lib/python/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 2976, in group_by_window
return _GroupByWindowDataset(
File "/Users/garnet/Library/Python/3.8/lib/python/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 5841, in __init__
self._make_reduce_func(reduce_func, input_dataset)
File "/Users/garnet/Library/Python/3.8/lib/python/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 5890, in _make_reduce_func
self._reduce_func = structured_function.StructuredFunctionWrapper(
File "/Users/garnet/Library/Python/3.8/lib/python/site-packages/tensorflow/python/data/ops/structured_function.py", line 271, in __init__
self._function = fn_factory()
File "/Users/garnet/Library/Python/3.8/lib/python/site-packages/tensorflow/python/eager/function.py", line 2610, in get_concrete_function
graph_function = self._get_concrete_function_garbage_collected(
File "/Users/garnet/Library/Python/3.8/lib/python/site-packages/tensorflow/python/eager/function.py", line 2576, in _get_concrete_function_garbage_collected
graph_function, _ = self._maybe_define_function(args, kwargs)
File "/Users/garnet/Library/Python/3.8/lib/python/site-packages/tensorflow/python/eager/function.py", line 2760, in _maybe_define_function
graph_function = self._create_graph_function(args, kwargs)
File "/Users/garnet/Library/Python/3.8/lib/python/site-packages/tensorflow/python/eager/function.py", line 2670, in _create_graph_function
func_graph_module.func_graph_from_py_func(
File "/Users/garnet/Library/Python/3.8/lib/python/site-packages/tensorflow/python/framework/func_graph.py", line 1247, in func_graph_from_py_func
func_outputs = python_func(*func_args, **func_kwargs)
File "/Users/garnet/Library/Python/3.8/lib/python/site-packages/tensorflow/python/data/ops/structured_function.py", line 248, in wrapped_fn
ret = wrapper_helper(*args)
File "/Users/garnet/Library/Python/3.8/lib/python/site-packages/tensorflow/python/data/ops/structured_function.py", line 177, in wrapper_helper
ret = autograph.tf_convert(self._func, ag_ctx)(*nested_args)
File "/Users/garnet/Library/Python/3.8/lib/python/site-packages/tensorflow/python/autograph/impl/api.py", line 689, in wrapper
return converted_call(f, args, kwargs, options=options)
File "/Users/garnet/Library/Python/3.8/lib/python/site-packages/tensorflow/python/autograph/impl/api.py", line 377, in converted_call
return _call_unconverted(f, args, kwargs, options)
File "/Users/garnet/Library/Python/3.8/lib/python/site-packages/tensorflow/python/autograph/impl/api.py", line 458, in _call_unconverted
return f(*args, **kwargs)
File "/Users/garnet/Library/Python/3.8/lib/python/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 3111, in batching_fn
shapes = make_padded_shapes(
File "/Users/garnet/Library/Python/3.8/lib/python/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 3083, in make_padded_shapes
shape = [
File "/Users/garnet/Library/Python/3.8/lib/python/site-packages/tensorflow/python/framework/tensor_shape.py", line 882, in __iter__
raise ValueError("Cannot iterate over a shape with unknown rank.")
ValueError: Cannot iterate over a shape with unknown rank.
Likewise, my code works fine if I uncomment that line but remove the call to bucket_by_sequence_length() and limit my training data to images with identical dimensions.
It seems that bucket_by_sequence_length() and tf.py_function() don't play nice together, even with eager mode enabled. I need to do some image augmentations/transformations which the standard tf.image functions don't provide. Any ideas?