def METHOD_NAME(
    image: TensorLike, bins: int = 256, name: Optional[str] = None
) -> tf.Tensor:
    """Equalize image(s)
    Args:
      images: A tensor of shape
          `(num_images, num_rows, num_columns, num_channels)` (NHWC), or
          `(num_rows, num_columns, num_channels)` (HWC), or
          `(num_rows, num_columns)` (HW). The rank must be statically known (the
          shape is not `TensorShape(None)`).
      bins: The number of bins in the histogram.
      name: The name of the op.
    Returns:
      Image(s) with the same type and shape as `images`, equalized.
    """
    with tf.name_scope(name or "equalize"):
        image_dims = tf.rank(image)
        image = to_4D_image(image)
        fn = partial(_equalize_image, bins=bins)
        image = tf.map_fn(fn, image)
        return from_4D_image(image, image_dims)