[TPU][Pallas] relax tolerances and fix Pallas autotuning OOM in layer_norm

[yarongmu-google]

1. Fixes bfloat16 accuracy validation: Relaxes the rtol and atol from 1e-3 to 1e-2 in examples/layer_norm.py to prevent spurious validation failures caused by bfloat16 precision variations on TPUs.
2. Fixes VMEM Out-of-Memory during autotuning: The autotuner's default baseline config generation (block_sizes=[32, 128]) was attempting to allocate overly large chunks of the 10,240-element feature dimension into VMEM, causing the process to crash before tuning could even begin. Bypassed this by providing autotune_baseline_fn for both the forward and backward kernels.
3. Adds a safe fallback config: Added config=helion.Config(block_sizes=[32, 1024]) to the backward pass kernel to ensure a memory-safe configuration is available even if the search space fails to converge.