Remove Nx.map/2 (#1493)

Author: jonatanklosko

exla/lib/exla/backend.ex

@@ -353,7 +353,6 @@ defmodule EXLA.Backend do
       {:window_product, [:tensor, :shape, :opts], [:tensor]},
       {:window_max, [:tensor, :shape, :opts], [:tensor]},
       {:window_min, [:tensor, :shape, :opts], [:tensor]},
-      {:map, [:tensor, :opts, :fun], [:tensor]},
       {:sort, [:tensor, :opts], [:tensor]},
       {:argsort, [:tensor, :opts], [:tensor]},
       {:window_scatter_max, [:tensor, :source, :init_value, :window_dims, :opts],

exla/lib/exla/defn.ex

@@ -1207,11 +1207,6 @@ defmodule EXLA.Defn do
     mlir_scatter(tensors, out, :put)
   end
 
-  defp to_operator(:map, [%Value{} = arg, _opts, fun], ans, _state) do
-    arg = to_type(arg, ans.type)
-    Value.map(fun, [arg], Nx.axes(ans.shape), expr_to_typespec(ans))
-  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

exla/lib/exla/mlir/value.ex

@@ -660,24 +660,6 @@ defmodule EXLA.MLIR.Value do
     )
   end
 
-  def map(
-        %Region{ref: mapper},
-        [%Value{function: func} | _] = inputs,
-        dimensions,
-        typespec
-      ) do
-    result_types = typespecs_to_mlir_types([typespec])
-
-    attributes = [
-      dimensions: attr_array_i64_elements(dimensions)
-    ]
-
-    regions = [mapper]
-
-    op(func, "stablehlo.map", inputs, result_types, attributes: attributes, regions: regions)
-    |> one!()
-  end
-
   def if_op(
         %Value{function: func} = pred,
         %Region{ref: on_true},

exla/test/exla/defn/expr_test.exs

@@ -1499,64 +1499,23 @@ defmodule EXLA.Defn.ExprTest do
     end
   end
 
-  describe "map" do
-    defn map_plus(t), do: Nx.map(t, fn x -> x + 1 end)
-    defn map_equal(t), do: Nx.map(t, [type: {:f, 64}], fn x -> Nx.equal(x, 1) end)
-    defn map_exp(t), do: Nx.map(t, [type: {:f, 64}], fn x -> Nx.exp(x) end)
-
-    @tag :unsupported_64_bit_op
-    test "maps a function over the tensor" do
-      assert_equal(map_plus(Nx.tensor([[1, 2, 3], [4, 5, 6]])), Nx.tensor([[2, 3, 4], [5, 6, 7]]))
-    end
-
-    @tag :unsupported_64_bit_op
-    test "maps a function with an output type" do
-      assert_equal(
-        map_equal(Nx.tensor([[1, 2, 3], [4, 5, 6]])),
-        Nx.tensor([[1.0, 0.0, 0.0], [0.0, 0.0, 0.0]], type: {:f, 64})
-      )
-
-      assert_equal(
-        map_exp(Nx.tensor([[1, 2, 3], [4, 5, 6]])),
-        Nx.tensor(
-          [
-            [2.718281828459045, 7.38905609893065, 20.085536923187668],
-            [54.598150033144236, 148.4131591025766, 403.4287934927351]
-          ],
-          type: {:f, 64}
-        )
-      )
-    end
-
-    defn map_conditional(t), do: Nx.map(t, fn x -> if x > 0, do: x, else: -x end)
-
-    @tag :conditional_inside_map_reduce
-    @tag :unsupported_64_bit_op
-    test "maps a function with conditional" do
-      assert_equal(
-        map_conditional(Nx.tensor([-2, -1, 0, 1, 2])),
-        Nx.tensor([2, 1, 0, 1, 2])
-      )
-    end
-
-    defn while_inside_if(pred, x) do
-      if pred do
-        {x, _} =
-          while {x, i = 0}, i < 10 do
-            {x, i + 1}
-          end
+  defn while_inside_if(pred, x) do
+    if pred do
+      {x, _} =
+        while {x, i = 0}, i < 10 do
+          {x, i + 1}
+        end
 
-        x
-      else
-        x
-      end
+      x
+    else
+      x
     end
+  end
 
-    test "while inside if" do
-      assert %{a: a, b: b} = while_inside_if(1, %{a: 1, b: 2.0})
-      assert_all_close(a, 1)
-      assert_all_close(b, 2.0)
-    end
+  test "while inside if" do
+    assert %{a: a, b: b} = while_inside_if(1, %{a: 1, b: 2.0})
+    assert_all_close(a, 1)
+    assert_all_close(b, 2.0)
   end
 
   describe "reduce" do

nx/lib/nx.ex

@@ -11912,94 +11912,6 @@ defmodule Nx do
     end)
   end
 
-  @doc """
-  Maps the given scalar function over the entire
-  tensor.
-
-  The type of the returned tensor will be of the same type
-  as the input tensor, unless the `:type` option is given.
-  Therefore, you may need to explicitly cast the tensor to
-  avoid errors. For example, if you have an integer tensor
-  and you convert it to a float, as below, it will fail:
-
-      tensor = Nx.tensor([[1, 2, 3], [4, 5, 6]]),
-      Nx.map(tensor, fn x -> Nx.multiply(x, 1.0) end)
-
-  You need to explicitly pass the output type in such cases:
-
-      iex> tensor = Nx.tensor([[1, 2, 3], [4, 5, 6]])
-      iex> Nx.map(tensor, [type: :f32], fn x -> Nx.multiply(x, 1.0) end)
-      #Nx.Tensor<
-        f32[2][3]
-        [
-          [1.0, 2.0, 3.0],
-          [4.0, 5.0, 6.0]
-        ]
-      >
-
-  ## Limitations
-
-  Given this function relies on anonymous functions, it
-  may not be available or efficient on all Nx backends.
-  Therefore, you should avoid using `map/2` whenever possible
-  and use other functions in the `Nx` module to achieve the
-  desired result.
-
-  Inside `defn`, consider using `Nx.Defn.Kernel.while/4` instead.
-
-  ## Examples
-
-      iex> Nx.map(Nx.tensor([[1, 2, 3], [4, 5, 6]]), fn x -> Nx.add(x, 1) end)
-      #Nx.Tensor<
-        s64[2][3]
-        [
-          [2, 3, 4],
-          [5, 6, 7]
-        ]
-      >
-
-      iex> Nx.map(Nx.tensor(1), fn x -> Nx.add(x, 1) end)
-      #Nx.Tensor<
-        s64
-        2
-      >
-
-      iex> Nx.map(Nx.tensor([[1, 2, 3], [4, 5, 6]]), [type: :f64], fn x -> Nx.add(x, 1) end)
-      #Nx.Tensor<
-        f64[2][3]
-        [
-          [2.0, 3.0, 4.0],
-          [5.0, 6.0, 7.0]
-        ]
-      >
-
-  ## Vectorized tensors
-
-  `map/3` behaves the same as with non-vectorized tensors, applying
-  `fun` in an element-wise fashion.
-
-      iex> Nx.map(Nx.tensor([[1, 2, 3], [4, 5, 6]]) |> Nx.vectorize(:x), [type: :f64], &Nx.add(&1, 1))
-      #Nx.Tensor<
-        vectorized[x: 2]
-        f64[3]
-        [
-          [2.0, 3.0, 4.0],
-          [5.0, 6.0, 7.0]
-        ]
-      >
-  """
-  @doc type: :element
-  def map(tensor, opts \\ [], fun) do
-    apply_vectorized(tensor, fn tensor ->
-      %T{type: type} = tensor
-
-      opts = keyword!(opts, type: type)
-      output_type = Nx.Type.normalize!(opts[:type])
-      out = %{tensor | type: output_type}
-      impl!(tensor).map(out, tensor, opts, fun)
-    end)
-  end
-
   ## Matrix ops
 
   @doc """

nx/lib/nx/backend.ex

@@ -93,7 +93,6 @@ defmodule Nx.Backend do
   @callback window_product(out :: tensor, tensor, shape, keyword) :: tensor
   @callback window_max(out :: tensor, tensor, shape, keyword) :: tensor
   @callback window_min(out :: tensor, tensor, shape, keyword) :: tensor
-  @callback map(out :: tensor, tensor, keyword, fun) :: tensor
   @callback sort(out :: tensor, tensor, keyword) :: tensor
   @callback argsort(out :: tensor, tensor, keyword) :: tensor
   @callback window_scatter_max(out :: tensor, tensor, tensor, tensor, shape, keyword) :: tensor

nx/lib/nx/binary_backend.ex

@@ -1561,20 +1561,6 @@ defmodule Nx.BinaryBackend do
     window_reduce(out, tensor, init_value, window_dimensions, opts, fun)
   end
 
-  @impl true
-  def map(%{type: output_type} = out, %{type: {_, size}} = tensor, _opts, fun) do
-    data = to_binary(tensor)
-    template = %{tensor | shape: {}}
-
-    output_data =
-      for <<bin::size(size)-bitstring <- data>>, into: <<>> do
-        tensor = put_in(template.data.state, bin)
-        number_to_binary(scalar_to_number(fun.(tensor)), output_type)
-      end
-
-    from_binary(out, output_data)
-  end
-
   @impl true
   def window_scatter_max(out, tensor, source, init_value, window_dimensions, opts) do
     select_and_scatter(

nx/lib/nx/defn/expr.ex

@@ -985,13 +985,6 @@ defmodule Nx.Defn.Expr do
     expr(out, context, :window_reduce, [tensor, acc, window_dims, opts, fun])
   end
 
-  @impl true
-  def map(%{type: type} = out, tensor, opts, fun) do
-    args = [parameter(new_context(:map), type, {}, 0)]
-    %{data: %{context: context}} = tensor = to_expr(tensor)
-    expr(out, context, :map, [tensor, opts, apply_fun(context, fun, args, type)])
-  end
-
   @impl true
   def window_scatter_max(out, tensor, source, init_value, window_dims, opts) do
     {[tensor, source, init_value], context} = to_exprs([tensor, source, init_value])

nx/test/nx/defn/evaluator_test.exs

@@ -269,16 +269,6 @@ defmodule Nx.Defn.EvaluatorTest do
     test "calls external anonymous function via reduce" do
       assert calls_reduce_fun(&Nx.add/2, Nx.tensor([1, 2, 3])) == Nx.tensor(6)
     end
-
-    defn calls_map_fun(t) do
-      Nx.map(t, fn x ->
-        if Nx.equal(x, 0), do: 1, else: -x
-      end)
-    end
-
-    test "calls internal anonymous function via map" do
-      assert calls_map_fun(Nx.tensor([0, 1, 2])) == Nx.tensor([1, -1, -2])
-    end
   end
 
   describe "access" do
@@ -691,7 +681,7 @@ defmodule Nx.Defn.EvaluatorTest do
       t = Nx.iota({2, 3}, vectorized_axes: [a: 1], type: :s64)
 
       message = """
-      test/nx/defn/evaluator_test.exs:660: the do-block in while must return tensors with the same shape, type, and names as the initial arguments.
+      test/nx/defn/evaluator_test.exs:650: the do-block in while must return tensors with the same shape, type, and names as the initial arguments.
 
       {\e[32m
        <<<<< Body (do-block) <<<<<
@@ -723,7 +713,7 @@ defmodule Nx.Defn.EvaluatorTest do
 
       error =
         """
-        test/nx/defn/evaluator_test.exs:660: condition must be a scalar tensor, got: #Nx.Tensor<
+        test/nx/defn/evaluator_test.exs:650: condition must be a scalar tensor, got: #Nx.Tensor<
           vectorized[x: 1]
           u8[1]
         \s\s