Re-implement tensor shape/rank operators using TFE runtime APIs.

[eaplatanios]

@rxwei Ported from swiftlang/swift#24949.

Gives me about 98% speedup for Tensor.rank and 85% speedup for Tensor.shape (for a 2x3 tensor) on a MacBook Pro and all tests pass.

[rxwei]

Thanks Anthony! In my earlier experiment it crashed on macOS only. Glad to hear you didn't have this issue, but I'll build and verify this locally.

[bgogul]

Nice @eaplatanios ! One thing to keep in mind is that using eager API directly will affect the performance of Lazy Tensor as it triggers materialization of symbolic tensors and break traces into smaller ones.

[rxwei]

Nice @eaplatanios ! One thing to keep in mind is that using eager API directly will affect the performance of Lazy Tensor as it triggers materialization of symbolic tensors and break traces into smaller ones.

Would it be possible for you to profile your lazy tensor implementation against changes made by this PR?

[eaplatanios]

@bgogul @rxwei I think that's fine because if you want to use the lazy tensor implementation you can still use rankTensor and shapeTensor.

[rxwei]

@bgogul @rxwei I think that's fine because if you want to use the lazy tensor implementation you can still use rankTensor and shapeTensor.

Ah yep, totally.

[eaplatanios]

I believe that to test this and all follow-up PRs we need to update the toolchain used by the CI server. Otherwise, we'll be getting failures related to TensorGroup.

[rxwei]

A new CI toolchain will be ready in an hour.

[rxwei]

All tests passed. Thank you!