def crop(example: dict[str, tf.Tensor], TypeError: 'type' object is not subscriptable

):
"""
Crops the image tensor in the input dictionary to a random size and aspect ratio.

Args:
- example: A dictionary containing the image tensor to be cropped.

Returns:
- A dictionary containing the cropped image tensor.
"""
image = example['image']
shape = tf.shape(image)
height, width = shape[0], shape[1]
aspect_ratio = tf.cast(width, tf.float32) / tf.cast(height, tf.float32)

# Define the range of possible aspect ratios and choose one randomly
min_aspect_ratio = tf.constant(0.8, dtype=tf.float32)
max_aspect_ratio = tf.constant(1.2, dtype=tf.float32)
aspect_ratio_range = tf.linspace(min_aspect_ratio, max_aspect_ratio, 100)
random_aspect_ratio = tf.random.shuffle(aspect_ratio_range)[0]

# Define the range of possible heights and choose one randomly
max_height = tf.cast(tf.math.ceil(tf.cast(width, tf.float32) / random_aspect_ratio), tf.int32)
height_range = tf.linspace(tf.cast(tf.math.floor(tf.cast(height, tf.float32) * 0.8), tf.int32), max_height, 100)
random_height = tf.random.shuffle(height_range)[0]

# Define the range of possible widths based on the chosen height and aspect ratio
random_width = tf.cast(tf.math.ceil(tf.cast(random_height, tf.float32) * random_aspect_ratio), tf.int32)

# Crop the image tensor to the randomly chosen size
cropped_image = tf.image.random_crop(image, [random_height, random_width, 3])

# Update the dictionary with the cropped image tensor
example['image'] = cropped_image
return example