Here’s my suggestion:

  def traverse([template_container | _] = list_of_arbitrary_nested_maps, fun) when is_list(list_of_arbitrary_nested_maps) do
    zipped_containers = 
      list_of_arbitrary_nested_maps
      |> Enum.map(&Nx.Defn.Composite.flatten_list[&1])
      |> Enum.zip_with(fun)
    
    {v, []} =
        Nx.Defn.Composite.traverse(template_container, zipped_containers, fn _, [h | t] ->
          {h, t}
        end)

   v
  end

Couldn’t find any so I hacked a quick one that fits my use case. Any feedback or alternatives is more than welcome as I am pretty new to Elixir.

defmodule TreeMap do
  import Nx.Defn
  import Nx, only: [is_tensor: 1]

  defn fun({a, b, c}) do
    Nx.multiply(c, Nx.add(a, b))
  end

  def process_list_of_arbitrary_nested_maps(list_of_arbitrary_nested_maps) do
    list_of_arbitrary_nested_maps
    |> Enum.map(&Map.to_list/1)
    |> Enum.map(&Enum.sort/1)
    |> Enum.zip()
    |> Enum.map(&Tuple.to_list/1)
    |> Enum.map(fn elem ->
      elem |> Enum.map_reduce("", fn x, _ -> {x |> elem(1), x |> elem(0)} end)
    end)
  end

  def traverse(list_of_arbitrary_nested_maps, fun) do
    condition =
      list_of_arbitrary_nested_maps
      |> Enum.map(&Map.values/1)
      |> Enum.map(fn x -> x |> Enum.all?(&is_tensor/1) end)
      |> Enum.all?()

    cond do
      condition ->
        list_of_arbitrary_nested_maps
        |> process_list_of_arbitrary_nested_maps()
        |> Enum.map(fn {l, k} -> %{"#{k}" => fun.(l |> List.to_tuple())} end)
        |> Enum.reduce(&Map.merge/2)

      true ->
        list_of_arbitrary_nested_maps
        |> process_list_of_arbitrary_nested_maps()
        |> Enum.map(fn {l, k} ->
          %{
            "#{k}" => traverse(l, &fun/1)
          }
        end)
        |> Enum.reduce(&Map.merge/2)
    end
  end
end

Maybe Nx.Defn.Composite.traverse?