Accept sub-byte bitstrings in Nx.from_binary

nx/lib/nx.ex

@@ -2021,6 +2021,10 @@ defmodule Nx do
 
   If the binary size does not match its type, an error is raised.
 
+  For sub-byte types (u2/u4/s2/s4), the input may also be a bitstring
+  whose bit count is not divisible by 8, as produced by `to_binary/2`
+  for those types.
+
   ## Examples
 
       iex> Nx.from_binary(<<1, 2, 3, 4>>, :s8)
@@ -2049,7 +2053,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)

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