exla

defmodule EXLA.Defn do
  @moduledoc false

  require Logger
  require EXLA.Defn.Outfeed, as: Outfeed
  alias Nx.Defn.{Composite, Expr, Tree}
  alias Nx.Tensor, as: T

  alias EXLA.Typespec
  alias EXLA.MLIR.Value
  alias EXLA.MLIR.Function
  alias EXLA.CustomCall.Spec, as: CustomCallSpec

  @doc false
  def __partitions_options__(options) do
    client_name = Keyword.get_lazy(options, :client, &EXLA.Client.default_name/0)
    device_count = EXLA.Client.fetch!(client_name).device_count

    Enum.map(1..device_count//1, &Keyword.put(options, :device_id, &1 - 1))
  end

  @doc false
  def __to_backend__(options) do
    client_name = Keyword.get_lazy(options, :client, &EXLA.Client.default_name/0)

    device_id =
      Keyword.get_lazy(options, :device_id, fn ->
        EXLA.Client.fetch!(client_name).default_device_id
      end)

    {EXLA.Backend, [client: client_name, device_id: device_id]}
  end

  @doc false
  def __jit__(key, vars, fun, args_list, options) do
    __compile__(key, vars, fun, options).(args_list)
  end

  @doc false
  def __shard_jit__(key, mesh, vars, fun, args_list, options) do
    input_shardings = options[:input_shardings]

    # Validate that input_shardings is a list
    if not is_list(input_shardings) do
      raise ArgumentError,
            "input_shardings are required for sharding in EXLA, see EXLA.shard_jit/3 for more information"
    end

    # Validate number of input_shardings matches number of parameters
    num_params = length(vars)
    num_shardings = length(input_shardings)

    if num_shardings != num_params do
      raise ArgumentError,
            "expected #{num_params} input sharding configuration(s), got #{num_shardings}"
    end

    # Convert map format to list format and validate
    {input_shardings_list, unsharded_shape_multipliers} =
      convert_and_validate_input_shardings!(mesh, input_shardings, vars)

    vars = calculate_unsharded_inputs(vars, unsharded_shape_multipliers)

    options =
      options
      |> Keyword.put(:mesh, mesh)
      |> Keyword.put(:input_shardings, input_shardings_list)

    compiled_fun = __compile__(key, vars, fun, options)

    compiled_fun.([args_list])
  end

  defp convert_and_validate_input_shardings!(mesh, input_shardings, vars) do
    # Create a set of all available mesh axes
    all_axes =
      mesh.shape
      |> tuple_size()
      |> axes_for_rank()
      |> MapSet.new()

    # Convert and validate each input tensor's sharding spec
    results =
      Enum.zip_with(input_shardings, vars, fn input_sharding, var ->
        unless is_map(input_sharding) do
          raise ArgumentError,
                "expected input sharding to be a map, got: #{inspect(input_sharding)}"
        end

        # Convert map format to list format
        # Map: %{tensor_dim => [mesh_axes]}
        # List: [[mesh_axes], [mesh_axes], ...]
        tensor_rank = tuple_size(var.shape)

        # Build a mapping from dimension indices to mesh axes
        dim_to_axes_map =
          Enum.reduce(input_sharding, %{}, fn {tensor_dim, mesh_axes}, acc ->
            # Normalize tensor dimension (could be integer, negative integer, or atom name)
            dim_index = Nx.axis_index(var, tensor_dim)

            # Check for duplicate dimension specifications
            if Map.has_key?(acc, dim_index) do
              raise ArgumentError,
                    "tensor dimension #{inspect(tensor_dim)} (index #{dim_index}) was specified multiple times"
            end

            Map.put(acc, dim_index, mesh_axes)
          end)

        # Convert to list format (one entry per tensor dimension)
        dim_shardings_list =
          for dim_idx <- 0..(tensor_rank - 1)//1 do
            Map.get(dim_to_axes_map, dim_idx, [])
          end

        # Track which axes have been used for this specific tensor
        {dim_shardings, _used_axes} =
          Enum.map_reduce(dim_shardings_list, MapSet.new(), fn axis_list, used_axes ->
            # For each dimension, check that axes are valid and not reused
            Enum.each(axis_list, fn axis ->
              if axis not in all_axes do
                raise ArgumentError,
                      "axis #{axis} is not valid for mesh with #{MapSet.size(all_axes)} axes"
              end

              if axis in used_axes do
                raise ArgumentError, "axis #{axis} was used twice in the same input sharding"
              end
            end)

            # Mark these axes as used
            new_used_axes = Enum.reduce(axis_list, used_axes, &MapSet.put(&2, &1))

            accumulated_dim_shardings =
              Enum.reduce(axis_list, 1, fn axis, acc ->
                acc * elem(mesh.shape, axis)
              end)

            {accumulated_dim_shardings, new_used_axes}
          end)

        {dim_shardings_list, dim_shardings}
      end)

    # Separate the list format (for MLIR) and multipliers (for shape calculation)
    {list_format, multipliers} = Enum.unzip(results)
    {list_format, multipliers}
  end

  @doc false
  def __compile__(key, vars, fun, options) do
    {run_options, compile_options} = Keyword.pop(options, :run_options, [])
    debug? = Keyword.get(compile_options, :debug, false)

    callback = &to_computation(&1, &2, &3, &4, &5, compile_options)

    {executable, {used_inputs, outputs, outfeed, _input_typespecs?}} =
      compile(key, vars, fun, compile_options, 0, [], callback)

    if compile_options[:module_compilation] == :to_mlir do
      throw({:mlir_module, executable.ref, MapSet.new(Map.keys(used_inputs)), outputs})
    end

    fn [args] ->
      {time, lock} =
        :timer.tc(fn ->
          EXLA.Defn.Lock.lock(run_key(executable))
        end)

      debug? && Logger.debug("EXLA device #{executable.device_id} lock in #{us_to_ms(time)}ms")

      {time, res} =
        :timer.tc(fn ->
          maybe_outfeed(
            lock,
            executable,
            args,
            used_inputs,
            outputs,
            outfeed,
            run_options,
            !!options[:mesh]
          )
        end)

      debug? &&
        Logger.debug("EXLA execution on device #{executable.device_id} in #{us_to_ms(time)}ms")

      res
    end
  end

  defp to_computation(
         %Function{} = function,
         expr,
         used_typespecs,
         outfeed,
         client,
         options
       ) do
    [callback_pid_value | function_args] = Function.get_arguments(function)

    params =
      Enum.zip_with(used_typespecs, function_args, fn {pos, _typespec}, arg ->
        {pos, arg}
      end)

    unless client do
      raise ArgumentError, "missing client"
    end

    state = %{
      client: client,
      precision: Keyword.get(options, :precision, :default),
      builder: function,
      params: Map.new(params ++ outfeed.infeeds),
      scope_ids: Tree.scope_ids(expr),
      callback_pid_value: callback_pid_value
    }

    {res, cache} = recur_flatten(expr, state, new_cache(outfeed))
    outfeed = cache |> get_outfeed() |> Outfeed.close(function)
    Value.func_return(function, res)
    outfeed
  end

  defp maybe_outfeed(
         lock,
         executable,
         args,
         used_inputs,
         outputs,
         outfeed,
         run_options,
         is_sharded?
       )
       when Outfeed.will_outfeed(outfeed) or Outfeed.has_runtime_calls(outfeed) do
    if is_sharded? do
      raise ArgumentError, "outfeed is not supported for sharded execution yet"
    end

    {buffers, infeeds} =
      EXLA.Defn.Buffers.split_by_value(args, used_inputs, fn
        arg, _i, nil -> EXLA.Defn.Buffers.from_nx!(arg, executable, true)
        arg, i, _depth -> {i, EXLA.Defn.Buffers.from_nx!(arg, executable, false)}
      end)

    infeeds = Map.new(infeeds)

    {:ok, outfeed_pid} =
      Outfeed.start_child(executable, outfeed, Process.group_leader(), infeeds)

    ref = Process.monitor(outfeed_pid)

    run_options = Keyword.put(run_options, :callback_server_pid, outfeed_pid)

    {:ok, runner} =
      EXLA.Defn.Runner.start_link(lock, fn ->
        EXLA.Executable.run(executable, [Enum.reverse(buffers)], run_options)
      end)

    _ = EXLA.Defn.Lock.transfer(lock, fn -> send(runner, lock) end, outfeed_pid)

    receive do
      {:DOWN, ^ref, _, _, _} ->
        results = EXLA.Defn.Runner.read(runner)

        Enum.map(results, fn result ->
          EXLA.Defn.Buffers.to_nx!(result, outputs, executable.mesh)
        end)
    end
  end

  defp maybe_outfeed(
         lock,
         executable,
         args,
         used_inputs,
         outputs,
         _outfeed,
         run_options,
         is_sharded?
       ) do
    try do
      args = if is_sharded?, do: args, else: [args]

      # Check if args are pre-sliced (list of arglists for each partition)
      input_lists =
        Enum.map(args, fn partition_args ->
          EXLA.Defn.Buffers.filter_by_indexes(partition_args, used_inputs, fn arg, _i ->
            EXLA.Defn.Buffers.from_nx!(arg, executable)
          end)
        end)

      EXLA.Executable.run(executable, input_lists, run_options)
    else
      [_ | _] = results ->
        # For sharded execution, we get a list of results (one per partition)
        Enum.map(results, fn result ->
          EXLA.Defn.Buffers.to_nx!(result, outputs, executable.mesh)
        end)
    after
      EXLA.Defn.Lock.unlock(lock)
    end
  end

  defp run_key(%{client: %{platform: :host}, device_id: device_id}), do: [:host | device_id]
  defp run_key(%{client: %{ref: ref}, device_id: device_id}), do: [ref | device_id]

  ## Compile

  defp compile(
         key,
         vars,
         fun,
         options,
         used_buffers,
         used_inputs,
         to_computation
       ) do
    {cache, options} = Keyword.pop(options, :cache, true)
    {hooks, options} = Keyword.pop(options, :hooks, %{})
    {debug?, options} = Keyword.pop(options, :debug, false)
    {lazy_transfers, options} = Keyword.pop(options, :lazy_transfers, :opt_in)

    {client_name, options} = Keyword.pop_lazy(options, :client, &EXLA.Client.default_name/0)
    client = EXLA.Client.fetch!(client_name)
    mesh = options[:mesh]
    input_shardings = Keyword.get(options, :input_shardings, [])

    if !mesh and input_shardings != [] do
      raise ArgumentError,
            "input sharding configuration provided but no device mesh was provided"
    end

    {args_key, reverse_args_identifiers} =
      Enum.map_reduce(vars, [], fn var, acc ->
        Nx.Defn.Composite.traverse(var, acc, fn
          %T{vectorized_axes: vectorized_axes} = t, acc ->
            %T{type: type, shape: shape, names: names} = Nx.devectorize(t)
            identifier = {type, shape, names}
            cache_key = {type, shape, names, vectorized_axes}
            {cache_key, [identifier | acc]}
        end)
      end)

    disk_key = %{
      client: client.name,
      args: args_key,
      lazy_transfers: lazy_transfers,
      hooks: Map.keys(hooks),
      options: options
    }

    EXLA.Defn.Disk.cache(cache, client, disk_key, debug?, fn ->
      {{expr_cache_fun, comp_cache_fun}, options} =
        if cache do
          Keyword.pop(options, EXLA, {&EXLA.Defn.LockedCache.run/2, &EXLA.Defn.LockedCache.run/2})
        else
          cache_fun = fn _key, fun -> fun.() end
          {{cache_fun, cache_fun}, Keyword.delete(options, EXLA)}
        end

      {eval_time, {expr, {ref, outputs, {used_inputs, defined_hooks}}}} =
        :timer.tc(fn ->
          expr_cache_fun.({key, args_key, lazy_transfers}, fn ->
            expr = fun.(vars)
            inputs_and_hooks = Outfeed.used_inputs_and_hooks(expr, used_inputs, lazy_transfers)
            {expr, {make_ref(), Nx.to_template(expr), inputs_and_hooks}}
          end)
        end)

      if debug? do
        hit_or_miss = if expr, do: "miss", else: "hit"

        Logger.debug(
          "EXLA defn evaluation #{inspect(key)} cache #{hit_or_miss} in #{us_to_ms(eval_time)}ms"
        )
      end

      outfeed = Outfeed.new(hooks, defined_hooks)
      comp_key = {ref, client.name, outfeed.used_hooks, lazy_transfers, options}

      {comp_time, {evaled, {xla_time, executable, inputs_and_typespecs, outfeed}}} =
        :timer.tc(fn ->
          comp_cache_fun.(comp_key, fn ->
            {reverse_inputs_and_typespecs, reverse_infeeds} =
              reverse_args_identifiers
              |> Enum.reverse()
              |> EXLA.Defn.Buffers.split_by_value(used_inputs, fn
                {type, shape, _names}, i, nil -> {i, Typespec.tensor(type, shape)}
                {type, shape, _names}, i, depth -> {i, depth, Typespec.tensor(type, shape)}
              end)

            inputs_and_typespecs = Enum.reverse(reverse_inputs_and_typespecs)

            out_typespecs =
              [outputs]
              |> Nx.Defn.Composite.flatten_list()
              |> Enum.map(fn t ->
                t
                |> Nx.devectorize()
                |> then(&Typespec.tensor(&1.type, &1.shape))
              end)

            expr = expr || fun.(vars)
            expr = Nx.Defn.Composite.traverse(expr, &Nx.devectorize/1)
            callback_pid_typespec = EXLA.Executable.callback_server_pid_typespec()

            comp_typespecs =
              for {i, typespec} <- inputs_and_typespecs, i >= used_buffers, do: typespec

            comp_typespecs = [callback_pid_typespec | comp_typespecs]

            EXLA.MLIR.Module.new(comp_typespecs, out_typespecs, fn builder ->
              # Add device mesh to module if provided
              if mesh do
                EXLA.MLIR.Module.add_mesh(builder.module, mesh)

                # Apply sharding annotations to function arguments if provided
                Enum.with_index(input_shardings, fn sharding, arg_index ->
                  Function.set_arg_sharding(builder, arg_index + 1, {mesh, sharding})
                end)
              end

              # Only create the token when we know it will actually be
              # used, that is: streaming, lazy transfers or hooks
              outfeed =
                if reverse_infeeds != [] or hooks != %{} or defined_hooks != %{} do
                  outfeed
                  |> Outfeed.with_token(Value.create_token(builder))
                  |> Outfeed.add_infeeds(builder, reverse_infeeds)
                else
                  outfeed
                end

              outfeed = to_computation.(builder, expr, inputs_and_typespecs, outfeed, client)

              # Compute num_partitions from mesh and enable SPMD if mesh is provided
              options =
                if mesh do
                  num_partitions = Tuple.product(mesh.shape)

                  options
                  |> Keyword.put(:num_partitions, num_partitions)
                  |> Keyword.put(:use_spmd, true)
                else
                  options
                end

              {xla_time, executable} =
                :timer.tc(fn ->
                  EXLA.MLIR.Module.compile(
                    builder.module,
                    client,
                    comp_typespecs,
                    builder.return_typespecs,
                    options
                  )
                end)

              {:ok, {xla_time, executable, inputs_and_typespecs, %{outfeed | infeeds: []}}}
            end)
          end)
        end)

      cond do
        not debug? ->
          :ok

        evaled ->
          Logger.debug(
            "EXLA compilation #{inspect(key)} cache miss in #{us_to_ms(comp_time)}ms (#{us_to_ms(xla_time)}ms in XLA)"
          )

        true ->
          Logger.debug("EXLA compilation #{inspect(key)} cache hit in #{us_to_ms(comp_time)}ms")
      end

      if expr || evaled do
        measurements = %{
          eval_time: eval_time,
          compile_time: comp_time,
          total_time: eval_time + comp_time
        }

        :telemetry.execute([:exla, :compilation], measurements, %{key: key})
      end

      if evaled && cache, do: check_recompilation(key, args_key)

      outfeed = Outfeed.with_user_hooks(outfeed, hooks)
      {executable, {used_inputs, outputs, outfeed, inputs_and_typespecs}}
    end)
  end

  @recompilation_threshold 10

  defp check_recompilation(key, args_key) do
    {:module, mod} = :erlang.fun_info(key, :module)
    {:new_index, idx} = :erlang.fun_info(key, :new_index)
    count = EXLA.Defn.LockedCache.count({mod, idx, args_key})

    if count == @recompilation_threshold do
      Logger.warning(
        "EXLA has compiled #{inspect(key)} #{count} times with the same input " <>
          "shapes. This typically means tensor values are being captured inside a " <>
          "closure passed to defn, jit, or value_and_grad. Each distinct captured " <>
          "value forces a full recompilation. Pass changing tensors as explicit " <>
          "function arguments instead."
      )
    end
  end

  defp us_to_ms(time), do: Float.round(time / 1000, 1)

  ## Operator handling

  defp recur_flatten(composite, state, cache) do
    {acc, cache} =
      Composite.reduce(composite, {[], cache}, fn %T{} = expr, {acc, cache} ->
        {expr, cache} = recur_operator(expr, state, cache)
        {[acc, expr], cache}
      end)

    {List.flatten(acc), cache}
  end

  defp recur_operator(%T{data: %Expr{id: id, op: op}} = expr, state, cache) do
    case cache do
      %{^id => res} ->
        {res, cache}

      %{} ->
        {res, cache} = cached_recur_operator(op, Nx.devectorize(expr), state, cache)
        {res, Map.put(cache, id, res)}
    end
  end

  defp cached_recur_operator(
         :while,
         %T{data: %Expr{args: args}},
         %{builder: %Function{} = function} = state,
         cache
       ) do
    [initial_arg, _arg, pred, body] = args

    {initial, cache} = recur_composite(initial_arg, state, cache)

    initial =
      if token = get_token(cache) do
        [token | initial]
      else
        initial
      end

    {pred_computation, cache} = mlir_while_computation(pred, initial, {:pred, 8}, state, cache)
    {body_computation, cache} = mlir_while_computation(body, initial, :with_token, state, cache)

    results =
      Value.while(function, pred_computation, body_computation, List.flatten(initial))

    if get_token(cache) do
      [token | results] = results
      result = wrap_tuple_result(results, initial_arg)
      {result, update_token(cache, token)}
    else
      result = wrap_tuple_result(results, initial_arg)
      {result, cache}
    end
  end

  defp cached_recur_operator(:cond, %T{data: %Expr{args: args}} = t, state, cache) do
    [clauses, last] = args

    case clauses do
      [{pred, on_true}] ->
        to_if(pred, on_true, last, state, cache)

      _ ->
        # We convert cond into a nested tree of conds in order to compile it to ifs
        %T{data: %Expr{args: [[{pred, on_true}], on_false]}} =
          clauses
          |> Enum.reverse()
          |> Enum.reduce(last, fn {pred, on_true}, on_false ->
            update_in(t.data, fn data ->
              %{data | args: [[{pred, on_true}], on_false], id: make_ref()}
            end)
          end)

        to_if(pred, on_true, on_false, state, cache)
    end
  end

  defp cached_recur_operator(:parameter, %T{data: %Expr{args: [i]}}, state, cache) do
    {Map.fetch!(state.params, i), cache}
  end

  defp cached_recur_operator(:fun, %T{data: %Expr{args: args}, type: type}, state, cache) do
    [args, expr, {_, _, _}] = args
    {fun_computation(args, expr, type, state), cache}
  end

  # StableHLO-style lowering (gather, top_k, fft): not the C custom_call path;
  # see `EXLA.CustomCall` for blocks that delegate to native CPU kernels.

  defp cached_recur_operator(
         :block,
         %T{
           data: %Expr{
             args: [%Nx.Block.Take{axis: axis}, [tensor, indices], expr, _callback]
           }
         },
         state,
         cache
       ) do
    {tensor, cache} = recur_operator(tensor, state, cache) |> unwrap_single_tensor!()
    {indices, cache} = recur_operator(indices, state, cache) |> unwrap_single_tensor!()

    tensor_rank = tensor |> op_shape() |> tuple_size()
    indices_rank = indices |> op_shape() |> tuple_size()
    result_rank = tensor_rank - 1 + indices_rank

    index_vector_dim = indices_rank
    slice_sizes = tensor |> op_shape() |> put_elem(axis, 1) |> Tuple.to_list()

    {left, right} = result_rank |> axes_for_rank() |> Enum.split(axis)
    offset_dims = left ++ Enum.drop(right, indices_rank)

    collapsed_slice_dims = [axis]
    start_index_map = [axis]

    result =
      Value.gather(
        tensor,
        indices,
        index_vector_dim,
        slice_sizes,
        offset_dims,
        collapsed_slice_dims,
        start_index_map,
        expr_to_typespec(expr)
      )

    {result, cache}
  end

  defp cached_recur_operator(
         :block,
         %T{data: %Expr{args: [%Nx.Block.TopK{k: k}, [tensor], expr, _callback]}},
         state,
         cache
       ) do
    {tensor, cache} = recur_operator(tensor, state, cache) |> unwrap_single_tensor!()
    {values, idx} = expr
    typespecs = [expr_to_typespec(values), expr_to_typespec(idx)]
    results = Value.top_k(tensor, k, typespecs)
    {results, cache}
  end

  defp cached_recur_operator(
         :block,
         %T{data: %Expr{args: [%Nx.Block.FFT2{} = fft2_struct, [tensor], expr, _callback]}},
         state,
         cache
       ) do
    {tensor, cache} = recur_operator(tensor, state, cache) |> unwrap_single_tensor!()

    opts = [lengths: fft2_struct.lengths, axes: fft2_struct.axes]
    opts = if eps = fft2_struct.eps, do: Keyword.put(opts, :eps, eps), else: opts

    {fft2(&Value.fft(&1, :fft, &2, &3), [tensor, opts], expr), cache}
  end

  defp cached_recur_operator(
         :block,
         %T{data: %Expr{args: [%Nx.Block.IFFT2{} = ifft2_struct, [tensor], expr, _callback]}},
         state,
         cache
       ) do
    {tensor, cache} = recur_operator(tensor, state, cache) |> unwrap_single_tensor!()

    opts = [lengths: ifft2_struct.lengths, axes: ifft2_struct.axes]
    opts = if eps = ifft2_struct.eps, do: Keyword.put(opts, :eps, eps), else: opts

    {fft2(&Value.fft(&1, :ifft, &2, &3), [tensor, opts], expr), cache}
  end

  defp cached_recur_operator(
         :block,
         %T{data: %Expr{args: [%Nx.Block.RFFT{} = rfft_struct, [tensor], expr, _callback]}},
         state,
         cache
       ) do
    {tensor, cache} = recur_operator(tensor, state, cache) |> unwrap_single_tensor!()

    opts = [length: rfft_struct.length, axis: rfft_struct.axis]
    opts = if eps = rfft_struct.eps, do: Keyword.put(opts, :eps, eps), else: opts

    input_type = Nx.Type.to_real(expr.type)

    {fft(&Value.fft(&1, :rfft, &2, &3), input_type, expr.type, [tensor, opts], expr), cache}
  end

  defp cached_recur_operator(
         :block,
         %T{data: %Expr{args: [%Nx.Block.IRFFT{} = irfft_struct, [tensor], expr, _callback]}},
         state,
         cache
       ) do
    {tensor, cache} = recur_operator(tensor, state, cache) |> unwrap_single_tensor!()

    opts = [length: irfft_struct.length, axis: irfft_struct.axis]
    opts = if eps = irfft_struct.eps, do: Keyword.put(opts, :eps, eps), else: opts

    n = irfft_struct.length
    input_type = Nx.Type.to_complex(expr.type)

    {fft(
       &Value.fft(&1, :irfft, &2, &3),
       input_type,
       expr.type,
       div(n, 2) + 1,
       [tensor, opts],
       expr
     ), cache}
  end

  # C-backed custom_call blocks (QR, Eigh, …): `EXLA.CustomCall`; else compile default callback.
  defp cached_recur_operator(
         :block,
         %T{data: %Expr{args: [struct, in_args, out, _callback]}},
         %{client: client, builder: %Function{}} = state,
         cache
       ) do
    {call_args, cache} = Enum.map_reduce(in_args, cache, &recur_operator(&1, state, &2))

    case EXLA.CustomCall.call(struct, out, in_args, client) do
      :skip ->
        default_block_implementation(struct, call_args, out, state, cache)

      {:ok, %CustomCallSpec{} = spec} ->
        dictionary_entries =
          case spec.attributes do
            list when is_list(list) ->
              list

            other ->
              raise ArgumentError,
                    "EXLA.CustomCall.Spec attributes must be a list of {binary_key, binary_attr} pairs, got: #{inspect(other)}"
          end

        call_args = cast_custom_call_operands(call_args, spec.operand_element_types)

        out_typespecs =
          [out]
          |> Composite.flatten_list()
          |> Enum.map(&expr_to_typespec/1)

        lowered =
          Value.custom_call(call_args, out_typespecs, spec.call_target_name, dictionary_entries)
          |> wrap_tuple_result(out)

        {lowered, cache}

      other ->
        raise ArgumentError,
              "EXLA.CustomCall.call/4 must return :skip or {:ok, %EXLA.CustomCall.Spec{}}, got: #{inspect(other)}"
    end
  end

  defp cached_recur_operator(
         :runtime_call,
         %T{data: %Expr{id: id, args: [tensor_expr, fun, out_template, opts]}} = expr,
         %{client: %EXLA.Client{platform: :host}, callback_pid_value: callback_pid_value} =
           state,
         cache
       ) do
    # Flatten the tensor_or_container expression into its tensor leaves so we
    # can compile each as an independent operand to the host callback.
    tensor_exprs = Composite.flatten_list([tensor_expr])

    {arg_values, cache} =
      Enum.map_reduce(tensor_exprs, cache, fn arg, cache ->
        recur_operator(arg, state, cache) |> unwrap_single_tensor!()
      end)

    # Build a template container for the tensor_or_container argument so the
    # callback server can reconstruct the full structure from a flat list of
    # decoded tensors.
    arg_template = Nx.to_template(tensor_expr)

    cache = add_runtime_callback(cache, {id, fun, out_template, arg_template, opts})

    typespecs = container_to_typespecs(out_template)

    unless callback_pid_value do
      raise "internal bug: runtime_call callback pid operand is missing"
    end

    results =
      Value.runtime_call([callback_pid_value | arg_values], typespecs, id)

    {wrap_tuple_result(results, expr), cache}
  end

  defp cached_recur_operator(
         :runtime_call,
         %T{data: %Expr{id: id, args: [tensor_expr, fun, out_template, opts]}} = expr,
         %{client: %EXLA.Client{platform: :cuda}, callback_pid_value: callback_pid_value} =
           state,
         cache
       ) do
    tensor_exprs = Composite.flatten_list([tensor_expr])

    {arg_values, cache} =
      Enum.map_reduce(tensor_exprs, cache, fn arg, cache ->
        recur_operator(arg, state, cache) |> unwrap_single_tensor!()
      end)

    arg_template = Nx.to_template(tensor_expr)

    cache = add_runtime_callback(cache, {id, fun, out_template, arg_template, opts})

    typespecs = container_to_typespecs(out_template)

    unless callback_pid_value do
      raise "internal bug: runtime_call callback pid operand is missing"
    end

    results =
      Value.runtime_call([callback_pid_value | arg_values], typespecs, id)

    {wrap_tuple_result(results, expr), cache}
  end

  defp cached_recur_operator(
         :runtime_call,
         _expr,
         %{client: %EXLA.Client{platform: platform}},
         _cache
       ) do
    raise """
    Nx.runtime_call/4 is currently only supported for EXLA CPU (platform: :host) and CUDA (platform: :cuda),
    but the active EXLA client is configured for platform #{inspect(platform)}.
    Please run on the :host or :cuda client or wait for future segmentation-based support.
    """
  end

  defp cached_recur_operator(
         :lu,
         %T{
           data: %Expr{args: [{p_expr, l_expr, u_expr}, %{type: {type_kind, _}} = tensor, _opts]}
         },
         %{client: %{platform: :host}} = state,
         cache
       )
       when type_kind != :c do
    # We only want to accelerate the LU operation for real inputs on the host device.
    # Otherwise, we use the default implementation in Nx.
    {tensor, cache} = recur_operator(tensor, state, cache) |> unwrap_single_tensor!()

    tensor =
      if op_type(tensor) != u_expr.type do
        to_type(tensor, u_expr.type)
      else
        tensor
      end

    {p, l, u} =
      Value.lu(
        tensor,
        expr_to_typespec(p_expr),
        expr_to_typespec(l_expr),
        expr_to_typespec(u_expr)
      )

    {[p, l, u], cache}
  end

  defp cached_recur_operator(:attach_token, %T{data: %Expr{args: [token, expr]}}, state, cache) do
    {op, cache} = recur_operator(expr, state, cache)
    {_, cache} = recur_operator(token, state, cache)
    {op, cache}
  end

  defp cached_recur_operator(:token, %T{data: %Expr{args: [token]}}, state, cache) do
    builder = state.builder

    cache =
      List.foldr(token.hooks, cache, fn %{name: name, expr: expr}, cache ->
        # First traverse the child because if it has hooks,
        # we need to handle them first
        {tuple, cache} = recur_flatten(expr, state, cache)

        cache
        |> get_outfeed()
        |> Outfeed.maybe_add_function_hook(builder, tuple, name, expr)
        |> then(&put_outfeed(cache, &1))
      end)

    {[], cache}
  end

  defp cached_recur_operator(op, expr, state, cache) do
    {args, cache} = Tree.apply_args(expr, cache, &recur_operator(&1, state, &2))
    {to_operator(op, args, expr, state), cache}
  end

  defp cast_custom_call_operands(call_args, :default), do: call_args

  defp cast_custom_call_operands(call_args, types) when is_list(types) do
    n = length(call_args)

    if length(types) != n do
      raise ArgumentError,
            "EXLA.CustomCall.Spec operand_element_types must be a list of length #{n} (one per block input), got length #{length(types)}"
    end

    Enum.zip_with(call_args, types, fn arg, desired ->
      ts = Value.get_typespec(arg)

      if ts.type == desired do
        arg
      else
        Value.convert(arg, Typespec.tensor(desired, ts.shape))
      end
    end)
  end

  defp cast_custom_call_operands(_call_args, other) do
    raise ArgumentError,
          "EXLA.CustomCall.Spec operand_element_types must be :default or a list of Nx types, got: #{inspect(other)}"
  end

  defp default_block_implementation(struct, call_args, expr, state, cache) do
    %module{} = struct
    key = computation_key(module, [struct | call_args])

    {call_body, cache} =
      case cache do
        %{^key => computation} ->
          {computation, cache}

        %{} ->
          {computation, cache} =
            block_computation(
              block_subfunction_description(struct),
              call_args,
              expr,
              state,
              cache
            )

          {computation, Map.put(cache, key, computation)}
      end

    if token = get_token(cache) do
      typespecs = [Typespec.token() | container_to_typespecs(expr)]
      [token | result] = Value.call(state.builder, [token | call_args], call_body, typespecs)
      {wrap_tuple_result(result, expr), update_token(cache, token)}
    else
      typespecs = container_to_typespecs(expr)
      result = Value.call(state.builder, call_args, call_body, typespecs)
      {wrap_tuple_result(result, expr), cache}
    end
  end

  ## to_operator creation

  defp to_operator(:constant, [constant], ans, state) do
    op = to_constant(state.builder, constant, ans.type)

    if ans.shape == {} do
      op
    else
      Value.broadcast_in_dim(op, [], expr_to_typespec(ans))
    end
  end

  defp to_operator(:tensor, [tensor], _ans, state) do
    tensor = Nx.devectorize(tensor)

    case tensor.shape do
      {} ->
        to_constant(state.builder, Nx.to_number(tensor), tensor.type)

      shape ->
        Value.constant(
          state.builder,
          Nx.to_flat_list(tensor),
          Typespec.tensor(tensor.type, shape)
        )
    end
  end

  defp to_operator(:iota, [axis], ans, state) do
    EXLA.Lib.iota(state.builder, axis, expr_to_typespec(ans))
  end

  defp to_operator(:eye, [], %{type: type, shape: shape}, state) do
    iota_type = Nx.Type.merge_number({:u, 8}, Tuple.product(shape))
    iota_typespec = Typespec.tensor(iota_type, shape)
    rank = tuple_size(shape)

    i0 = Value.iota(state.builder, rank - 2, iota_typespec)
    i1 = Value.iota(state.builder, rank - 1, iota_typespec)

    typespec = Typespec.tensor({:pred, 8}, shape)
    Value.equal(i0, i1, typespec) |> to_type(type)
  end

  ## to_operator shape

  defp to_operator(:reshape, [%Value{} = op], ans, _state) do
    Value.reshape(op, expr_to_typespec(ans))
  end

  defp to_operator(:pad, [%Value{} = op, %Value{} = value, padding_config], ans, _state) do
    Value.pad(
      to_type(op, ans.type),
      to_type(value, ans.type),
      padding_config,
      expr_to_typespec(ans)
    )
  end

  defp to_operator(:broadcast, [%Value{} = op, _shape, axes], ans, _state) do
    Value.broadcast_in_dim(to_type(op, ans.type), axes, expr_to_typespec(ans))
  end

  defp to_operator(:transpose, [%Value{} = op, axes], ans, _state) do
    Value.transpose(op, axes, expr_to_typespec(ans))
  end

  defp to_operator(:squeeze, [%Value{} = op, _axes], ans, _state) do
    Value.reshape(op, expr_to_typespec(ans))
  end

  ## to_operator others

  defp to_operator(:metadata, [op, _metadata], _ans, _state) do
    case op do
      %Value{} ->
        op

      op when is_tuple(op) ->
        Tuple.to_list(op)
    end
  end

  defp to_operator(:elem, [op, index], _ans, _state) when is_list(op) do
    Enum.fetch!(op, index)
  end

  defp to_operator(
         :dot,
         [
           %Value{} = left,
           contract_axes1,
           batch_axes1,
           %Value{} = right,
           contract_axes2,
           batch_axes2
         ],
         ans,
         state
       ) do
    precision = state.precision

    Value.dot_general(
      left,
      right,
      {contract_axes1, batch_axes1, contract_axes2, batch_axes2},
      precision,
      expr_to_typespec(ans)
    )
  end

  defp to_operator(
         :conv,
         [operand, kernel, opts],
         ans,
         state
       ) do
    padding = opts[:padding]
    strides = opts[:strides]
    input_dilation = opts[:input_dilation]
    kernel_dilation = opts[:kernel_dilation]
    feature_group_count = opts[:feature_group_size]
    batch_group_count = opts[:batch_group_size]

    %{type: output_type} = ans

    dimension_numbers =
      {opts[:input_permutation], opts[:kernel_permutation], opts[:output_permutation]}

    # Ensure both types are floating
    operand = to_type(operand, output_type)
    kernel = to_type(kernel, output_type)

    Value.convolution(
      operand,
      kernel,
      strides,
      padding,
      input_dilation,
      kernel_dilation,
      dimension_numbers,
      feature_group_count,
      batch_group_count,
      state.precision,
      expr_to_typespec(ans)
    )
  end

  defp to_operator(
         :select,
         [%Value{} = pred, %Value{} = on_true, %Value{} = on_false],
         %{type: type, shape: shape} = ans,
         _state
       ) do
    pred = to_type(pred, {:pred, 8})

    typespec = expr_to_typespec(ans)

    on_true =
      on_true
      |> to_type(type)
      |> Value.broadcast_in_dim(broadcast_axes(op_shape(on_true), shape), typespec)

    on_false =
      on_false
      |> to_type(type)
      |> Value.broadcast_in_dim(broadcast_axes(op_shape(on_false), shape), typespec)

    Value.select(pred, on_true, on_false, typespec)
  end

  defp to_operator(:triangular_solve, [%Value{} = a, b, opts], %{type: type} = ans, _state) do
    left_side = Keyword.fetch!(opts, :left_side)
    lower = Keyword.fetch!(opts, :lower)
    transform = Keyword.fetch!(opts, :transform_a)

    # StableHLO triangular_solve doesn't support conjugate-only transform,
    # so we pre-conjugate A and use :none instead.
    # For real types, conjugate is a no-op.
    {a, transform} =
      case transform do
        :conjugate ->
          if Nx.Type.complex?(type) do
            a_typespec = Value.get_typespec(a)
            {Value.conjugate(to_type(a, type), Typespec.to_type(a_typespec, type)), :none}
          else
            {a, :none}
          end

        other ->
          {a, other}
      end

    a_shape = Value.get_typespec(a).shape
    b_shape = Value.get_typespec(b).shape

    if tuple_size(a_shape) > tuple_size(b_shape) do
      b_shape = Tuple.insert_at(b_shape, tuple_size(b_shape), 1)

      b =
        b
        |> to_type(type)
        |> Value.reshape(Typespec.tensor(type, b_shape))

      typespec = Typespec.tensor(type, b_shape)

      to_type(a, type)
      |> Value.triangular_solve(b, left_side, lower, transform, typespec)
      |> Value.reshape(Typespec.tensor(type, ans.shape))
    else
      typespec = Typespec.tensor(type, ans.shape)

      to_type(a, type)
      |> Value.triangular_solve(to_type(b, type), left_side, lower, transform, typespec)
    end
  end

  ## to_operator element-wise

  defp to_operator(:negate, [%Value{} = op], ans, _state),
    do: Value.negate(op, expr_to_typespec(ans))

  defp to_operator(:abs, [%Value{} = op], ans, _state), do: Value.abs(op, expr_to_typespec(ans))

  defp to_operator(:sign, [%Value{} = op], ans, state) do
    typespec = expr_to_typespec(ans)

    case typespec.type do
      {:u, _} ->
        one = Value.constant(state.builder, [1], Typespec.to_shape(typespec, {}))

        one
        |> Value.broadcast_in_dim([], typespec)
        |> Value.min(op, typespec)

      _ ->
        Value.sign(op, typespec)
    end
  end

  defp to_operator(:right_shift, [%Value{} = left, %Value{} = right], out, _state) do
    op =
      if match?({:u, _}, out.type),
        do: :right_shift_logical,
        else: :right_shift_arithmetic

    apply_mlir_broadcasted_bin_op(op, out, left, right)
  end

  @bin_op [:add, :subtract, :multiply, :min, :max, :remainder, :pow, :divide, :atan2] ++
            [:bitwise_and, :bitwise_or, :bitwise_xor, :left_shift]

  defp to_operator(op, [%Value{} = left, %Value{} = right], out, _state)
       when op in @bin_op do
    apply_mlir_broadcasted_bin_op(op, out, left, right)
  end

  defp to_operator(:quotient, [left, right], ans, state) do
    to_operator(:divide, [to_type(left, ans.type), to_type(right, ans.type)], ans, state)
  end

  @bin_comp_op [:equal, :not_equal, :greater, :less, :greater_equal, :less_equal]

  defp to_operator(op, [%Value{} = left, %Value{} = right], ans, _state)
       when op in @bin_comp_op do
    apply_mlir_broadcasted_bin_op(op, ans, left, right)
  end

  @bin_pred_op [logical_and: :bitwise_and, logical_or: :bitwise_or, logical_xor: :bitwise_xor]

  for {logical, bitwise} <- @bin_pred_op do
    defp to_operator(unquote(logical), [%Value{} = left, %Value{} = right], ans, _state) do
      apply_mlir_broadcasted_bin_op(
        unquote(bitwise),
        ans,
        to_mlir_logical(left),
        to_mlir_logical(right)
      )
    end
  end

  @unary_op [:exp, :expm1, :log, :log1p, :sigmoid, :cos, :sin, :tan, :tanh, :sqrt, :rsqrt, :cbrt] ++
              [:bitwise_not, :count_leading_zeros, :population_count, :cosh, :sinh, :acos] ++
              [:asin, :atan, :floor, :ceil, :round, :acosh, :asinh, :atanh, :erf] ++
              [:erfc, :erf_inv, :conjugate]

  defp to_operator(op, [%Value{} = arg], %{type: type} = ans, _state)
       when op in @unary_op do
    apply(Value, op, [to_type(arg, type), expr_to_typespec(ans)])
  end

  defp to_operator(:fft, [%Value{} | _] = args, out, _state),
    do: fft(&Value.fft(&1, :fft, &2, &3), out.type, out.type, args, out)

  defp to_operator(:ifft, [%Value{} | _] = args, out, _state),
    do: fft(&Value.fft(&1, :ifft, &2, &3), out.type, out.type, args, out)

  defp to_operator(:is_nan, [%Value{} = arg], out, _state),
    do: Value.is_nan(arg, expr_to_typespec(out))

  defp to_operator(:is_infinity, [%Value{} = arg], out, _state),
    do: Value.is_infinity(arg, expr_to_typespec(out))

  # These operations do the type conversion implicitly, and so
  # we cannot mess with the output type (e.g. the to_type conversion)
  # because it will throw an error
  @complex_op [:real, :imag]

  defp to_operator(op, [%Value{} = arg], ans, _state)
       when op in @complex_op do
    maybe_cast_arg =
      if Nx.Type.integer?(op_type(arg)) do
        to_type(arg, ans.type)
      else
        arg
      end

    apply(Value, op, [maybe_cast_arg, expr_to_typespec(ans)])
  end

  defp to_operator(:as_type, [arg], %{type: type}, _state) do
    to_type(arg, type)
  end

  defp to_operator(:bitcast, [%Value{} = arg], ans, _state) do
    if op_type(arg) == ans.type do
      arg
    else
      Value.bitcast_convert(arg, expr_to_typespec(ans))
    end
  end

  ## to_operator reduction

  defp to_operator(:all, [arg, opts], %{shape: shape}, %{builder: %Function{}} = state) do
    to_aggregate(:bitwise_and, {:u, 8}, shape, to_mlir_logical(arg), 1, opts, state)
  end

  defp to_operator(:any, [arg, opts], %{shape: shape}, %{builder: %Function{}} = state) do
    to_aggregate(:bitwise_or, {:u, 8}, shape, to_mlir_logical(arg), 0, opts, state)
  end

  defp to_operator(:all, [arg, opts], %{shape: shape}, state) do
    to_aggregate(:bitwise_and, {:pred, 8}, shape, arg, 1, opts, state)
  end

  defp to_operator(:any, [arg, opts], %{shape: shape}, state) do
    to_aggregate(:bitwise_or, {:pred, 8}, shape, arg, 0, opts, state)
  end

  defp to_operator(:sum, [arg, opts], %{type: type, shape: shape}, state) do
    to_aggregate(:add, type, shape, arg, 0, opts, state)
  end

  defp to_operator(:product, [arg, opts], %{type: type, shape: shape}, state) do
    to_aggregate(:multiply, type, shape, arg, 1, opts, state)
  end

  defp to_operator(:reduce_max, [arg, opts], %{type: type, shape: shape}, state) do
    min_number = EXLA.Lib.min_number(state.builder, type)
    to_aggregate(:max, type, shape, arg, min_number, opts, state)
  end

  defp to_operator(:reduce_min, [arg, opts], %{type: type, shape: shape}, state) do
    max_number = EXLA.Lib.max_number(state.builder, type)
    to_aggregate(:min, type, shape, arg, max_number, opts, state)
  end

  defp to_operator(
         :reduce,
         [%Value{} = arg, %Value{} = acc, opts, fun],
         %{type: type, shape: shape},
         _state
       ) do
    arg = to_type(arg, type)
    keep_axes = opts[:keep_axes]
    reduce_axes = reduce_axes(arg, opts[:axes])

    typespec = Typespec.tensor(type, remove_axes(op_shape(arg), reduce_axes))
    [result] = Value.reduce(fun, [to_type(acc, type)], [arg], reduce_axes, [typespec])

    if keep_axes do
      Value.reshape(result, Typespec.tensor(type, shape))
    else
      result
    end
  end

  defp to_operator(:window_sum, [arg, window_dims, opts], ans, state) do
    to_window_aggregate(:add, ans, arg, 0, window_dims, opts, state)
  end

  defp to_operator(:window_max, [arg, window_dims, opts], %{type: type} = ans, state) do
    min_number = EXLA.Lib.min_number(state.builder, type)
    to_window_aggregate(:max, ans, arg, min_number, window_dims, opts, state)
  end

  defp to_operator(:window_min, [arg, window_dims, opts], %{type: type} = ans, state) do
    max_number = EXLA.Lib.max_number(state.builder, type)
    to_window_aggregate(:min, ans, arg, max_number, window_dims, opts, state)
  end

  defp to_operator(:window_product, [arg, window_dims, opts], ans, state) do
    to_window_aggregate(:multiply, ans, arg, 1, window_dims, opts, state)
  end

  defp to_operator(
         :window_reduce,
         [arg, acc, window_dimensions, opts, fun],
         %{type: type} = ans,
         %{builder: %Function{}}
       ) do
    padding_config = opts[:padding]
    strides = opts[:strides]
    window_dilations = opts[:window_dilations]
    arg = to_type(arg, type)
    acc = to_type(acc, type)

    [result] =
      Value.window_reduce(
        fun,
        [acc],
        [arg],
        Tuple.to_list(window_dimensions),
        strides,
        List.duplicate(1, tuple_size(op_shape(arg))),
        window_dilations,
        padding_config,
        [expr_to_typespec(ans)]
      )

    result
  end

  defp to_operator(
         :window_scatter_max,
         [%Value{} = arg, %Value{} = source, %Value{} = init_value, window_dimensions, opts],
         %{type: type} = ans,
         _state
       ) do
    padding_config = opts[:padding]
    strides = opts[:strides]

    arg = to_type(arg, type)
    source = to_type(source, type)
    init_value = to_type(init_value, type)

    Value.select_and_scatter(
      arg,
      source,
      init_value,
      :gt,
      Tuple.to_list(window_dimensions),
      strides,
      padding_config,
      expr_to_typespec(ans)
    )
  end

  defp to_operator(
         :window_scatter_min,
         [%Value{} = arg, %Value{} = source, %Value{} = init_value, window_dimensions, opts],
         %{type: type} = ans,
         _state
       ) do
    padding_config = opts[:padding]
    strides = opts[:strides]

    arg = to_type(arg, type)
    source = to_type(source, type)
    init_value = to_type(init_value, type)

    Value.select_and_scatter(
      arg,
      source,
      init_value,
      :lt,
      Tuple.to_list(window_dimensions),
      strides,
      padding_config,
      expr_to_typespec(ans)
    )
  end

  defp to_operator(:indexed_add, [%Value{} | _] = tensors, out, _state) do
    mlir_scatter(tensors, out, :add)
  end

  defp to_operator(:indexed_put, [%Value{} | _] = tensors, out, _state) do
    mlir_scatter(tensors, out, :put)
  end

  defp to_operator(op, [arg, opts], ans, state) when op in [:argmax, :argmin] do
    apply(EXLA.Lib, op, [state.builder, arg, ans.type, opts])
  end

  defp to_operator(:clip, [%Value{} = operand, %Value{} = min, %Value{} = max], ans, _state) do
    min = to_type(min, ans.type)
    max = to_type(max, ans.type)
    operand = to_type(operand, ans.type)

    Value.clamp(operand, min, max, expr_to_typespec(ans))
  end

  defp to_operator(:slice, [%Value{} = tensor, start_indices, lengths, strides], ans, _state) do
    all_static? = Enum.all?(start_indices, &is_integer/1)

    if all_static? do
      limit_indices = Enum.zip_with(start_indices, lengths, fn i, len -> i + len end)
      Value.slice(tensor, start_indices, limit_indices, strides, expr_to_typespec(ans))
    else
      sample = Enum.find(start_indices, &(not is_integer(&1)))

      type =
        Enum.reduce(start_indices, op_type(sample), fn
          index, acc when is_integer(index) -> acc
          value, acc -> merge_type(op_type(value), acc)
        end)

      start_indices = Enum.map(start_indices, &to_type(&1, type))
      zeros = List.duplicate(0, tuple_size(ans.shape))

      typespec = Typespec.tensor(ans.type, List.to_tuple(lengths))
      slice = Value.dynamic_slice(tensor, start_indices, lengths, typespec)

      if Enum.all?(strides, &(&1 == 1)) do
        slice
      else
        Value.slice(slice, zeros, lengths, strides, expr_to_typespec(ans))
      end
    end
  end

  defp to_operator(:put_slice, [%Value{} = tensor, start_indices, slice], ans, _state) do
    tensor = to_type(tensor, ans.type)
    slice = to_type(slice, ans.type)
    Value.dynamic_update_slice(tensor, slice, start_indices, expr_to_typespec(ans))
  end

  defp to_operator(:gather, [%Value{} = tensor, indices, opts], ans, _state) do
    axes = Keyword.fetch!(opts, :axes)
    tensor_shape = op_shape(tensor)
    tensor_rank = tuple_size(tensor_shape)
    tensor_axes = axes_for_rank(tensor_rank)
    indices_rank = tuple_size(op_shape(indices))
    index_vector_dim = indices_rank - 1

    slice_sizes =
      for i <- tensor_axes do
        if i in axes, do: 1, else: elem(tensor_shape, i)
      end

    batch_size = tensor_rank - length(axes)
    offset_dims = count_up(batch_size, index_vector_dim)

    Value.gather(
      tensor,
      indices,
      index_vector_dim,
      slice_sizes,
      offset_dims,
      axes,
      axes,
      expr_to_typespec(ans)
    )
  end

  defp to_operator(:reverse, [%Value{} = tensor, axes], ans, _state) do
    Value.reverse(tensor, axes, expr_to_typespec(ans))
  end

  defp to_operator(:concatenate, [[%Value{} | _rest] = tensors, axis], ans, _state) do
    tensors = Enum.map(tensors, &to_type(&1, ans.type))
    Value.concatenate(tensors, axis, expr_to_typespec(ans))
  end

  defp to_operator(:stack, [[%Value{} | _rest] = tensors, axis], ans, _state) do
    reshape_typespec = Typespec.tensor(ans.type, put_elem(ans.shape, axis, 1))
    tensors = Enum.map(tensors, &(&1 |> to_type(ans.type) |> Value.reshape(reshape_typespec)))
    Value.concatenate(tensors, axis, expr_to_typespec(ans))
  end

  defp to_operator(:sort, [%Value{} = tensor, opts], ans, state) do
    dimension = opts[:axis]

    op =
      case opts[:direction] do
        :asc -> :less
        :desc -> :greater
      end

    arg_typespec = Typespec.tensor(ans.type, {})
    arg_typespecs = [arg_typespec, arg_typespec]

    comp = sort_computation(op, ans.type, arg_typespecs, state)

    Value.sort([tensor], comp, dimension, opts[:stable] == true, [expr_to_typespec(ans)]) |> hd()
  end

  defp to_operator(:argsort, [tensor, opts], ans, state) do
    dimension = opts[:axis]
    stable = opts[:stable] == true

    op =
      case opts[:direction] do
        :asc -> :less
        :desc -> :greater
      end

    type = op_type(tensor)

    value_typespec = Typespec.tensor(type, {})
    idx_typespec = Typespec.tensor(ans.type, {})
    arg_typespecs = [value_typespec, value_typespec, idx_typespec, idx_typespec]

    comp = sort_computation(op, type, arg_typespecs, state)

    EXLA.Lib.argsort(state.builder, tensor, dimension, stable, comp, ans.type)
  end

  defp fft(exla_op, input_type, output_type, pad_n \\ nil, [%Value{} = tensor, opts], ans) do
    fft_n = opts[:length]
    pad_n = pad_n || fft_n
    axis = opts[:axis]
    tensor = to_type(tensor, input_type)

    shape = op_shape(tensor)
    m = elem(shape, axis)

    tensor = fft_pad_or_slice(tensor, m, pad_n, axis, shape, input_type)

    last_axis = tuple_size(shape) - 1

    if axis != last_axis do
      permutation =
        Enum.map(0..last_axis, fn
          ^axis -> last_axis
          ^last_axis -> axis
          ax -> ax
        end)

      padded_shape = op_shape(tensor)

      {transposed_input_shape, _} =
        Nx.Shape.transpose(
          padded_shape,
          permutation,
          List.duplicate(nil, tuple_size(padded_shape))
        )

      transposed_input_typespec = Typespec.tensor(input_type, transposed_input_shape)

      {transposed_output_shape, _} = Nx.Shape.transpose(ans.shape, permutation, ans.names)
      transposed_output_typespec = Typespec.tensor(output_type, transposed_output_shape)

      tensor
      |> Value.transpose(permutation, transposed_input_typespec)
      |> exla_op.([fft_n], transposed_output_typespec)
      |> Value.transpose(permutation, expr_to_typespec(ans))
    else
      exla_op.(tensor, [fft_n], expr_to_typespec(ans))
    end
  end

  defp fft2(exla_op, [%Value{} = tensor, opts], %{type: type} = ans) do
    [l1, l2] = lengths = opts[:lengths]
    [ax1, ax2] = axes = opts[:axes]
    output_type = Nx.Type.to_complex(type)
    tensor = to_type(tensor, output_type)

    shape = op_shape(tensor)
    m1 = elem(shape, ax1)
    m2 = elem(shape, ax2)

    tensor = fft_pad_or_slice(tensor, m1, l1, ax1, shape, output_type)
    tensor = fft_pad_or_slice(tensor, m2, l2, ax2, op_shape(tensor), output_type)

    last_axis = tuple_size(shape) - 1
    penultimate_axis = last_axis - 1
    last_axes = [penultimate_axis, last_axis]

    if axes != last_axes do
      permutation =
        Enum.map(0..last_axis, fn
          ^ax1 -> penultimate_axis
          ^penultimate_axis -> ax1
          ^ax2 -> last_axis
          ^last_axis -> ax2
          ax -> ax
        end)

      {transposed_shape, _} = Nx.Shape.transpose(ans.shape, permutation, ans.names)
      transposed_typespec = Typespec.tensor(ans.type, transposed_shape)

      tensor
      |> Value.transpose(permutation, transposed_typespec)
      |> exla_op.(lengths, transposed_typespec)
      |> Value.transpose(permutation, expr_to_typespec(ans))
    else
      exla_op.(tensor, lengths, expr_to_typespec(ans))
    end
  end

  defp fft_pad_or_slice(%Value{function: builder} = tensor, m, n, axis, shape, output_type) do
    cond do
      m == n ->
        tensor

      m > n ->
        lengths =
          shape
          |> Tuple.insert_at(axis + 1, n)
          |> Tuple.delete_at(axis)
          |> Tuple.to_list()

        starts = List.duplicate(0, tuple_size(shape))
        strides = List.duplicate(1, tuple_size(shape))

        limit_indices = Enum.zip_with(starts, lengths, fn i, len -> i + len end)

        {_, shape} = Nx.Shape.slice(shape, starts, limit_indices, strides)
        typespec = Typespec.tensor(output_type, shape)
        Value.slice(tensor, starts, limit_indices, strides, typespec)

      m < n ->
        zero_value = if Nx.Type.complex?(output_type), do: Complex.new(0), else: 0

        zero =
          Value.constant(builder, [zero_value], Typespec.tensor(output_type, {}))

        padding_config =
          {0, 0, 0}
          |> List.duplicate(tuple_size(shape))
          |> List.replace_at(axis, {0, n - m, 0})

        shape = Nx.Shape.pad(shape, padding_config)
        typespec = Typespec.tensor(output_type, shape)
        Value.pad(tensor, zero, padding_config, typespec)
    end
  end

  defp mlir_scatter([target, indices, updates, opts], %{type: type} = ans, kind)
       when kind in [:add, :put] do
    target = to_type(target, type)
    updates = to_type(updates, type)
    update_rank = updates |> op_shape() |> tuple_size()
    update_axes = tl(axes_for_rank(update_rank))
    index_axes = Keyword.fetch!(opts, :axes)

    Value.scatter(
      target,
      indices,
      updates,
      kind,
      1,
      update_axes,
      index_axes,
      index_axes,
      expr_to_typespec(ans)
    )
  end

  ## Cache and hook helpers helpers

  defp no_token_cache(),
    do: %{__MODULE__ => Outfeed.empty()}

  defp new_cache(outfeed),
    do: %{__MODULE__ => outfeed}

  defp merge_outfeed(%{__MODULE__ => outfeed} = cache, %{__MODULE__ => new_outfeed}),
    do: %{cache | __MODULE__ => Outfeed.with_token(new_outfeed, outfeed.token)}

  defp reset_token(%{__MODULE__ => outfeed}, token),
    do: %{__MODULE__ => Outfeed.with_token(outfeed, token)}

  defp update_token(%{__MODULE__ => outfeed} = cache, token),
    do: %{cache | __MODULE__ => Outfeed.with_token(outfeed, token)}

  defp get_token(%{__MODULE__ => outfeed}), do: outfeed.token

  defp get_outfeed(%{__MODULE__ => value}), do: value

  defp put_outfeed(cache, outfeed), do: %{cache | __MODULE__ => outfeed}

  defp add_runtime_callback(cache, runtime_callback) do
    cache
    |> get_outfeed()
    |> Outfeed.add_runtime_callback(runtime_callback)
    |> then(&put_outfeed(cache, &1))
  end

  ## Computation helpers

  defp sort_computation(operator, type, arg_typespecs, %{
         builder: %EXLA.MLIR.Function{} = function
       }) do
    {region, [lhs, rhs | _]} = Function.push_region(function, arg_typespecs)

    typespec = Typespec.tensor({:pred, 8}, {})

    {lhs, rhs} =
      if Nx.Type.integer?(type) do
        {lhs, rhs}
      else
        {sort_computation_canonicalize_float(lhs), sort_computation_canonicalize_float(rhs)}
      end

    op = apply(Value, operator, [lhs, rhs, typespec, [total_order: true]])

    Value.return(function, [op])
    Function.pop_region(function)
    region
  end

  defp sort_computation_canonicalize_float(%Value{function: func} = op) do
    # Standardize the representation of NaNs (-NaN, NaN) and zeros (-0, 0).
    # See https://github.com/google/jax/blob/e81c82605f0e1813080cfe1037d043b27b38291d/jax/_src/lax/lax.py#L4248-L4253

    op_typespec = Value.get_typespec(op)

    zero = Value.constant(func, [0], Typespec.to_shape(op_typespec, {}))
    zeros = Value.constant(func, [0], op_typespec)
    nans = Value.constant(func, [:nan], op_typespec)

    pred_typespec = Typespec.tensor({:pred, 8}, {})
    op = Value.select(Value.equal(op, zero, pred_typespec), zeros, op, op_typespec)
    Value.select(Value.is_nan(op, pred_typespec), nans, op, op_typespec)
  end

  defp op_computation(
         op,
         arg_typespecs,
         %{builder: %EXLA.MLIR.Function{} = builder},
         prepare_args
       ) do
    {region, args} = Function.push_region(builder, arg_typespecs)
    op = apply(Value, op, prepare_args.(args) ++ [hd(arg_typespecs)])
    Value.return(builder, [op])
    Function.pop_region(builder)
    region
  end

  defp fun_computation(args, expr, type, %{builder: %Function{} = function} = state) do
    arg_typespecs =
      Enum.map(args, fn %{type: type, shape: shape} -> Typespec.tensor(type, shape) end)

    {region, mlir_args} = Function.push_region(function, arg_typespecs)

    arg_params = Enum.zip(args, mlir_args)

    params = Enum.flat_map(arg_params, &computation_arg_param/1)

    state = %{
      state
      | builder: function,
        params: Map.new(params),
        scope_ids: Tree.scope_ids(expr)
    }

    {res, _} = recur_composite(expr, state, no_token_cache())
    Value.return(function, Enum.map(res, &to_type(&1, type)))
    Function.pop_region(function)
    region
  end

  defp mlir_while_computation(expr, initial, type, state, cache) do
    arg_typespecs = Enum.map(List.flatten(initial), &Value.get_typespec/1)

    {region, args} = Function.push_region(state.builder, arg_typespecs)

    outer_token = get_token(cache)

    {inner_token, arg_params} =
      if outer_token do
        [arg_token | arg_params] = args
        {arg_token, arg_params}
      else
        {nil, args}
      end

    params = Enum.with_index(arg_params, &{&2, &1})

    state = %{
      state
      | params: Map.new(params),
        scope_ids: Tree.scope_ids(expr)
    }

    expr =
      if type == {:pred, 8} and expr.type == {:u, 8} do
        %{expr | type: {:pred, 8}}
      else
        expr
      end

    {res, comp_cache} = recur_composite(expr, & &1, state, reset_token(cache, inner_token))

    res =
      if type == :with_token do
        if outer_token do
          [get_token(comp_cache) | List.flatten(res)]
        else
          List.flatten(res)
        end
      else
        Enum.map(res, &to_type(&1, type))
      end

    Value.return(state.builder, res)
    Function.pop_region(state.builder)

    {region, merge_outfeed(cache, comp_cache)}
  end

  defp block_subfunction_description(%module{} = _) do
    module
    |> Atom.to_string()
    |> String.replace(".", "_")
  end

  defp block_computation(description, args, expr, %{builder: %Function{}} = state, cache) do
    %Function{module: module, name: name} = subbuilder(state.builder, description)

    arg_typespecs = Enum.map(args, &Value.get_typespec/1)
    out_typespecs = container_to_typespecs(expr)

    outer_token = get_token(cache)
    token_typespec = Typespec.token()

    {arg_typespecs, out_typespecs} =
      if outer_token do
        {[token_typespec | arg_typespecs], [token_typespec | out_typespecs]}
      else
        {arg_typespecs, out_typespecs}
      end

    function = EXLA.MLIR.Module.add_function(module, name, arg_typespecs, out_typespecs)
    args = EXLA.MLIR.Function.get_arguments(function)

    {inner_token, args} =
      if outer_token do
        [arg_token | args] = args
        {arg_token, args}
      else
        {nil, args}
      end

    params = Enum.with_index(args, fn param, i -> {i, param} end)

    state = %{
      state
      | builder: function,
        params: Map.new(params),
        scope_ids: Tree.scope_ids(expr)
    }

    {res, comp_cache} = recur_composite(expr, state, reset_token(cache, inner_token))

    if outer_token do
      Value.func_return(function, [get_token(comp_cache) | List.flatten(res)])
    else
      Value.func_return(function, List.flatten(res))
    end

    {function, merge_outfeed(cache, comp_cache)}
  end

  # The cache is built on top of call args because we need to handle pred/u8.
  defp computation_key(op, args) do
    keys =
      Enum.map(args, fn
        %Value{} = op ->
          %Typespec{type: type, shape: shape} = Value.get_typespec(op)
          {shape, type}

        opts ->
          opts
      end)

    {op, keys}
  end

  defp computation_arg_param({tuple, params}) when is_tuple(tuple) and is_list(params) do
    tuple
    |> Tuple.to_list()
    |> Enum.zip(params)
    |> Enum.flat_map(&computation_arg_param/1)
  end

  defp computation_arg_param({%T{data: %Expr{op: :parameter, args: [pos]}}, [param]}) do
    [{pos, param}]
  end

  defp computation_arg_param({%T{data: %Expr{op: :parameter, args: [pos]}}, param}) do
    [{pos, param}]
  end

  defp recur_composite(composite, state, cache) do
    recur_composite(composite, & &1, state, cache)
  end

  defp recur_composite(tuple, transform, state, cache) when is_tuple(tuple) do
    list = Tuple.to_list(tuple)

    if expr = full_tuple(list) do
      recur_composite(expr, transform, state, cache)
    else
      Enum.map_reduce(list, cache, &recur_composite(&1, transform, state, &2))
    end
  end

  defp recur_composite(%Value{} = op, transform, _state, cache) do
    {[transform.(op)], cache}
  end

  defp recur_composite(expr, transform, state, cache) do
    {op, cache} = recur_operator(expr, state, cache)

    result =
      if is_list(op) do
        Enum.map(op, transform)
      else
        [transform.(op)]
      end

    {result, cache}
  end

  # If each element of the tuple is just a reference to the parent expression,
  # discard the tuple elements and return the parent expression.
  defp full_tuple(list) do
    with [%T{data: %Expr{op: :elem, args: [%T{data: %Expr{id: id} = expr}, 0]}} | _] <- list,
         true <- full_tuple?(list, 0, id) do
      expr
    else
      _ -> nil
    end
  end

  defp full_tuple?([arg | args], index, id) do
    case arg do
      %T{data: %Expr{op: :elem, args: [%T{data: %Expr{id: ^id}, shape: shape}, ^index]}} ->
        if shape == {:tuple, index} do
          true
        else
          full_tuple?(args, index + 1, id)
        end

      _ ->
        false
    end
  end

  # We got until the end without traversing the whole tuple
  defp full_tuple?([], _index, _id), do: false

  ## Aggregation

  defp to_aggregate(op, type, shape, %Value{} = arg, initial, opts, state) do
    arg = to_type(arg, type)

    acc =
      case initial do
        %Value{} = initial ->
          initial

        initial when is_number(initial) ->
          Value.constant(state.builder, [initial], Typespec.tensor(type, {}))
      end

    args = [Typespec.tensor(type, {}), Typespec.tensor(type, {})]
    comp = op_computation(op, args, state, &Enum.reverse/1)

    keep_axes = opts[:keep_axes]
    reduce_axes = reduce_axes(arg, opts[:axes])

    typespec = Typespec.tensor(type, remove_axes(op_shape(arg), reduce_axes))
    [result] = Value.reduce(comp, [acc], [arg], reduce_axes, [typespec])

    if keep_axes do
      Value.reshape(result, Typespec.tensor(type, shape))
    else
      result
    end
  end

  defp to_window_aggregate(op, %{type: type} = ans, arg, initial, window_dimensions, opts, state) do
    arg = to_type(arg, type)

    acc =
      case initial do
        %Value{} = initial ->
          initial

        initial when is_number(initial) ->
          Value.constant(state.builder, [initial], Typespec.tensor(type, {}))
      end

    args = [Typespec.tensor(type, {}), Typespec.tensor(type, {})]
    # We reverse the argument order because :nan + :infinity
    # returns :nan but :infinity + :nan returns :infinity.
    # So we want to keep the current value as first argument
    # to preserve such properties.
    comp = op_computation(op, args, state, &Enum.reverse/1)

    strides = opts[:strides]
    padding = opts[:padding]
    window_dilations = opts[:window_dilations]

    [result] =
      Value.window_reduce(
        comp,
        [acc],
        [arg],
        Tuple.to_list(window_dimensions),
        strides,
        List.duplicate(1, tuple_size(op_shape(arg))),
        window_dilations,
        padding,
        [expr_to_typespec(ans)]
      )

    result
  end

  ## Cond

  defp to_if(pred, on_true, on_false, %{builder: %Function{}} = state, cache) do
    {pred_op, cache} = recur_operator(pred, state, cache) |> unwrap_single_tensor!()

    true_ids = Tree.scope_ids(on_true)
    false_ids = Tree.scope_ids(on_false)

    cache = recur_shared_ids(on_true, false_ids, state, cache)
    cache = recur_shared_ids(on_false, true_ids, state, cache)

    out_typespecs = container_to_typespecs(on_true)

    outer_token = get_token(cache)

    result_typespecs =
      if outer_token do
        [Typespec.token() | out_typespecs]
      else
        out_typespecs
      end

    {true_computation, cache} = to_mlir_if_branch(on_true, true_ids, state, cache)
    {false_computation, cache} = to_mlir_if_branch(on_false, false_ids, state, cache)
    if_results = Value.if_op(pred_op, true_computation, false_computation, result_typespecs)

    if outer_token do
      [token | results] = if_results
      {wrap_tuple_result(results, on_true), update_token(cache, token)}
    else
      {wrap_tuple_result(if_results, on_true), cache}
    end
  end

  defp recur_shared_ids(
         expr,
         other_ids,
         %{scope_ids: ids} = state,
         cache
       ) do
    {_, cache} =
      Composite.reduce(expr, {%{}, cache}, fn node, acc ->
        do_recur_shared_ids(node, state, acc, {ids, other_ids})
      end)

    cache
  end

  defp shared?(_id, :parameter, _args, _shared_ids),
    do: true

  # We never pass reference to tuples around, only through their elements,
  # so if a tuple is in a predicate, then it all must be in a predicate.
  defp shared?(_id, :elem, [%T{data: %Expr{id: tuple_id}}, _pos], {parent_ids, sibling_ids})
       when is_map_key(parent_ids, tuple_id) or is_map_key(sibling_ids, tuple_id),
       do: true

  defp shared?(id, _op, _args, {parent_ids, sibling_ids}),
    do: is_map_key(parent_ids, id) or is_map_key(sibling_ids, id)

  defp do_recur_shared_ids(
         %T{data: %Expr{id: id, op: op, args: args}} = expr,
         state,
         {visited, cache},
         shared_ids
       ) do
    cond do
      Map.has_key?(visited, id) ->
        {visited, cache}

      op == :constant or shared?(id, op, args, shared_ids) ->
        {_, cache} = recur_operator(expr, state, cache)
        {Map.put(visited, id, true), cache}

      true ->
        {_, {visited, cache}} =
          Tree.apply_args(
            expr,
            :scope,
            {visited, cache},
            &{&1, do_recur_shared_ids(&1, state, &2, shared_ids)}
          )

        {Map.put(visited, id, true), cache}
    end
  end

  defp to_mlir_if_branch(expr, current_ids, state, cache) do
    {region, []} = Function.push_region(state.builder, [])

    comp_state = %{state | scope_ids: current_ids}

    {res, res_cache} = recur_composite(expr, & &1, comp_state, cache)

    if token = get_token(cache) do
      Value.return(state.builder, [token | List.flatten(res)])
    else
      Value.return(state.builder, List.flatten(res))
    end

    Function.pop_region(state.builder)

    {region, merge_outfeed(cache, res_cache)}
  end

  ## Axes helpers

  defp broadcast_axes(left, right) do
    {min, max} = if left <= right, do: {left, right}, else: {right, left}
    min_size = tuple_size(min)
    max_size = tuple_size(max)

    # To reproduce Nx broadcast, we simply match the lower dimensions to the highest ones.
    count_up(min_size, max_size - min_size)
  end

  defp reduce_axes(op, axes) do
    if axes do
      Enum.sort(axes)
    else
      Nx.axes(op_shape(op))
    end
  end

  defp count_up(0, _n), do: []
  defp count_up(i, n), do: [n | count_up(i - 1, n + 1)]

  defp axes_for_rank(0), do: []

  defp axes_for_rank(rank) do
    Enum.to_list(0..(rank - 1))
  end

  ## Op Helpers

  defp op_type(%Value{} = op), do: Value.get_typespec(op).type

  defp op_shape(%Value{} = op), do: Value.get_typespec(op).shape

  defp to_type(%Value{} = op, type) do
    typespec = Value.get_typespec(op)

    if typespec.type == type do
      op
    else
      Value.convert(op, Typespec.to_type(typespec, type))
    end
  end

  defp merge_type({:pred, 8}, {:pred, 8}), do: {:pred, 8}
  defp merge_type(left, right), do: Nx.Type.merge(to_nx_type(left), to_nx_type(right))

  defp to_nx_type({:pred, 8}), do: {:u, 8}
  defp to_nx_type(type), do: type

  defp to_constant(%EXLA.MLIR.Function{} = function, constant, type) do
    Value.constant(function, [constant], Typespec.tensor(type, {}))
  end

  defp subbuilder(%EXLA.MLIR.Function{name: name} = function, description) do
    suffix = System.unique_integer([:positive])
    %{function | name: name <> "_" <> description <> "_" <> Integer.to_string(suffix)}
  end

  # Helpers

  defp apply_mlir_broadcasted_bin_op(op, out, left, right) do
    left_typespec = Value.get_typespec(left)
    right_typespec = Value.get_typespec(right)
    left_dims = broadcast_axes(left_typespec.shape, out.shape)
    right_dims = broadcast_axes(right_typespec.shape, out.shape)

    type = merge_type(left_typespec.type, right_typespec.type)
    type = merge_type(type, out.type)

    left = to_type(left, type)

    left =
      if left_typespec.shape == out.shape do
        left
      else
        Value.broadcast_in_dim(left, left_dims, Typespec.tensor(type, out.shape))
      end

    right = to_type(right, type)

    right =
      if right_typespec.shape == out.shape do
        right
      else
        Value.broadcast_in_dim(right, right_dims, Typespec.tensor(type, out.shape))
      end

    Value
    |> apply(op, [left, right, Typespec.tensor(type, out.shape)])
    |> to_type(out.type)
  end

  defp to_mlir_logical(%Value{} = value) do
    to_type(value, {:pred, 8})
  end

  defp container_to_typespecs(container) do
    [container]
    |> Nx.Defn.Composite.flatten_list()
    |> Enum.flat_map(fn
      %Nx.Tensor{type: {:tuple, _}, data: %{args: values}} ->
        Enum.flat_map(values, &container_to_typespecs/1)

      t ->
        [Typespec.tensor(t.type, t.shape)]
    end)
  end

  defp wrap_tuple_result(list, template) when is_tuple(template) do
    list
  end

  defp wrap_tuple_result(list, %Nx.Tensor{type: {:tuple, _}}) do
    list
  end

  defp wrap_tuple_result([value], _), do: value

  defp unwrap_single_tensor!({[%Value{} = op], cache}), do: {op, cache}
  defp unwrap_single_tensor!({%Value{} = op, cache}), do: {op, cache}

  defp remove_axes(shape, axes) do
    axes
    |> Enum.reverse()
    |> Enum.reduce(shape, &Tuple.delete_at(&2, &1))
  end

  defp expr_to_typespec(expr) do
    Typespec.tensor(expr.type, expr.shape)
  end

  defp calculate_unsharded_inputs(vars, input_shardings) do
    # We use only the first input list in the collection,
    # and we just assume they all have the same deep shapes.

    Enum.zip_with(vars, input_shardings, fn var, sharding ->
      Nx.Defn.Composite.traverse(var, fn
        %T{shape: shape} = t ->
          # TODO: we need to decide how vectorization interacts with input shardings
          updated_shape =
            shape |> Tuple.to_list() |> Enum.zip_with(sharding, &*/2) |> List.to_tuple()

          %{t | shape: updated_shape}

        var ->
          var
      end)
    end)
  end
end