[Issue]: MIGraphX incorrectly determines output tensors' strides

Problem description

MIGraphX incorrectly determines output tensors' strides when running the attached model.
The model will be correctly compiled with static shapes, but when the compiled model is runned, the following exception is raised:

Traceback (most recent call last):
  File "issue.py", line 126, in <module>
    model.run_async(arguments, stream.cuda_stream, "ihipStream_t")
    ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: src/program.cpp:539: operator(): Incorrect shape {half_type, {4, 12, 64}, {768, 64, 1}} for parameter: main:#output_2 should be: half_type, {4, 12, 64}, {295680, 64, 1}

The exception doesn't make sense as we are supplying a tensor that has a correct stride for its shape.
The first dimension of the stride that the error is referring to (295680) can't be obtained even if all input shapes are multiplied together (3072).
Also, the exception seems to be flaky as it sometimes refers to different output variables (e.g. main:#output_1, main:#output_3, ...).

A curious observation is that if you change the BATCH_SIZE constant (from 4 to e.g. 8), the first stride mentioned in the error will stay the same invalid value (295680).

We have tested that version https://github.com/ROCm/AMDMIGraphX/commit/5c50d0cb1d2743e6c2ed063a17dd3f0e8d7c2f3f works and will output ALL OK when running the example script.
We have also tested that the version shipped with ROCm 10.0.0 doesn't work: https://github.com/ROCm/AMDMIGraphX/commit/becdb3da862f2297041b746b90bc6130e2b1d1f7
Also, the latest develop from a couple of days ago doesn't work: https://github.com/ROCm/AMDMIGraphX/commit/cee403e851ed2623342eef5371348ff7af932ce9
An identical error is produced by both non-working versions.

We haven't bisected to the exact commit as we can't do that easily automatically.

Steps to reproduce

import migraphx
import torch

DEVICE = "cuda:0"
BATCH_SIZE = 4

torch.cuda.set_device(DEVICE)

model = migraphx.parse_onnx("shape_issue.onnx", map_input_dims={
    "inputs_embeds": [BATCH_SIZE, 1, 768],
    "inputs_embeds_position_mask": [BATCH_SIZE, 1, 385],

    "decoder_kv_mask": [BATCH_SIZE, 384],

    "decoder.0.key": [BATCH_SIZE, 384, 12, 64],
    "decoder.1.key": [BATCH_SIZE, 384, 12, 64],
    "decoder.2.key": [BATCH_SIZE, 384, 12, 64],
    "decoder.3.key": [BATCH_SIZE, 384, 12, 64],
    "decoder.4.key": [BATCH_SIZE, 384, 12, 64],
    "decoder.5.key": [BATCH_SIZE, 384, 12, 64],

    "decoder.0.value": [BATCH_SIZE, 384, 12, 64],
    "decoder.1.value": [BATCH_SIZE, 384, 12, 64],
    "decoder.2.value": [BATCH_SIZE, 384, 12, 64],
    "decoder.3.value": [BATCH_SIZE, 384, 12, 64],
    "decoder.4.value": [BATCH_SIZE, 384, 12, 64],
    "decoder.5.value": [BATCH_SIZE, 384, 12, 64],

    "encoder.0.key": [BATCH_SIZE, 591, 12, 64],
    "encoder.1.key": [BATCH_SIZE, 591, 12, 64],
    "encoder.2.key": [BATCH_SIZE, 591, 12, 64],
    "encoder.3.key": [BATCH_SIZE, 591, 12, 64],
    "encoder.4.key": [BATCH_SIZE, 591, 12, 64],
    "encoder.5.key": [BATCH_SIZE, 591, 12, 64],

    "encoder.0.value": [BATCH_SIZE, 591, 12, 64],
    "encoder.1.value": [BATCH_SIZE, 591, 12, 64],
    "encoder.2.value": [BATCH_SIZE, 591, 12, 64],
    "encoder.3.value": [BATCH_SIZE, 591, 12, 64],
    "encoder.4.value": [BATCH_SIZE, 591, 12, 64],
    "encoder.5.value": [BATCH_SIZE, 591, 12, 64],
})

model.compile(migraphx.get_target("gpu"), offload_copy=False)

input_bindings = {
    "inputs_embeds": torch.zeros((BATCH_SIZE, 1, 768), dtype=torch.float16, device=DEVICE),
    "inputs_embeds_position_mask": torch.zeros((BATCH_SIZE, 1, 385), dtype=torch.float16, device=DEVICE),

    "decoder_kv_mask": torch.zeros((BATCH_SIZE, 384), dtype=torch.float16, device=DEVICE),

    "decoder.0.key": torch.zeros((BATCH_SIZE, 384, 12, 64), dtype=torch.float16, device=DEVICE),
    "decoder.0.value": torch.zeros((BATCH_SIZE, 384, 12, 64), dtype=torch.float16, device=DEVICE),
    "decoder.1.key": torch.zeros((BATCH_SIZE, 384, 12, 64), dtype=torch.float16, device=DEVICE),
    "decoder.1.value": torch.zeros((BATCH_SIZE, 384, 12, 64), dtype=torch.float16, device=DEVICE),
    "decoder.2.key": torch.zeros((BATCH_SIZE, 384, 12, 64), dtype=torch.float16, device=DEVICE),
    "decoder.2.value": torch.zeros((BATCH_SIZE, 384, 12, 64), dtype=torch.float16, device=DEVICE),
    "decoder.3.key": torch.zeros((BATCH_SIZE, 384, 12, 64), dtype=torch.float16, device=DEVICE),
    "decoder.3.value": torch.zeros((BATCH_SIZE, 384, 12, 64), dtype=torch.float16, device=DEVICE),
    "decoder.4.key": torch.zeros((BATCH_SIZE, 384, 12, 64), dtype=torch.float16, device=DEVICE),
    "decoder.4.value": torch.zeros((BATCH_SIZE, 384, 12, 64), dtype=torch.float16, device=DEVICE),
    "decoder.5.key": torch.zeros((BATCH_SIZE, 384, 12, 64), dtype=torch.float16, device=DEVICE),
    "decoder.5.value": torch.zeros((BATCH_SIZE, 384, 12, 64), dtype=torch.float16, device=DEVICE),

    "encoder.0.key": torch.zeros((BATCH_SIZE, 591, 12, 64), dtype=torch.float16, device=DEVICE),
    "encoder.0.value": torch.zeros((BATCH_SIZE, 591, 12, 64), dtype=torch.float16, device=DEVICE),
    "encoder.1.key": torch.zeros((BATCH_SIZE, 591, 12, 64), dtype=torch.float16, device=DEVICE),
    "encoder.1.value": torch.zeros((BATCH_SIZE, 591, 12, 64), dtype=torch.float16, device=DEVICE),
    "encoder.2.key": torch.zeros((BATCH_SIZE, 591, 12, 64), dtype=torch.float16, device=DEVICE),
    "encoder.2.value": torch.zeros((BATCH_SIZE, 591, 12, 64), dtype=torch.float16, device=DEVICE),
    "encoder.3.key": torch.zeros((BATCH_SIZE, 591, 12, 64), dtype=torch.float16, device=DEVICE),
    "encoder.3.value": torch.zeros((BATCH_SIZE, 591, 12, 64), dtype=torch.float16, device=DEVICE),
    "encoder.4.key": torch.zeros((BATCH_SIZE, 591, 12, 64), dtype=torch.float16, device=DEVICE),
    "encoder.4.value": torch.zeros((BATCH_SIZE, 591, 12, 64), dtype=torch.float16, device=DEVICE),
    "encoder.5.key": torch.zeros((BATCH_SIZE, 591, 12, 64), dtype=torch.float16, device=DEVICE),
    "encoder.5.value": torch.zeros((BATCH_SIZE, 591, 12, 64), dtype=torch.float16, device=DEVICE),
}

output_bindings = [
    ("logits", torch.zeros((BATCH_SIZE, 1, 51289), dtype=torch.float16, device=DEVICE)),

    ("next.decoder.0.key", torch.zeros((BATCH_SIZE, 12, 64), dtype=torch.float16, device=DEVICE)),
    ("next.decoder.0.value", torch.zeros((BATCH_SIZE, 12, 64), dtype=torch.float16, device=DEVICE)),
    ("next.decoder.1.key", torch.zeros((BATCH_SIZE, 12, 64), dtype=torch.float16, device=DEVICE)),
    ("next.decoder.1.value", torch.zeros((BATCH_SIZE, 12, 64), dtype=torch.float16, device=DEVICE)),
    ("next.decoder.2.key", torch.zeros((BATCH_SIZE, 12, 64), dtype=torch.float16, device=DEVICE)),
    ("next.decoder.2.value", torch.zeros((BATCH_SIZE, 12, 64), dtype=torch.float16, device=DEVICE)),
    ("next.decoder.3.key", torch.zeros((BATCH_SIZE, 12, 64), dtype=torch.float16, device=DEVICE)),
    ("next.decoder.3.value", torch.zeros((BATCH_SIZE, 12, 64), dtype=torch.float16, device=DEVICE)),
    ("next.decoder.4.key", torch.zeros((BATCH_SIZE, 12, 64), dtype=torch.float16, device=DEVICE)),
    ("next.decoder.4.value", torch.zeros((BATCH_SIZE, 12, 64), dtype=torch.float16, device=DEVICE)),
    ("next.decoder.5.key", torch.zeros((BATCH_SIZE, 12, 64), dtype=torch.float16, device=DEVICE)),
    ("next.decoder.5.value", torch.zeros((BATCH_SIZE, 12, 64), dtype=torch.float16, device=DEVICE)),
]

torch.cuda.synchronize()


def convert_tensor_to_argument(tensor):
    assert tensor.is_contiguous()
    return migraphx.argument_from_pointer(
        migraphx.shape(
            type="half_type",
            lens=list(tensor.size()),
            strides=list(tensor.stride()),
        ), tensor.data_ptr())


stream = torch.cuda.default_stream()

arguments = {}
for key, value in input_bindings.items():
    arguments[key] = convert_tensor_to_argument(value)
# NOTE: For some reason outputs can only be binded with a "magic" name, you
# can't use their actual names from the ONNX model.
for i, (key, value) in enumerate(output_bindings):
    # NOTE: For some reason outputs with a high enough number are left-padded
    # with zeroes and preceded with a colon...
    if i < 10:
        parameter_name = str(i)
    else:
        parameter_name = f":{i:05d}"
    arguments[f"main:#output_{parameter_name}"] = convert_tensor_to_argument(value)
# NOTE: Taken from:
# AMDMIGraphX/examples/diffusion/python_flux/models.py
# AMDMIGraphX/examples/diffusion/python_stable_diffusion_xl/txt2img.py
model.run_async(arguments, stream.cuda_stream, "ihipStream_t")
stream.synchronize()

print("ALL OK!")

Environment

OS: Debian GNU/Linux 13 (bookworm)
CPU: AMD Ryzen 9 9950X
GPU: AMD Radeon AI PRO R9700
ROCm version: 10.0.0

Model

https://drive.google.com/file/d/17XnXlj07cK8AtRRYHrIEw-ZfRYxc8C92/view?usp=drive_link