Swap operands of dot if the LHS is fed by a parameter

PiperOrigin-RevId: 642090766

Author: Google-ML-Automation

tests/shape_poly_test.py

@@ -3327,6 +3327,10 @@ def test_harness(self, harness: PolyHarness):
     if "random_gamma" in harness.group_name:
       config_flags = {**config_flags, "jax_debug_key_reuse": False}
 
+    # TPU precision is a little lower since we swap the order of matmul operands.
+    if "cholesky" in harness.group_name and jtu.test_device_matches(["tpu"]):
+      harness.tol = 5e-5
+
     with jtu.global_config_context(**config_flags):
       harness.run_test(self)