@LPanosTT, can you please take a look at this?

@mrakitaTT, we're getting closer to running more multi-device tests and we'd like to leverage Jax to do so. Do you know if this is still an issue?

@nsmithtt no one has worked on this yet and we still have a hardcoded int device_info_count_ = 1; hack in order to be able to use n300 devices to run tests. The last tech sync when Wooseok was presenting the multi device progress I was wondering how is he getting the hlo graphs and if he is using tt-xla for his work, but haven't yet gotten to asking him.

@ajakovljevicTT this might be a good PJRT task for you or Sanja, can you please investigate what fails when you disable the device count hack and how device is selected in PJRT etc?

I think we got multi device stableHLO IR from a multi-device google collab and have been working directly from the stablehlo. Or @wooseokTT did you hack PJRT to make it work for a one off test?

Thanks for looking into it @ajakovljevicTT!

@nsmithtt @mrakitaTT I have local fix that uses device ids instead of the hardcoded "1". Then, I used matmul python code to generate stableHLO mlir using tt-xla.

@wooseokTT Would you mind sharing that fix in a branch? Thanks!

@ajakovljevicTT It's not a big change. I just fixed code to the point where I can run multidevice, and didn't test any other things.

auto [system_desc, chip_ids] = tt::runtime::getCurrentSystemDesc();
#if TEST_MESH
int device_info_count_ = chip_ids.size();
#else
int device_info_count_ = 1;
#endif

auto [system_desc, chip_ids] = tt::runtime::getCurrentSystemDesc();
#if TEST_MESH
auto device = tt::runtime::openDevice(chip_ids);
#else
int dev_0 = chip_ids[0];
auto device = tt::runtime::openDevice({dev_0});
#endif