Fix axlearn CI failure due to pyarrow version upgrade (#1311)

Author: raytomato

pyproject.toml

@@ -30,6 +30,7 @@ core = [
     "nltk==3.7",  # for text preprocessing
     "optax==0.1.7",  # optimizers (0.1.0 has known bugs).
     "portpicker",
+    "pyarrow>=20.0.0,<21.0.0",  # Pin to v20.x to avoid PyExtensionType -> ExtensionType breaking change in v21
     "protobuf>=3.20.3",
     "tensorboard-plugin-profile==2.15.1",
     # This has both x86 and arm64 wheels. Underneath the hood it uses tensorflow-macos since 2.13.
@@ -125,6 +126,7 @@ vertexai_tensorboard = [
 ]
 # Dataflow dependencies.
 dataflow = [
+    "pyarrow>=20.0.0,<21.0.0",  # Pin to v20.x to avoid PyExtensionType -> ExtensionType breaking change in v21
     "apache-beam==2.55.1",
     "apache-beam[gcp]",
     "google-apitools",  # for beam pipeline