Fix jax2tf failure coming from dot_general

PiperOrigin-RevId: 688738110

Author: yashk2810

jax/experimental/jax2tf/impl_no_xla.py

@@ -364,6 +364,7 @@ def _conv_general_dilated(
 def _dot_general(lhs, rhs, *, dimension_numbers,
                  precision: tuple[PrecisionType, PrecisionType] | None,
                  preferred_element_type: DType | None,
+                 out_type=None,
                  _in_avals: Sequence[core.ShapedArray],
                  _out_aval: core.ShapedArray):
   """Implementation of lax.dot_general_p in terms of tf.linalg.einsum."""

jax/experimental/jax2tf/jax2tf.py

@@ -2180,9 +2180,10 @@ def gen_conv(lhs, rhs, preferred_element_type: DType | None):
 tf_impl_with_avals[lax.conv_general_dilated_p] = _conv_general_dilated
 
 
-def _dot_general(lhs, rhs, *, dimension_numbers, out_type,
+def _dot_general(lhs, rhs, *, dimension_numbers,
                  precision: lax_internal.CanonicalPrecision,
                  preferred_element_type: DType | None,
+                 out_type=None,
                  _in_avals: Sequence[core.ShapedArray],
                  _out_aval: core.ShapedArray):
   """Implementation of lax.dot_general_p in terms of tf.linalg.einsum."""