Default integers to 32-bit precision

Author: jonatanklosko

exla/lib/exla.ex

@@ -220,7 +220,7 @@ defmodule EXLA do
 
       iex> EXLA.jit(&Nx.add(&1, &1)).(Nx.tensor([1, 2, 3]))
       #Nx.Tensor<
-        s64[3]
+        s32[3]
         [2, 4, 6]
       >
 
@@ -265,7 +265,7 @@ defmodule EXLA do
 
       iex> EXLA.jit_apply(&Nx.add(&1, &1), [Nx.tensor([1, 2, 3])])
       #Nx.Tensor<
-        s64[3]
+        s32[3]
         [2, 4, 6]
       >
 
@@ -278,10 +278,10 @@ defmodule EXLA do
   @doc """
   A shortcut for `Nx.Defn.compile/3` with the EXLA compiler.
 
-      iex> fun = EXLA.compile(&Nx.add(&1, &1), [Nx.template({3}, {:s, 64})])
+      iex> fun = EXLA.compile(&Nx.add(&1, &1), [Nx.template({3}, {:s, 32})])
       iex> fun.(Nx.tensor([1, 2, 3]))
       #Nx.Tensor<
-        s64[3]
+        s32[3]
         [2, 4, 6]
       >
 
@@ -328,7 +328,7 @@ defmodule EXLA do
 
   Now let's invoke it:
 
-      stream = EXLA.stream(&Streamed.sum/2, [Nx.template({}, {:s, 64}), 0])
+      stream = EXLA.stream(&Streamed.sum/2, [Nx.template({}, {:s, 32}), 0])
 
       for i <- 1..5 do
         Nx.Stream.send(stream, i)

exla/test/exla/backend_test.exs

@@ -32,9 +32,9 @@ defmodule EXLA.BackendTest do
 
   test "Nx.to_binary/1" do
     t = Nx.tensor([1, 2, 3, 4], backend: EXLA.Backend)
-    assert Nx.to_binary(t) == <<1::64-native, 2::64-native, 3::64-native, 4::64-native>>
-    assert Nx.to_binary(t, limit: 2) == <<1::64-native, 2::64-native>>
-    assert Nx.to_binary(t, limit: 6) == <<1::64-native, 2::64-native, 3::64-native, 4::64-native>>
+    assert Nx.to_binary(t) == <<1::32-native, 2::32-native, 3::32-native, 4::32-native>>
+    assert Nx.to_binary(t, limit: 2) == <<1::32-native, 2::32-native>>
+    assert Nx.to_binary(t, limit: 6) == <<1::32-native, 2::32-native, 3::32-native, 4::32-native>>
   end
 
   test "Nx.backend_transfer/1" do
@@ -44,7 +44,7 @@ defmodule EXLA.BackendTest do
     assert %EXLA.Backend{buffer: %EXLA.DeviceBuffer{}} = et.data
 
     nt = Nx.backend_transfer(et)
-    assert Nx.to_binary(nt) == <<1::64-native, 2::64-native, 3::64-native, 4::64-native>>
+    assert Nx.to_binary(nt) == <<1::32-native, 2::32-native, 3::32-native, 4::32-native>>
 
     assert_raise RuntimeError, ~r"called on deleted or donated buffer", fn ->
       Nx.backend_transfer(et)
@@ -63,7 +63,7 @@ defmodule EXLA.BackendTest do
     assert old_buffer == new_buffer
 
     nt = Nx.backend_transfer(et)
-    assert Nx.to_binary(nt) == <<1::64-native, 2::64-native, 3::64-native, 4::64-native>>
+    assert Nx.to_binary(nt) == <<1::32-native, 2::32-native, 3::32-native, 4::32-native>>
 
     assert_raise RuntimeError, ~r"called on deleted or donated buffer", fn ->
       Nx.backend_transfer(et)
@@ -83,10 +83,10 @@ defmodule EXLA.BackendTest do
       assert old_buffer != new_buffer
 
       nt = Nx.backend_copy(et)
-      assert Nx.to_binary(nt) == <<1::64-native, 2::64-native, 3::64-native, 4::64-native>>
+      assert Nx.to_binary(nt) == <<1::32-native, 2::32-native, 3::32-native, 4::32-native>>
 
       nt = Nx.backend_copy(et)
-      assert Nx.to_binary(nt) == <<1::64-native, 2::64-native, 3::64-native, 4::64-native>>
+      assert Nx.to_binary(nt) == <<1::32-native, 2::32-native, 3::32-native, 4::32-native>>
     end
 
     test "different clients" do
@@ -102,10 +102,10 @@ defmodule EXLA.BackendTest do
       assert new_buffer.device_id == 0
 
       nt = Nx.backend_copy(et)
-      assert Nx.to_binary(nt) == <<1::64-native, 2::64-native, 3::64-native, 4::64-native>>
+      assert Nx.to_binary(nt) == <<1::32-native, 2::32-native, 3::32-native, 4::32-native>>
 
       nt = Nx.backend_copy(et)
-      assert Nx.to_binary(nt) == <<1::64-native, 2::64-native, 3::64-native, 4::64-native>>
+      assert Nx.to_binary(nt) == <<1::32-native, 2::32-native, 3::32-native, 4::32-native>>
     end
   end
 
@@ -153,7 +153,7 @@ defmodule EXLA.BackendTest do
     assert inspect(t) ==
              """
              #Nx.Tensor<
