def get_desired_batch_size(cls) -> Optional[int]:
    """Retrieves a user set fixed batch size, None if not set."""
    state = cls._get_topmost_state_frame()
    if state.desired_batch_size is not None:
      tf.compat.v1.logging.info('Using fixed batch size: %d',
                                state.desired_batch_size)
      return state.desired_batch_size
    return None
Retrieves a user set fixed batch size, None if not set.
tensorflow_transform/beam/context.py