def _convert_jax_impl(jax_impl: Callable, *, multiple_results=True) -> Callable:
    'Convert the JAX implementation of a primitive.\n\n  Args:\n    jax_impl: typically the impl-rule for a primitive, with signature\n      `(*args: JaxVal, **kwargs) -> Sequence[JaxVal]`. This function implements\n        a primitive in terms of other primitives.\n    multiple_results: whether `jax_impl` returns a sequence of results.\n\n  Returns:\n     a function with signature `(*args: TfVal, _in_avals, _out_aval, **kwargs)\n     -> Sequence[TfVal]`.\n  '

    def wrapped(*tf_args: TfVal, _in_avals: Sequence[core.AbstractValue], _out_aval: core.AbstractValue, **kwargs) -> Sequence[TfVal]:

        def jax_impl_jax_args(*jax_args):
            jax_results = jax_impl(*jax_args, **kwargs)
            return (jax_results if multiple_results else [jax_results])
        tf_results_with_avals = _interpret_fun(lu.wrap_init(jax_impl_jax_args), tf_args, _in_avals)
        (tf_results, _) = util.unzip2(tf_results_with_avals)
        return (tf_results if multiple_results else tf_results[0])
    return wrapped
Convert the JAX implementation of a primitive.

Args:
  jax_impl: typically the impl-rule for a primitive, with signature
    `(*args: JaxVal, **kwargs) -> Sequence[JaxVal]`. This function implements
      a primitive in terms of other primitives.
  multiple_results: whether `jax_impl` returns a sequence of results.

Returns:
   a function with signature `(*args: TfVal, _in_avals, _out_aval, **kwargs)
   -> Sequence[TfVal]`.
def _convert_jax_impl(jax_impl: Callable, *, multiple_results=True) -> Callable:
    'Convert the JAX implementation of a primitive.\n\n  Args:\n    jax_impl: typically the impl-rule for a primitive, with signature\n      `(*args: JaxVal, **kwargs) -> Sequence[JaxVal]`. This function implements\n        a primitive in terms of other primitives.\n    multiple_results: whether `jax_impl` returns a sequence of results.\n\n  Returns:\n     a function with signature `(*args: TfVal, _in_avals, _out_aval, **kwargs)\n     -> Sequence[TfVal]`.\n  '

    def wrapped(*tf_args: TfVal, _in_avals: Sequence[core.AbstractValue], _out_aval: core.AbstractValue, **kwargs) -> Sequence[TfVal]:

        def jax_impl_jax_args(*jax_args):
            jax_results = jax_impl(*jax_args, **kwargs)
            return (jax_results if multiple_results else [jax_results])
        tf_results_with_avals = _interpret_fun(lu.wrap_init(jax_impl_jax_args), tf_args, _in_avals)
        (tf_results, _) = util.unzip2(tf_results_with_avals)
        return (tf_results if multiple_results else tf_results[0])
    return wrapped