Jax input tensor test variations

[nsmithtt]

Integrate the tests from Tap's repo into tt-xla. Feel free to update this issue as it makes sense.

As I understand it there are currently 2 flavors of inputs:

• Manual sharding (annotated graph) + no input sharding
• Manual sharding (annotated graph) + input sharding via device put & mesh attribute

[nsmithtt]

We should mark the unsupported flavors as just unsupported for now. I've filed a follow on issue for supporting the second bullet:

• Support automatic sharded tensors in runtime #210

[tapspatel]

I added @ajakovljevicTT to my repo and shared all the test cases. Please let me know if you need anything else from my end. As well, I iterate through many different meshes, but we will only be supporting the following in runtime

N300: 1x2 not supported: 2x1​
T3K: 1x8, 2x4 not supported: 4x2, 8x1​
TG: 1x32, 8x4 not supported: 2x16, 4x8, 6x2, 32x1