From 0c5e913d623468e0415c9b1adda12a8e4966c0c5 Mon Sep 17 00:00:00 2001
From: Hiroki Furuta <hirokifuruta@google.com>
Date: Mon, 8 Jun 2026 19:57:19 -0700
Subject: [PATCH] Add sharding support to ChatSampler.

PiperOrigin-RevId: 928907308
---
 gemma/gm/text/_chat_sampler.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/gemma/gm/text/_chat_sampler.py b/gemma/gm/text/_chat_sampler.py
index 05710159..b230da6f 100644
--- a/gemma/gm/text/_chat_sampler.py
+++ b/gemma/gm/text/_chat_sampler.py
@@ -256,6 +256,7 @@ def _sample(
           return_state=True,
           last_state=last_state,
           stream=bool(stream),
+          sharding=sharding,
       )
 
   def chat(