[XLA:CPU] Relax test tolerances for tests using XLA:CPU.
[XLA:CPU] Relax test tolerances for tests using XLA:CPU.

An upcoming change to XLA:CPU will disable reassociation on floating point operators by default which is an unsound fast math optimization. This change is being made to fix numerical errors in softmax computations caused by reassocation. After that change, we will enable reassociation only in reduction operators where it is very important for performance and the XLA operator contract allows that.

Since this change alters the order of operations, it may cause small numerical changes leading to test failures. This change relaxes test tolerances to make tests pass.

PiperOrigin-RevId: 431475728

google/fedjax,google/fedjax