def _sharded_jit_lowering(ctx, *in_nodes,
                          in_parts, out_parts_thunk, nparts,
                          name, call_jaxpr, local_in_parts,
                          local_out_parts_thunk, local_nparts):
  # We assume any extra leading in_nodes are constants and replicate them.
  num_extra_nodes = len(in_nodes) - len(in_parts)
  assert num_extra_nodes >= 0
  in_parts = (None,) * num_extra_nodes + in_parts

  args = []
  for ns, sharding in safe_zip(
      safe_map(mlir.wrap_singleton_ir_values, in_nodes), in_parts):
    if sharding is not None:
      args.append(
          [mlir.wrap_with_sharding_op(n, xla.sharding_to_proto(sharding))
           for n in ns])
    else:
      args.append(ns)

  sub_ctx = ctx.module_context.replace(
      name_stack=new_name_stack(wrap_name(name, "sharded_jit")))
  fn = mlir.lower_jaxpr_to_fun(sub_ctx, f"sharded_jit_{name}",
                               core.ClosedJaxpr(call_jaxpr, ()))

  output_types = safe_map(mlir.aval_to_ir_types, ctx.avals_out)
  flat_output_types = util.flatten(output_types)
  call = func_dialect.CallOp(flat_output_types,
                             ir.FlatSymbolRefAttr.get(fn.name.value),
                             mlir.flatten_lowering_ir_args(args))
  out_nodes = util.unflatten(call.results, safe_map(len, output_types))

  out_parts = out_parts_thunk()
  outputs = []
  for ns, sharding in safe_zip(out_nodes, out_parts):
    if sharding is not None:
      outputs.append(
          [mlir.wrap_with_sharding_op(n, xla.sharding_to_proto(sharding))
           for n in ns])
    else:
      outputs.append(ns)
  return outputs