jax

490b6472 - Artificially limit SMEM capacity for tests

Artificially limit SMEM capacity for tests This removes a failure mode where new tests that pass on 10.x devices would fail on 12.y devices due to different SMEM capacity. Also adjust test cases to fit in 12.x SMEM limits.