[0z5a]

Global relax.median(..., axis=None, keepdims=True) flattens its input for sorting, then currently returns the rank-one intermediate. For example, a (2, 3) input is inferred to return (1, 1), but LegalizeOps changes the function's return shape to (1,).

Reshape the result to one dimension per input dimension, including the scalar case. Add regression coverage for the inferred return type and LLVM runtime values and shapes.

Fixes #20915.

Testing

• On GH200 (CPU/LLVM), the original implementation returned shape (1,) for (2, 3) and (2, 3, 4) inputs; the change returned (1, 1) and (1, 1, 1) with the expected lower-median values.
• pytest for the new regression file, test_transform_legalize_ops_search_statistical.py, and test_op_statistical.py: 62 passed.
• ruff check and ruff format --check on the changed files: passed.