[Perf] Optimize simple_gla CPU reference test execution time

[0xaskr]

Summary

The CPU reference test in tests/ref/simple_gla currently takes too long to complete, impacting CI efficiency and development velocity.

Type

• Performance regression (was faster before)
• Below expected performance target (not meeting 80% roofline)
• Optimization opportunity

Kernel / Operation

tests/ref/simple_gla — CPU reference test suite for Simple GLA.

Observed Performance

Test execution time is excessively long (exact numbers TBD via profiling).

Expected Performance

Significantly reduced test duration while maintaining correctness.

Environment

• Python version:
• JAX version:
• Hardware: CPU
• OS:

Reproduction

uv run pytest tests/ref/simple_gla -v

Tasks

• Profile the CPU reference test to determine major time-consuming sections
• Propose and implement optimizations to reduce total execution time
• Validate correctness of results post optimization
• Document analysis and improvements

Additional Context

Faster reference testing will improve CI efficiency and support development velocity.

[0xaskr]

@