Add support for `int32_t` indices in TBE training (2D/N) (#3374) #515

Job	Run time
build-bazel	0s
build-linux	0s
build-macos (macos-13, static)	16m 21s
build-windows (windows-2022, static)	12m 19s
build-macos (macos-13, shared)	14m 39s
build-windows (windows-2022, shared)	12m 46s
	56m 5s

Provide feedback