-               s64[4]
+               s32[4]
                [1, 2, 3, 4]
              >\
              """

exla/test/exla/defn/expr_test.exs

@@ -1807,7 +1807,7 @@ defmodule EXLA.Defn.ExprTest do
       indices = Nx.tensor([[0]])
       updates = Nx.tensor([1])
 
-      assert_equal(indexed_add(target, indices, updates), Nx.tensor([1], type: {:s, 64}))
+      assert_equal(indexed_add(target, indices, updates), Nx.tensor([1], type: {:s, 32}))
 
       target = Nx.tensor([0])
       indices = Nx.tensor([[0]])
@@ -1879,7 +1879,7 @@ defmodule EXLA.Defn.ExprTest do
       indices = Nx.tensor([[0]])
       updates = Nx.tensor([1])
 
-      assert_equal(indexed_put(target, indices, updates), Nx.tensor([1], type: {:s, 64}))
+      assert_equal(indexed_put(target, indices, updates), Nx.tensor([1], type: {:s, 32}))
 
       target = Nx.tensor([0])
       indices = Nx.tensor([[0]])
@@ -1963,7 +1963,7 @@ defmodule EXLA.Defn.ExprTest do
     test "computes the sum across types" do
       assert_equal(Nx.tensor([1, 2, 3]) |> sum(), Nx.tensor(6))
       assert_equal(Nx.tensor([1, 2, 3], type: {:s, 8}) |> sum(), Nx.tensor(6))
-      assert_equal(Nx.tensor([1, 2, 3], type: {:u, 8}) |> sum(), Nx.tensor(6, type: {:u, 64}))
+      assert_equal(Nx.tensor([1, 2, 3], type: {:u, 8}) |> sum(), Nx.tensor(6, type: {:u, 32}))
       assert_equal(Nx.tensor([1.0, 2.0, 3.0]) |> sum(), Nx.tensor(6.0))
 
       assert_equal(
@@ -1986,9 +1986,9 @@ defmodule EXLA.Defn.ExprTest do
     defn sum_equal(t), do: Nx.sum(Nx.equal(t, 1.0))
 
     test "does not overflow" do
-      assert_equal(sum_equal(Nx.tensor(1)), Nx.tensor(1, type: {:u, 64}))
-      assert_equal(sum_equal(Nx.tensor([1, 1, 1])), Nx.tensor(3, type: {:u, 64}))
-      assert_equal(sum_equal(Nx.tensor([1, 2, 3])), Nx.tensor(1, type: {:u, 64}))
+      assert_equal(sum_equal(Nx.tensor(1)), Nx.tensor(1, type: {:u, 32}))
+      assert_equal(sum_equal(Nx.tensor([1, 1, 1])), Nx.tensor(3, type: {:u, 32}))
+      assert_equal(sum_equal(Nx.tensor([1, 2, 3])), Nx.tensor(1, type: {:u, 32}))
     end
 
     defn sum_keep(t), do: Nx.sum(t, keep_axes: true)
@@ -2011,7 +2011,7 @@ defmodule EXLA.Defn.ExprTest do
     test "computes the product across types" do
       assert_equal(Nx.tensor([1, 2, 3]) |> product(), Nx.tensor(6))
       assert_equal(Nx.tensor([1, 2, 3], type: {:s, 8}) |> product(), Nx.tensor(6))
-      assert_equal(Nx.tensor([1, 2, 3], type: {:u, 8}) |> product(), Nx.tensor(6, type: {:u, 64}))
+      assert_equal(Nx.tensor([1, 2, 3], type: {:u, 8}) |> product(), Nx.tensor(6, type: {:u, 32}))
       assert_equal(Nx.tensor([1.0, 2.0, 3.0]) |> product(), Nx.tensor(6.0))
 
       assert_equal(
@@ -2034,9 +2034,9 @@ defmodule EXLA.Defn.ExprTest do
     defn product_equal(t), do: Nx.product(Nx.equal(t, 1.0))
 
     test "does not overflow" do
-      assert_equal(product_equal(Nx.tensor(1)), Nx.tensor(1, type: {:u, 64}))
-      assert_equal(product_equal(Nx.tensor([1, 1, 1])), Nx.tensor(1, type: {:u, 64}))
-      assert_equal(product_equal(Nx.tensor([1, 2, 3])), Nx.tensor(0, type: {:u, 64}))
+      assert_equal(product_equal(Nx.tensor(1)), Nx.tensor(1, type: {:u, 32}))
+      assert_equal(product_equal(Nx.tensor([1, 1, 1])), Nx.tensor(1, type: {:u, 32}))
+      assert_equal(product_equal(Nx.tensor([1, 2, 3])), Nx.tensor(0, type: {:u, 32}))
     end
 
     defn product_keep(t), do: Nx.product(t, keep_axes: true)
@@ -2416,12 +2416,12 @@ defmodule EXLA.Defn.ExprTest do
         window_max2(Nx.tensor([[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]])),
         Nx.tensor([
           [
-            [-9_223_372_036_854_775_808, -9_223_372_036_854_775_808],
-            [-9_223_372_036_854_775_808, 6]
+            [-2_147_483_648, -2_147_483_648],
+            [-2_147_483_648, 6]
           ],
           [
-            [-9_223_372_036_854_775_808, -9_223_372_036_854_775_808],
-            [-9_223_372_036_854_775_808, 6]
+            [-2_147_483_648, -2_147_483_648],
+            [-2_147_483_648, 6]
           ]
         ])
       )
@@ -2482,12 +2482,12 @@ defmodule EXLA.Defn.ExprTest do
         window_min2(Nx.tensor([[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]])),
         Nx.tensor([
           [
-            [9_223_372_036_854_775_807, 9_223_372_036_854_775_807],
-            [9_223_372_036_854_775_807, 3]
+            [2_147_483_647, 2_147_483_647],
+            [2_147_483_647, 3]
           ],
           [
-            [9_223_372_036_854_775_807, 9_223_372_036_854_775_807],
-            [9_223_372_036_854_775_807, 3]
+            [2_147_483_647, 2_147_483_647],
+            [2_147_483_647, 3]
           ]
         ])
       )