Fixes pickling error in saved_transform_io_v2.py with Python 3.11.

PiperOrigin-RevId: 610903193

Author: iindyk

tensorflow_transform/py_func/pyfunc_helper.py

@@ -15,6 +15,7 @@
 
 import dill
 import tensorflow as tf
+from tfx_bsl import beam as tfx_bsl_beam
 # TODO(b/243513856): Switch to `collections.namedtuple` or `typing.NamedTuple`
 # once the Spark issue is resolved.
 from tfx_bsl.types import tfx_namedtuple
@@ -26,6 +27,8 @@
 
 _PYFUNC_COLLECTION_KEY = 'pyfuncs'
 
+tfx_bsl_beam.fix_code_type_pickling()
+
 
 class _PyFuncDef(tfx_namedtuple.namedtuple('_PyFuncDef', ['token', 'func'])):
   """An internal wrapper around tuple(token, func).