xla

f53d7431 - Migrate to upstream tensor_util classes and functions (#3317)

Migrate to upstream tensor_util classes and functions (#3317) * Migrate to upstream tensor_util classes and functions * Fix linter issues * Remove PTXLA OptionalOr definition * Pull latest changes