Modify `augment.transform` to use `transforms.shape.rank` instead of `tf.rank`.

PiperOrigin-RevId: 313435294

official/vision/image_classification/augment.py

@@ -169,7 +169,7 @@ def transform(image: tf.Tensor, transforms) -> tf.Tensor:
   """Prepares input data for `image_ops.transform`."""
   original_ndims = tf.rank(image)
   transforms = tf.convert_to_tensor(transforms, dtype=tf.float32)
-  if tf.rank(transforms) == 1:
+  if transforms.shape.rank == 1:
     transforms = transforms[None]
   image = to_4d(image)
   image = image_ops.transform(