Makes CreateTpuFeeds() work for non-2D inputs for use_partitioned_infeed_queue=True.
Makes CreateTpuFeeds() work for non-2D inputs for use_partitioned_infeed_queue=True.

PiperOrigin-RevId: 312475129

tensorflow/lingvo,tensorflow/lingvo,tensorflow/lingvo,tensorflow/lingvo