code: def _call_exported_lowering(
      ctx: mlir.LoweringRuleContext, *args, exported: Exported
  ):
      if exported.uses_shape_polymorphism:
          ctx.module_context.shape_poly_state.uses_dim_vars = True
  
      axis_context = ctx.module_context.axis_context
      if isinstance(axis_context, sharding_impls.ShardingContext):
          num_devices = axis_context.num_devices
      elif isinstance(axis_context, sharding_impls.SPMDAxisContext):
          num_devices = axis_context.mesh.size
      else:
          raise NotImplementedError(type(axis_context))
      if num_devices != exported.nr_devices:
          raise NotImplementedError(
              f"Exported module {exported.fun_name} was lowered for "
              f"{exported.nr_devices} devices and is called in a context with "
              f"{num_devices} devices"
          )
  
      # Apply in_shardings
      args = tuple(
          wrap_with_sharding(ctx, x, x_aval, x_sharding)
          for x, x_aval, x_sharding in zip(
              args, ctx.avals_in, exported.in_shardings
          )
      )
      submodule = ir.Module.parse(exported.mlir_module())
      symtab = ir.SymbolTable(submodule.operation)
  
      # The called function may have been exported with polymorphic shapes and called
      # now with more refined shapes. We insert hlo.ConvertOp to ensure the module
      # is valid.
      def convert_shape(
          x: ir.Value,
          x_aval: core.AbstractValue,
          new_aval: core.AbstractValue,
      ) -> ir.Value:
          new_ir_type = mlir.aval_to_ir_type(new_aval)
          if x.type != new_ir_type:
              return hlo.convert(mlir.aval_to_ir_type(new_aval), x)
          else:
              return x
  
      callee_type = symtab["main"].type
      # TODO: maybe cache multiple calls
      fn = mlir.merge_mlir_modules(
          ctx.module_context.module,
          f"call_exported_{exported.fun_name}",
          submodule,
      )
  
      submodule_args = []
      # All the platforms for the current lowering must be among the platforms
      # for which the callee was lowered.
      lowering_platforms = ctx.module_context.platforms
  
      callee_lowering_platform_index: list[int] = []
      for platform in lowering_platforms:
          if platform in exported.lowering_platforms:
              callee_lowering_platform_index.append(
                  exported.lowering_platforms.index(platform)
              )
          elif (
              DisabledSafetyCheck.platform()
              in exported.disabled_safety_checks
          ):
              callee_lowering_platform_index.append(0)
          else:
              raise ValueError(
                  f"The exported function '{exported.fun_name}' was lowered for "
                  f"platforms '{exported.lowering_platforms}' but it is used "
                  f"on '{lowering_platforms}'."
              )
  
      if len(exported.lowering_platforms) > 1:
          # The exported module takes a platform index argument
          if len(lowering_platforms) > 1:
              current_platform_idx = ctx.dim_var_values[0]
          else:
              current_platform_idx = mlir.ir_constant(np.int32(0))
          # Compute the rule index based on the current platform
          i32_type = mlir.aval_to_ir_types(
              core.ShapedArray((), dtype=np.int32)
          )[0]
          if current_platform_idx.type != i32_type:
              current_platform_idx = hlo.ConvertOp(
                  i32_type, current_platform_idx
              )
          callee_platform_idx = hlo.CaseOp(
              [i32_type],
              index=current_platform_idx,
              num_branches=len(lowering_platforms),
          )
          for i in range(len(lowering_platforms)):
              branch = callee_platform_idx.regions[i].blocks.append()
              with ir.InsertionPoint(branch):
                  hlo.return_(
                      mlir.ir_constants(
                          np.int32(callee_lowering_platform_index[i])
                      )
                  )
          if callee_platform_idx.result.type != callee_type.inputs[0]:
              callee_platform_idx = hlo.ConvertOp(
                  callee_type.inputs[0], callee_platform_idx
              )
  
          submodule_args.append(callee_platform_idx)
      else:
          assert len(lowering_platforms) == 1
  
      if _keep_main_tokens(exported.mlir_module_serialization_version):
          ordered_effects = exported.ordered_effects
      else:
          ordered_effects = ()
      for eff in ordered_effects:
          token_in = ctx.tokens_in.get(eff)[0]
          submodule_args.append(token_in)
      kept_args = [
          convert_shape(a, a_aval, exported_in_aval)
          for i, (a, a_aval, exported_in_aval) in enumerate(
              zip(args, ctx.avals_in, exported.in_avals)
          )
          if i in exported.module_kept_var_idx
      ]
      submodule_args = submodule_args + kept_args
  
      call = func_dialect.CallOp(
          callee_type.results,
          ir.FlatSymbolRefAttr.get(fn),
          submodule_args,
      )
      if ordered_effects:
          tokens_out = {
              eff: (call.results[effect_idx],)
              for effect_idx, eff in enumerate(ordered_effects)
          }
          ctx.set_tokens_out(mlir.TokenSet(tokens_out))
      # The ctx.avals_out already contain the abstract values refined by
      # _call_exported_abstract_eval.
      results = tuple(
          convert_shape(out, out_aval, refined_out_aval)
          for out, out_aval, refined_out_aval in zip(
              call.results[len(ordered_effects) :],
              exported.out_avals,
              ctx.avals_out,
          )
      )
      # Apply out_shardings
      results = tuple(
          wrap_with_sharding(ctx, x, x_aval, x_sharding)
          for x, x_aval, x_sharding in zip(
              results, ctx.avals_out, exported.out_shardings
          )
      )
      return results
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import copy
  import dataclasses
  import functools
  import itertools
  import re
  from typing import Any, Callable, Union
  import warnings
  from absl import logging
  import numpy
  import jax
  from jax import sharding
  from jax._src import ad_util
  from jax._src import config
  from jax._src import core
  from jax._src import dispatch
  from jax._src import dtypes
  from jax._src import effects
  from jax._src.interpreters import mlir
  from jax._src.interpreters import pxla
  from jax._src.lib import xla_client
  from jax._src.lib.mlir import ir
  from jax._src.lib.mlir.dialects import hlo
  from jax._src.lib.mlir.dialects import func
  from jax._src import pjit
  from jax._src import sharding_impls
  from jax._src import source_info_util
  from jax._src import tree_util
  from jax._src import util
  from jax._src import xla_bridge
  from jax.experimental.export import _shape_poly
label_desc: Use image node_id 18 for calling a global function with example usage: _call_exported_lowering(ctx) and returns: results, hlo, x