refactor: Nx.Random.key as deftransform (#1525)

Co-authored-by: Jonatan Kłosko <[email protected]>

Author: polvalente

nx/lib/nx/random.ex

@@ -80,22 +80,33 @@ defmodule Nx.Random do
         [0, 12]
       >
 
-  A single key effectively consists of 64 bits, so all possible values
-  of a 64-bit integer result in a different key. However, when passing
-  an integer literal, Nx implicitly assumes its type to be `:s32`,
-  which would result in an overflow for large integers. Therefore,
-  when dealing with large seeds, make sure to explicitly use a 64 bit
-  type:
-
-      iex> Nx.Random.key(Nx.u64(999999999999))
+      iex> Nx.Random.key(999999999999)
       #Nx.Tensor<
         u32[2]
         [232, 3567587327]
       >
   """
-  defn key(seed) do
-    seed = Nx.as_type(seed, :u64)
+  deftransform key(seed) do
+    seed =
+      case seed do
+        seed when is_integer(seed) ->
+          Nx.u64(seed)
 
+        %Nx.Tensor{} = seed when seed.type == {:u, 64} ->
+          seed
+
+        %Nx.Tensor{} = seed when seed.type == {:s, 64} ->
+          Nx.bitcast(seed, {:u, 64})
+
+        other ->
+          raise ArgumentError,
+                "expected seed to be an integer, u64 tensor or s64 tensor, got: #{inspect(other)}"
+      end
+
+    key_n(seed)
+  end
+
+  defnp key_n(seed) do
     k1 = Nx.right_shift(seed, 32)
     k2 = Nx.bitwise_and(seed, Nx.u64(0xFFFFFFFF))
 

nx/test/nx/random_test.exs

@@ -67,7 +67,7 @@ defmodule Nx.RandomTest do
   describe "distributions" do
     defp distribution_case(name, args: args, expected: expected) do
       seed = :erlang.adler32("#{name}threefry2x32")
-      key = Nx.Random.key(Nx.u64(seed))
+      key = Nx.Random.key(seed)
       actual = apply(Nx.Random, name, [key | args])
 
       assert_all_close(actual, expected)
@@ -261,7 +261,7 @@ defmodule Nx.RandomTest do
       |> assert_all_close(apply(expected_func, expected_args), rtol: 0.1)
 
       seed = :erlang.adler32("uniformthreefry2x32")
-      key = Nx.Random.key(Nx.tensor(seed, type: :u64))
+      key = Nx.Random.key(seed)
       t = apply(Nx.Random, name, [key | args])
 
       apply(Nx, moment, [t])

torchx/test/torchx/random_test.exs

@@ -55,7 +55,7 @@ defmodule Torchx.Nx.RandomTest do
   describe "distributions" do
     defp distribution_case(name, args: args, expected: expected) do
       seed = :erlang.adler32("#{name}threefry2x32")
-      key = Nx.Random.key(Nx.u64(seed))
+      key = Nx.Random.key(seed)
       actual = apply(Nx.Random, name, [key | args])
 
       assert_all_close(actual, expected)
@@ -209,7 +209,7 @@ defmodule Torchx.Nx.RandomTest do
       |> assert_all_close(apply(expected_func, expected_args), rtol: 0.1)
 
       seed = :erlang.adler32("#{name}threefry2x32")
-      key = Nx.Random.key(Nx.tensor(seed, type: :u64))
+      key = Nx.Random.key(seed)
       t = apply(Nx.Random, name, [key | args])
 
       apply(Nx, moment, [t])