[kfac-jax] Update _Linear layer to support the new "algorithm" tuning…

… parameters for dot_general that will be included in the next JAX release.

PiperOrigin-RevId: 674292750

tests/models.py

@@ -92,12 +92,19 @@ def __call__(self, inputs: LayerInputs, *_) -> LayerInputs:  # pytype: disable=s
       else:
         assert all(p.dtype == y.dtype for p in params if p is not None)
         preferred_element_type = y.dtype
+      # The dot algorithm and transpose_algorithm parameters are available (and
+      # required) after JAX version 0.4.33.
+      if jax_version <= (0, 4, 33):
+        algorithm_kwargs = {}
+      else:
+        algorithm_kwargs = dict(algorithm=None, transpose_algorithm=None)
 
       y = tags.register_dense(
           y, x, *params,
           dimension_numbers=(((1,), (0,)), ((), ())),
           precision=(jax.lax.Precision.HIGHEST, jax.lax.Precision.HIGHEST),
           preferred_element_type=preferred_element_type,
+          **algorithm_kwargs,
       )
     layer_values.append((x, y))