jax

a5d92a06 - PR #34434: scaled_matmul_stablehlo_test: do not assert ordering within HLO

PR #34434: scaled_matmul_stablehlo_test: do not assert ordering within HLO Imported from GitHub PR https://github.com/jax-ml/jax/pull/34434 https://github.com/openxla/xla/pull/35348 had the effect of swapping the order in the HLO dump of the collectives and custom calls being regexed for. This relaxes the test logic to require that the collective and custom call patterns match separately and in either order. These tests require 4 or more Blackwell GPUs. Copybara import of the project: -- 4b3b049b6425e0c66e8c0107893c5b20a4e7f147 by Olli Lupton <olupton@nvidia.com>: scaled_matmul_stablehlo_test: do not assert ordering within HLO Merging this change closes #34434 COPYBARA_INTEGRATE_REVIEW=https://github.com/jax-ml/jax/pull/34434 from olupton:relax-ordering 4b3b049b6425e0c66e8c0107893c5b20a4e7f147 PiperOrigin-RevId: 857192546