diff --git a/nx/lib/nx.ex b/nx/lib/nx.ex index 93a58207ff..c27f010208 100644 --- a/nx/lib/nx.ex +++ b/nx/lib/nx.ex @@ -2020,7 +2020,6 @@ defmodule Nx do Creates a one-dimensional tensor from a `binary` with the given `type`. If the binary size does not match its type, an error is raised. - ## Examples iex> Nx.from_binary(<<1, 2, 3, 4>>, :s8) @@ -2049,7 +2048,7 @@ defmodule Nx do is ignored inside `defn` """ @doc type: :creation - def from_binary(binary, type, opts \\ []) when is_binary(binary) do + def from_binary(binary, type, opts \\ []) when is_bitstring(binary) do opts = keyword!(opts, [:backend]) {_, size} = type = Nx.Type.normalize!(type) dim = div(Kernel.bit_size(binary), size) diff --git a/nx/test/nx_test.exs b/nx/test/nx_test.exs index 0448a2d2a7..a5069aed37 100644 --- a/nx/test/nx_test.exs +++ b/nx/test/nx_test.exs @@ -1489,6 +1489,20 @@ defmodule NxTest do Nx.from_binary("", {:u, 32}) end) end + + test "round-trips sub-byte tensors whose binary is not byte-aligned" do + for {type, values} <- [ + {{:u, 2}, [1, 2, 3]}, + {{:u, 4}, [5, 10, 15]}, + {{:s, 2}, [-1, 0, 1]} + ] do + tensor = Nx.tensor(values, type: type) + data = Nx.to_binary(tensor) + + refute is_binary(data) + assert Nx.from_binary(data, type) == tensor + end + end end describe "to_batched/3" do