From f5a25dfc615dbfc9904b8308d90b58d1d9528a8f Mon Sep 17 00:00:00 2001
From: Bradley Lewis Fargo <blewfargs@gmail.com>
Date: Mon, 17 Aug 2026 12:16:47 -0500
Subject: [PATCH 1/2] Accept sub-byte bitstrings in Nx.from_binary
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

to_binary's docs note that sub-byte types (u2/u4/s2/s4) may yield a
bitstring whose bit count is not divisible by 8 — but from_binary's
is_binary guard rejected exactly those values, so the documented round
trip crashed. The body already measures in bits (bit_size + rem
check); the guard becomes is_bitstring. Docs now note the bitstring
case, mirroring to_binary's existing info box.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 nx/lib/nx.ex        |  6 +++++-
 nx/test/nx_test.exs | 14 ++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/nx/lib/nx.ex b/nx/lib/nx.ex
index 93a58207ff..f687637cd6 100644
--- a/nx/lib/nx.ex
+++ b/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)
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

From b810ac166dd7635cc7339b281422f04cecd2b550 Mon Sep 17 00:00:00 2001
From: Paulo Valente <16843419+polvalente@users.noreply.github.com>
Date: Mon, 17 Aug 2026 14:22:53 -0300
Subject: [PATCH 2/2] Apply suggestion from @polvalente

---
 nx/lib/nx.ex | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/nx/lib/nx.ex b/nx/lib/nx.ex
index f687637cd6..c27f010208 100644
--- a/nx/lib/nx.ex
+++ b/nx/lib/nx.ex
@@ -2020,11 +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.
-
-  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)