code: def _call_tf_lowering(
      ctx: mlir.LoweringRuleContext,
      *args_op,
      platform,
      function_flat_tf,
      args_flat_sig_tf,
      has_side_effects,
      ordered,
      call_tf_graph,
      output_avals,
      **_,
  ):
      # We use the same TF lowering device as for the embedding JAX computation.
      # One example when this is needed is when the code refers to variables on one
      # device. Or, for sharding annotations (only supported on TPU).
  
      if platform in ["cpu", "tpu"]:
          tf_platform = platform.upper()
      elif platform == "cuda":
          tf_platform = "GPU"
      else:
          raise ValueError("platform {platform} not supported")
  
      concrete_function_flat_tf = _get_concrete_function_tf(
          function_flat_tf, args_flat_sig_tf
      )
  
      captured_inputs = []
      if concrete_function_flat_tf.captured_inputs:
          # The function uses either captured variables or tensors.
          msg = (
              "call_tf works best with a TensorFlow function that does not capture "
              "variables or tensors from the context. "
              "See https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#limitations-of-call_tf for a discussion. "
              f"The following captures were found {concrete_function_flat_tf.captured_inputs}"
          )
          logging.warning(msg)
          for inp in concrete_function_flat_tf.captured_inputs:
              if (
                  inp.dtype == tf.resource
              ):  # A variable; lookup by handle
                  inp_vars = [
                      v
                      for v in concrete_function_flat_tf.variables
                      if inp is v.handle
                  ]
                  assert len(inp_vars) == 1, f"Found {inp_vars}"
                  captured_inputs.append(inp_vars[0])
              else:
                  captured_inputs.append(inp)
  
      captured_ops = tuple(
          mlir.ir_constant(np.asarray(inp)) for inp in captured_inputs
      )
  
      if call_tf_graph:
          with jax2tf_internal.inside_call_tf():
              return emit_tf_embedded_graph_custom_call(
                  ctx,
                  concrete_function_flat_tf,
                  tuple(args_op) + captured_ops,
                  has_side_effects,
                  ordered,
                  output_avals,
              )
  
      def convert_to_spec(x):
          if isinstance(x, tf.TensorSpec):
              return x
          else:
              return tf.TensorSpec.from_tensor(x)
  
      args_tf_flat = [convert_to_spec(a) for a in args_flat_sig_tf]
  
      with jax2tf_internal.inside_call_tf():
          # When the TF computation uses variables on a particular device, we must
          # get_compiler_ir for that exact device.
          tf_device_name = f"/device:{tf_platform}:0"
          try:
              func_tf_hlo = (
                  function_flat_tf.experimental_get_compiler_ir(
                      *args_tf_flat
                  )(stage="hlo_serialized", device_name=tf_device_name)
              )
          except Exception as e:
              msg = (
                  "Error compiling TensorFlow function (see below for the caught exception)."
                  + "\ncall_tf can used "
                  + "in a staged context (under jax.jit, lax.scan, etc.) only with "
                  + "compilable functions with static output shapes.\n"
                  + "See https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#limitations-of-call_tf for a discussion."
                  + "\n\nCaught TensorFlow exception: "
                  + str(e)
              )
              raise ValueError(msg) from e
  
      xla_comp = xla_client.XlaComputation(func_tf_hlo)
  
      # Canonicalize the results; e.g., makes them x32 if JAX is in 32-bit mode
      def canonical_res_aval(
          res_shape: xla_client.Shape,
      ) -> core.ShapedArray:
          if not res_shape.is_static():
              msg = (
                  "Compiled TensorFlow function has dynamic output shape "
                  + f"{res_shape}. call_tf can used "
                  + "in a staged context (under jax.jit, lax.scan, etc.) only with "
                  + "compilable functions with static output shapes. "
                  + "See https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#limitations-of-call_tf for a discussion."
              )
              raise ValueError(msg)
  
          res_dtype = res_shape.numpy_dtype()
          jax_res_dtype = dtypes.canonicalize_dtype(res_dtype)
          return core.ShapedArray(res_shape.dimensions(), jax_res_dtype)
  
      result_shape = xla_comp.program_shape().result_shape()
      if not result_shape.is_tuple():
          # TF does not wrap singletons as tuples, but JAX expects tuples because
          # call_tf is a multiple_results primitive.
          result_shapes = (result_shape,)
      else:
          result_shapes = result_shape.tuple_shapes()  # type: ignore
  
      result_avals = tuple(map(canonical_res_aval, result_shapes))  # type: ignore
  
      submodule = mlir.xla_computation_to_mlir_module(xla_comp)
      symtab = ir.SymbolTable(submodule.operation)
      callee_result_types = symtab["main"].type.results
      fn = mlir.merge_mlir_modules(
          ctx.module_context.module,
          f"call_tf_{function_flat_tf.name}",
          submodule,
      )
      call = func_dialect.CallOp(
          callee_result_types,
          ir.FlatSymbolRefAttr.get(fn),
          tuple(args_op) + captured_ops,
      )
      if result_shape.is_tuple():
          flat_results = [
              hlo.get_tuple_element(call, mlir.i32_attr(i))
              for i in range(len(result_shapes))
          ]
      else:
          flat_results = call.results
  
      if ordered:
          raise NotImplementedError(
              "ordered=True is not supported in the jitted context without"
              " `call_tf_graph=True`"
          )
  
      outputs = []
      for op, res_aval, res_shape in zip(
          flat_results, result_avals, result_shapes
      ):
          if res_aval.dtype != res_shape.numpy_dtype():
              op = hlo.ConvertOp(
                  mlir.aval_to_ir_type(res_aval), op
              ).result
          outputs.append(op)
      return outputs
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import dataclasses
  import functools
  from typing import Any, Callable, Optional
  from absl import logging
  import jax
  from jax import dlpack
  from jax import dtypes
  from jax import numpy
  from jax import tree_util
  from jax._src import ad_util
  from jax._src import core
  from jax._src import effects
  from jax._src import util
  from jax._src.lib import xla_client
  from jax._src.lib.mlir import ir
  from jax._src.lib.mlir.dialects import func
  from jax._src.lib.mlir.dialects import hlo
  from jax.experimental.jax2tf import jax2tf
  from jax.interpreters import mlir
  import numpy
  import tensorflow
label_desc: Use image node_id 6 for calling a global function with example usage: _call_tf_lowering(ctx) and returns: outputs, core, emit_tf_embedded_graph_custom_call, x, tf