Adding improvements to the Pallas code and a few other minor updates.

Author: botev

colabs/gsm8k_eval.ipynb

@@ -140,7 +140,6 @@
       "source": [
         "VARIANT = '2b' # @param ['2b', '2b-it'] {type:\"string\"}\n",
         "weights_dir = kagglehub.model_download(f'google/recurrentgemma/flax/{VARIANT}')\n",
-        "\n",
         "weights_dir = pathlib.Path(weights_dir)\n",
         "ckpt_path = weights_dir / VARIANT\n",
         "vocab_path = weights_dir / 'tokenizer.model'"

colabs/sampling_tutorial_jax.ipynb

@@ -255,7 +255,7 @@
         "id": "i15odZWiNpBn"
       },
       "source": [
-        "You should a description of the solar system.\n"
+        "You should get a description of the solar system.\n"
       ]
     }
   ],

recurrentgemma/jax/__init__.py

@@ -18,18 +18,17 @@
 from recurrentgemma.jax import griffin
 from recurrentgemma.jax import layers
 from recurrentgemma.jax import modules
-from recurrentgemma.jax import pallas
 from recurrentgemma.jax import sampler
+from recurrentgemma.jax import scan
 from recurrentgemma.jax import utils
 
 
 ScanType = common.ScanType
 TemporalBlockType = common.TemporalBlockType
 Preset = common.Preset
 GriffinConfig = common.GriffinConfig
-PallasShardingSpec = pallas.PallasShardingSpec
-sharded_lru = pallas.sharded_lru
-rnn_scan = layers.rnn_scan
+ShardingSpec = scan.ShardingSpec
+linear_scan = scan.linear_scan
 BlockDiagonalLinear = layers.BlockDiagonalLinear
 RGLRU = layers.RGLRU
 Conv1D = layers.Conv1D
@@ -47,9 +46,8 @@
 __all__ = (
     "ScanType",
     "TemporalBlockType",
-    "PallasShardingSpec",
-    "sharded_lru",
-    "rnn_scan",
+    "ShardingSpec",
+    "linear_scan",
     "BlockDiagonalLinear",
     "RGLRU",
     "Conv1D",