Fix jax2tf failure coming from dot_general

PiperOrigin-RevId: 688738110
jax-ml · Oct 22, 2024 · 4688da3 · 4688da3
1 parent f8a1f02
commit 4688da3
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/jax/experimental/jax2tf/impl_no_xla.py b/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."""

diff --git a/jax/experimental/jax2tf/jax2tf.py b/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."""