diff --git a/activation_store/collect.py b/activation_store/collect.py
index 6512cde..89427fa 100644
--- a/activation_store/collect.py
+++ b/activation_store/collect.py
@@ -1,47 +1,112 @@
-from transformers import AutoModelForCausalLM
-import torch
-from datasets import Dataset
-from tqdm.auto import tqdm
-from torch.utils.data import DataLoader
-from loguru import logger
+import gc
 from pathlib import Path
-from baukit.nethook import TraceDict, recursive_copy
-from einops import rearrange
-from datasets.arrow_writer import ArrowWriter, ParquetWriter
-from datasets.fingerprint import Hasher
-from transformers.modeling_outputs import ModelOutput
-
-from activation_store.helpers.torch import clear_mem
 from typing import Dict, Generator
+import copy
+import torch
+from baukit.nethook import TraceDict, recursive_copy
+from datasets import Dataset
+from datasets.arrow_writer import ParquetWriter
+from datasets.fingerprint import Hasher
+from einops import rearrange
+from loguru import logger
 from torch import Tensor
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+from transformers import AutoModelForCausalLM
+from transformers.modeling_outputs import ModelOutput
 
 default_output_folder = (Path(__file__).parent.parent / "outputs").resolve()
 
-def default_postprocess_result(input: dict, trace: TraceDict, output: ModelOutput) -> Dict[str, Tensor]:
+
+def clear_mem():
+    gc.collect()
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+def default_postprocess_result(
+    input: dict, trace: TraceDict, output: ModelOutput, model: AutoModelForCausalLM
+) -> Dict[str, Tensor]:
     """add activations to output, and rearrange hidden states"""
 
     # Baukit records the literal layer output, which varies by model. Here we assume that the output or the first part are activations we want
-    acts = {f'act-{k}': 
-            v.output[0] if isinstance(v.output, tuple) else v.output
-            for k, v in trace.items()}
+    acts = {
+        f"act-{k}": v.output[0] if isinstance(v.output, tuple) else v.output
+        for k, v in trace.items()
+    }
+
+    output.hidden_states = rearrange(list(output.hidden_states), "l b t h -> b l t h")
+
+    o = dict(
+        attention_mask=input["attention_mask"],
+        **acts, **output
+    )
+    return o
+
+def to_cpu(x):
+    """
+    Trys to convert torch if possible a single item
+    """
+    if isinstance(x, torch.Tensor):
+        x = x.cpu()
+        return x
+    else:
+        return x
+
+def recursive_copy2(x, clone=None, detach=None, retain_grad=None):
+    """
+    from baukit with addition of deep copy for non tensors
     
-    output.hidden_states = rearrange(list(output.hidden_states), 'l b t h -> b l t h')
-
-    return dict(**acts, **output)
-
+    Copies a reference to a tensor, or an object that contains tensors,
+    optionally detaching and cloning the tensor(s).  If retain_grad is
+    true, the original tensors are marked to have grads retained.
+    """
+    if not clone and not detach and not retain_grad:
+        return x
+    if isinstance(x, torch.Tensor):
+        if retain_grad:
+            if not x.requires_grad:
+                x.requires_grad = True
+            x.retain_grad()
+        elif detach:
+            x = x.detach()
+        if clone:
+            x = x.clone()
+        return x
+    # Only dicts, lists, and tuples (and subclasses) can be copied.
+    if isinstance(x, dict):
+        return type(x)({k: recursive_copy(v, clone=clone, detach=detach, retain_grad=retain_grad) for k, v in x.items()})
+    elif isinstance(x, (list, tuple)):
+        return type(x)([recursive_copy(v, clone=clone, detach=detach, retain_grad=retain_grad) for v in x])
+    else:
+        return copy.deepcopy(x)
 
 @torch.no_grad
-def generate_batches(loader: DataLoader, model: AutoModelForCausalLM, layers, postprocess_result=default_postprocess_result) -> Generator[Dict[str, Tensor], None, None]:
+def generate_batches(
+    loader: DataLoader,
+    model: AutoModelForCausalLM,
+    layers,
+    postprocess_result=default_postprocess_result,
+) -> Generator[Dict[str, Tensor], None, None]:
     model.eval()
-    for batch in tqdm(loader, 'collecting activations'):
+    for batch in tqdm(loader, "collecting activations"):
         device = next(model.parameters()).device
         with torch.amp.autocast(device_type=device.type):
             with TraceDict(model, layers) as trace:
-                out = model(**batch, use_cache=False, output_hidden_states=True, return_dict=True)
-        o = postprocess_result(batch, trace, out)
+                out = model(
+                    **batch,
+                    use_cache=False,
+                    output_hidden_states=True,
+                    return_dict=True,
+                )
+        o = postprocess_result(batch, trace, out, model)
 
         # copy to avoid memory leaks
-        o = recursive_copy(o)
+        for k in o:
+            if not isinstance(o[k], torch.Tensor):
+                print('o', k, type(o[k]))
+        o = {k: to_cpu(v) for k, v in o.items()}
+        o = recursive_copy(o, clone=True, detach=True)
         out = trace = batch = None
         clear_mem()
         yield o
@@ -52,7 +117,17 @@ def dataset_hash(**kwargs):
     return suffix
 
 
-def activation_store(loader: DataLoader, model: AutoModelForCausalLM, dataset_name='', layers=[], dataset_dir=default_output_folder, writer_batch_size=1, postprocess_result=default_postprocess_result) -> Dataset:
+def activation_store(
+    loader: DataLoader,
+    model: AutoModelForCausalLM,
+    dataset_name="",
+    layers=[],
+    dataset_dir=default_output_folder,
+    writer_batch_size=1,
+    postprocess_result=default_postprocess_result,
+    features=None,
+    schema=None,
+) -> Dataset:
     """
     Collect activations from a model and store them in a dataset
 
@@ -76,20 +151,27 @@ def activation_store(loader: DataLoader, model: AutoModelForCausalLM, dataset_na
     f.parent.mkdir(exist_ok=True, parents=True)
     logger.info(f"creating dataset {f}")
 
-    iterator = generate_batches(loader, model, layers=layers, postprocess_result=postprocess_result) 
+    iterator = generate_batches(
+        loader, model, layers=layers, postprocess_result=postprocess_result
+    )
 
-    with ParquetWriter(path=f, writer_batch_size=writer_batch_size,
-                       embed_local_files=True
-                       ) as writer: 
+    # batch_1 = next(iterator)
+    # Features.encode_batch(batch_1)
+    # features = Features({'x': Array2D(shape=(1, 3), dtype='int32')})
+
+    with ParquetWriter(
+        path=f, writer_batch_size=writer_batch_size, embed_local_files=True,
+         features=features, schema=schema,
+    ) as writer:
+        # writer.write_batch(batch_1)
         for bo in iterator:
-
             bs = len(next(iter(bo.values())))
-            assert all(len(v) == bs for v in bo.values()), f"must return Dict[str,Tensor] and all tensors with same batch size a first dimension"
-
-            # or maybe better compression to `writer.write(example, key)` for each
+            assert all(len(v) == bs for v in bo.values()), (
+                "must return Dict[str,Tensor] and all tensors with same batch size a first dimension"
+            )
             writer.write_batch(bo)
-        writer.finalize() 
+        writer.finalize()
         writer.close()
-    
+
     # ds = Dataset.from_file(str(f)).with_format("torch")
     return f
diff --git a/activation_store/helpers/torch.py b/activation_store/helpers/torch.py
deleted file mode 100644
index e649b37..0000000
--- a/activation_store/helpers/torch.py
+++ /dev/null
@@ -1,8 +0,0 @@
-import torch
-import gc
-
-
-def clear_mem():
-    gc.collect()
-    torch.cuda.empty_cache()
-    gc.collect()
diff --git a/nbs/example.ipynb b/nbs/example.ipynb
index 2afe528..2037ee5 100644
--- a/nbs/example.ipynb
+++ b/nbs/example.ipynb
@@ -18,10 +18,10 @@
    "source": [
     "from datasets import load_dataset\n",
     "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
+    "from datasets import Dataset\n",
+    "import torch\n",
     "\n",
-    "from activation_store.collect import activation_store\n",
-    "\n",
-    "import torch"
+    "from activation_store.collect import activation_store\n"
    ]
   },
   {
@@ -64,8 +64,8 @@
      "data": {
       "text/plain": [
        "Dataset({\n",
-       "    features: ['input_ids', 'attention_mask'],\n",
-       "    num_rows: 20\n",
+       "    features: ['attention_mask', 'input_ids'],\n",
+       "    num_rows: 10\n",
        "})"
       ]
      },
@@ -75,8 +75,8 @@
     }
    ],
    "source": [
-    "N = 20\n",
-    "max_length = 256\n",
+    "N = 10\n",
+    "max_length = 128\n",
     "\n",
     "imdb = load_dataset('wassname/imdb_dpo', split=f'test[:{N}]', keep_in_memory=False)\n",
     "\n",
@@ -110,7 +110,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "<torch.utils.data.dataloader.DataLoader object at 0x76557465f770>\n"
+      "<torch.utils.data.dataloader.DataLoader object at 0x7089fb69f6e0>\n"
      ]
     }
    ],
@@ -136,7 +136,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [
     {
@@ -168,7 +168,7 @@
        " 'model.layers.23.mlp.down_proj']"
       ]
      },
-     "execution_count": 7,
+     "execution_count": 6,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -181,25 +181,25 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "\u001b[32m2025-02-15 21:58:37.654\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mactivation_store.collect\u001b[0m:\u001b[36mactivation_store\u001b[0m:\u001b[36m70\u001b[0m - \u001b[1mcreating dataset /media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__9b3f4b0da96e9ad5.parquet\u001b[0m\n"
+      "\u001b[32m2025-02-16 09:16:55.292\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mactivation_store.collect\u001b[0m:\u001b[36mactivation_store\u001b[0m:\u001b[36m122\u001b[0m - \u001b[1mcreating dataset /media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__4a18b59a7867ed48.parquet\u001b[0m\n"
      ]
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "fe95a697e5c0432e85d15707b07fd001",
+       "model_id": "90a9936ab9f94893a77fc79bf972a04f",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "collecting activations:   0%|          | 0/5 [00:00<?, ?it/s]"
+       "collecting activations:   0%|          | 0/3 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -213,15 +213,14 @@
      ]
     },
     {
-     "ename": "NameError",
-     "evalue": "name 'ds_a' is not defined",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[8], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m f \u001b[38;5;241m=\u001b[39m activation_store(ds, model, layers\u001b[38;5;241m=\u001b[39mlayers, writer_batch_size\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m10\u001b[39m)\n\u001b[0;32m----> 2\u001b[0m \u001b[43mds_a\u001b[49m\n",
-      "\u001b[0;31mNameError\u001b[0m: name 'ds_a' is not defined"
-     ]
+     "data": {
+      "text/plain": [
+       "PosixPath('/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__4a18b59a7867ed48.parquet')"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
     }
    ],
    "source": [
@@ -233,11 +232,25 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
    "outputs": [
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1e1429e1d3224a2b8a5398f7a414911d",
+       "model_id": "06fafa5231674f4da16d4ddfab520bd7",
        "version_major": 2,
        "version_minor": 0
       },
@@ -252,19 +265,352 @@
      "data": {
       "text/plain": [
        "Dataset({\n",
-       "    features: ['act-model.layers.0.mlp.down_proj', 'act-model.layers.1.mlp.down_proj', 'act-model.layers.2.mlp.down_proj', 'act-model.layers.3.mlp.down_proj', 'act-model.layers.4.mlp.down_proj', 'act-model.layers.5.mlp.down_proj', 'act-model.layers.6.mlp.down_proj', 'act-model.layers.7.mlp.down_proj', 'act-model.layers.8.mlp.down_proj', 'act-model.layers.9.mlp.down_proj', 'act-model.layers.10.mlp.down_proj', 'act-model.layers.11.mlp.down_proj', 'act-model.layers.12.mlp.down_proj', 'act-model.layers.13.mlp.down_proj', 'act-model.layers.14.mlp.down_proj', 'act-model.layers.15.mlp.down_proj', 'act-model.layers.16.mlp.down_proj', 'act-model.layers.17.mlp.down_proj', 'act-model.layers.18.mlp.down_proj', 'act-model.layers.19.mlp.down_proj', 'act-model.layers.20.mlp.down_proj', 'act-model.layers.21.mlp.down_proj', 'act-model.layers.22.mlp.down_proj', 'act-model.layers.23.mlp.down_proj', 'logits', 'hidden_states'],\n",
-       "    num_rows: 20\n",
+       "    features: ['attention_mask', 'act-model.layers.0.mlp.down_proj', 'act-model.layers.1.mlp.down_proj', 'act-model.layers.2.mlp.down_proj', 'act-model.layers.3.mlp.down_proj', 'act-model.layers.4.mlp.down_proj', 'act-model.layers.5.mlp.down_proj', 'act-model.layers.6.mlp.down_proj', 'act-model.layers.7.mlp.down_proj', 'act-model.layers.8.mlp.down_proj', 'act-model.layers.9.mlp.down_proj', 'act-model.layers.10.mlp.down_proj', 'act-model.layers.11.mlp.down_proj', 'act-model.layers.12.mlp.down_proj', 'act-model.layers.13.mlp.down_proj', 'act-model.layers.14.mlp.down_proj', 'act-model.layers.15.mlp.down_proj', 'act-model.layers.16.mlp.down_proj', 'act-model.layers.17.mlp.down_proj', 'act-model.layers.18.mlp.down_proj', 'act-model.layers.19.mlp.down_proj', 'act-model.layers.20.mlp.down_proj', 'act-model.layers.21.mlp.down_proj', 'act-model.layers.22.mlp.down_proj', 'act-model.layers.23.mlp.down_proj', 'logits', 'hidden_states'],\n",
+       "    num_rows: 10\n",
        "})"
       ]
      },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# load\n",
+    "ds_a = Dataset.from_parquet(str(f)).with_format(\"torch\")\n",
+    "ds_a"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "DatasetInfo(description='', citation='', homepage='', license='', features={'attention_mask': Sequence(feature=Value(dtype='int8', id=None), length=-1, id=None), 'act-model.layers.0.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.1.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.2.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.3.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.4.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.5.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.6.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.7.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.8.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.9.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.10.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.11.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.12.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.13.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.14.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.15.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.16.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.17.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.18.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.19.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.20.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.21.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.22.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'act-model.layers.23.mlp.down_proj': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'logits': Sequence(feature=Sequence(feature=Value(dtype='float16', id=None), length=-1, id=None), length=-1, id=None), 'hidden_states': Sequence(feature=Sequence(feature=Sequence(feature=Value(dtype='float32', id=None), length=-1, id=None), length=-1, id=None), length=-1, id=None)}, post_processed=None, supervised_keys=None, builder_name='parquet', dataset_name='parquet', config_name='default', version=0.0.0, splits={'train': SplitInfo(name='train', num_bytes=1391398926, num_examples=10, shard_lengths=[4, 6], dataset_name='parquet')}, download_checksums={'/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__4a18b59a7867ed48.parquet': {'num_bytes': 1363203837, 'checksum': None}}, download_size=1363203837, post_processing_size=None, dataset_size=1391398926, size_in_bytes=2754602763)"
+      ]
+     },
      "execution_count": 10,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "from datasets import Dataset\n",
-    "Dataset.from_parquet(str(f))"
+    "ds_a.info"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([2, 25, 453, 896])"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "ds_a[0:2]['hidden_states'].shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([2, 453, 896])"
+      ]
+     },
+     "execution_count": 12,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "ds_a[0:2]['act-model.layers.0.mlp.down_proj'].shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "ZeroDivisionError",
+     "evalue": "division by zero",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mZeroDivisionError\u001b[0m                         Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[9], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;241;43m1\u001b[39;49m\u001b[38;5;241;43m/\u001b[39;49m\u001b[38;5;241;43m0\u001b[39;49m\n",
+      "\u001b[0;31mZeroDivisionError\u001b[0m: division by zero"
+     ]
+    }
+   ],
+   "source": [
+    "1/0"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## With dtypes compression - wip"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "lost 2.02%\n"
+     ]
+    }
+   ],
+   "source": [
+    "def float_to_int8(x: torch.Tensor) -> torch.Tensor:\n",
+    "    \"\"\"Converts a floating point tensor to float16, then reinterprets as int16.\"\"\"\n",
+    "    downcast = x.type(torch.float8_e4m3fn)\n",
+    "    # if not downcast.isfinite().all():\n",
+    "    #     raise ValueError(\"Cannot convert to 16 bit: values are not finite\")\n",
+    "\n",
+    "    return downcast.view(torch.int8)\n",
+    "\n",
+    "def int8_to_float32(x: torch.Tensor) -> torch.Tensor:\n",
+    "    \"\"\"Converts int16 to float16, then reinterprets as float32.\"\"\"\n",
+    "    return x.view(torch.float8_e4m3fn).type(torch.float32)\n",
+    "\n",
+    "\n",
+    "x = torch.randn(2, 3, 4)\n",
+    "x2 = float_to_int8(x)\n",
+    "x3 = int8_to_float32(x2)\n",
+    "assert torch.isfinite(x3).all()\n",
+    "assert torch.allclose(x, x3, rtol=1e-1)\n",
+    "d = ((x-x3)/x).abs().mean()\n",
+    "print(f'lost {d:.2%}')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from activation_store.collect import default_postprocess_result\n",
+    "from datasets.features.features import cast_to_python_objects\n",
+    "# o = cast_to_python_objects(o, only_1d_for_numpy=True, optimize_list_casting=False)\n",
+    "\n",
+    "def float8_postprocess_result(\n",
+    "    input, trace, output, model\n",
+    "):\n",
+    "    o = default_postprocess_result(input, trace, output, model)\n",
+    "    # o = cast_to_python_objects(o, only_1d_for_numpy=False, optimize_list_casting=False)\n",
+    "\n",
+    "    for k, v in o.items():\n",
+    "        if k=='attention_mask':\n",
+    "            o[k] = v.to(torch.int8)\n",
+    "        if isinstance(v, torch.Tensor) and torch.is_floating_point(v):\n",
+    "            print(k, v.dtype, v.shape, 'to int8')\n",
+    "            o[k] = float_to_int8(v.float())\n",
+    "        else:\n",
+    "            print('no conv', k, type(v))\n",
+    "    # o = {k: float_to_int8(v) if isinstance(v, torch.Tensor) else v\n",
+    "        #   for k, v in o.items()}\n",
+    "    return o"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'attention_mask': Sequence(feature=Value(dtype='int8', id=None), length=-1, id=None),\n",
+       " 'act-model.layers.0.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.1.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.2.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.3.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.4.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.5.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.6.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.7.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.8.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.9.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.10.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.11.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.12.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.13.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.14.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.15.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.16.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.17.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.18.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.19.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.20.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.21.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.22.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'act-model.layers.23.mlp.down_proj': Array3D(shape=(-1, 453, 896), dtype='int8', id=None),\n",
+       " 'logits': Array3D(shape=(-1, 453, 151936), dtype='int8', id=None),\n",
+       " 'hidden_states': Array4D(shape=(-1, 25, 453, 896), dtype='int8', id=None)}"
+      ]
+     },
+     "execution_count": 28,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from datasets.arrow_writer import OptimizedTypedSequence, _ArrayXDExtensionType\n",
+    "from datasets.features.features import Features, Array2D, Array3D, Array4D, Array5D\n",
+    "\n",
+    "# manually build features\n",
+    "optimized_int_type_by_col = {\n",
+    "    \"attention_mask\": \"int8\",  # binary tensor\n",
+    "    \"special_tokens_mask\": \"int8\",\n",
+    "    \"input_ids\": \"int32\",  # typical vocab size: 0-50k (max ~500k, never > 1M)\n",
+    "    \"token_type_ids\": \"int8\",  # binary mask; some (XLNetModel) use an additional token represented by a 2\n",
+    "}\n",
+    "\n",
+    "def build_schema(d):\n",
+    "    inferred_features = Features()\n",
+    "    cols = d.keys()\n",
+    "    for col in cols:\n",
+    "        x = d[col]\n",
+    "        if col in optimized_int_type_by_col:\n",
+    "            dtype = optimized_int_type_by_col[col]\n",
+    "            typed_sequence = OptimizedTypedSequence(x, col=col)\n",
+    "            inferred_features[col] = typed_sequence.get_inferred_type()\n",
+    "        else:\n",
+    "            if x.ndim == 1:\n",
+    "                inferred_features[col] = OptimizedTypedSequence(x, col=col)\n",
+    "                inferred_features[col] = typed_sequence.get_inferred_type()\n",
+    "            shape=(-1,)+x.shape[1:]\n",
+    "            dtype = 'int8' if x.dtype == torch.float32 else x.dtype\n",
+    "            if x.ndim == 2:\n",
+    "                cls = Array2D\n",
+    "            elif x.ndim == 3:\n",
+    "                cls = Array3D\n",
+    "            elif x.ndim == 4:\n",
+    "                cls = Array4D\n",
+    "            elif x.ndim == 5:\n",
+    "                cls = Array5D\n",
+    "            else:\n",
+    "                raise ValueError(f\"Unsupported number of dimensions: {x.ndim}\")\n",
+    "            inferred_features[col] = cls(dtype=dtype, shape=shape)\n",
+    "    return inferred_features.arrow_schema\n",
+    "    # Features.from_arrow_schema(schema)\n",
+    "\n",
+    "d = ds_a[0:2]\n",
+    "schema = build_schema(d)\n",
+    "schema\n",
+    "Features.from_arrow_schema(schema)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-02-16 09:25:08.798\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mactivation_store.collect\u001b[0m:\u001b[36mactivation_store\u001b[0m:\u001b[36m152\u001b[0m - \u001b[1mcreating dataset /media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__c6184d05bf03be61.parquet\u001b[0m\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "6fc1c814733a4a7ab65468f0d2ad0b2b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "collecting activations:   0%|          | 0/3 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "no conv attention_mask <class 'torch.Tensor'>\n",
+      "act-model.layers.0.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.1.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.2.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.3.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.4.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.5.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.6.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.7.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.8.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.9.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.10.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.11.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.12.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.13.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.14.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.15.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.16.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.17.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.18.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.19.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.20.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.21.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.22.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "act-model.layers.23.mlp.down_proj torch.float16 torch.Size([4, 453, 896]) to int8\n",
+      "logits torch.float16 torch.Size([4, 453, 151936]) to int8\n",
+      "hidden_states torch.float32 torch.Size([4, 25, 453, 896]) to int8\n"
+     ]
+    },
+    {
+     "ename": "ArrowTypeError",
+     "evalue": "Could not convert tensor([[  48,   25,  -83,  ...,   31,   45,   41],\n        [ -76, -100,  -94,  ...,   26,  -84, -117],\n        [ -97,   26,   15,  ...,  -97, -107, -109],\n        ...,\n        [  44,  -94, -104,  ..., -110,   18,   27],\n        [ -77,   26,  -77,  ..., -100,   33,   43],\n        [ -98,   22, -111,  ..., -110,   14, -107]], dtype=torch.int8) with type Tensor: was not a sequence or recognized null for conversion to list type",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mArrowTypeError\u001b[0m                            Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[29], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m f2 \u001b[38;5;241m=\u001b[39m \u001b[43mactivation_store\u001b[49m\u001b[43m(\u001b[49m\u001b[43mds\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlayers\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mlayers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mwriter_batch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m10\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[1;32m      2\u001b[0m \u001b[43m                      \u001b[49m\u001b[43mschema\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mschema\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      3\u001b[0m \u001b[43m                        \u001b[49m\u001b[43mpostprocess_result\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfloat8_postprocess_result\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      4\u001b[0m f2\n\u001b[1;32m      5\u001b[0m ds_a2 \u001b[38;5;241m=\u001b[39m Dataset\u001b[38;5;241m.\u001b[39mfrom_parquet(\u001b[38;5;28mstr\u001b[39m(f2))\u001b[38;5;241m.\u001b[39mwith_format(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtorch\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
+      "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/activation_store/collect.py:172\u001b[0m, in \u001b[0;36mactivation_store\u001b[0;34m(loader, model, dataset_name, layers, dataset_dir, writer_batch_size, postprocess_result, features, schema)\u001b[0m\n\u001b[1;32m    168\u001b[0m     bs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlen\u001b[39m(\u001b[38;5;28mnext\u001b[39m(\u001b[38;5;28miter\u001b[39m(bo\u001b[38;5;241m.\u001b[39mvalues())))\n\u001b[1;32m    169\u001b[0m     \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mall\u001b[39m(\u001b[38;5;28mlen\u001b[39m(v) \u001b[38;5;241m==\u001b[39m bs \u001b[38;5;28;01mfor\u001b[39;00m v \u001b[38;5;129;01min\u001b[39;00m bo\u001b[38;5;241m.\u001b[39mvalues()), (\n\u001b[1;32m    170\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmust return Dict[str,Tensor] and all tensors with same batch size a first dimension\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    171\u001b[0m     )\n\u001b[0;32m--> 172\u001b[0m     \u001b[43mwriter\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwrite_batch\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbo\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    173\u001b[0m writer\u001b[38;5;241m.\u001b[39mfinalize()\n\u001b[1;32m    174\u001b[0m writer\u001b[38;5;241m.\u001b[39mclose()\n",
+      "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/datasets/arrow_writer.py:605\u001b[0m, in \u001b[0;36mArrowWriter.write_batch\u001b[0;34m(self, batch_examples, writer_batch_size)\u001b[0m\n\u001b[1;32m    603\u001b[0m         col_try_type \u001b[38;5;241m=\u001b[39m try_features[col] \u001b[38;5;28;01mif\u001b[39;00m try_features \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m col \u001b[38;5;129;01min\u001b[39;00m try_features \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    604\u001b[0m         typed_sequence \u001b[38;5;241m=\u001b[39m OptimizedTypedSequence(col_values, \u001b[38;5;28mtype\u001b[39m\u001b[38;5;241m=\u001b[39mcol_type, try_type\u001b[38;5;241m=\u001b[39mcol_try_type, col\u001b[38;5;241m=\u001b[39mcol)\n\u001b[0;32m--> 605\u001b[0m         arrays\u001b[38;5;241m.\u001b[39mappend(\u001b[43mpa\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43marray\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtyped_sequence\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[1;32m    606\u001b[0m         inferred_features[col] \u001b[38;5;241m=\u001b[39m typed_sequence\u001b[38;5;241m.\u001b[39mget_inferred_type()\n\u001b[1;32m    607\u001b[0m schema \u001b[38;5;241m=\u001b[39m inferred_features\u001b[38;5;241m.\u001b[39marrow_schema \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpa_writer \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mschema\n",
+      "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/pyarrow/array.pxi:252\u001b[0m, in \u001b[0;36mpyarrow.lib.array\u001b[0;34m()\u001b[0m\n",
+      "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/pyarrow/array.pxi:114\u001b[0m, in \u001b[0;36mpyarrow.lib._handle_arrow_array_protocol\u001b[0;34m()\u001b[0m\n",
+      "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/datasets/arrow_writer.py:218\u001b[0m, in \u001b[0;36mTypedSequence.__arrow_array__\u001b[0;34m(self, type)\u001b[0m\n\u001b[1;32m    215\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m    216\u001b[0m     \u001b[38;5;66;03m# custom pyarrow types\u001b[39;00m\n\u001b[1;32m    217\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(pa_type, _ArrayXDExtensionType):\n\u001b[0;32m--> 218\u001b[0m         storage \u001b[38;5;241m=\u001b[39m \u001b[43mto_pyarrow_listarray\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdata\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mpa_type\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    219\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m pa\u001b[38;5;241m.\u001b[39mExtensionArray\u001b[38;5;241m.\u001b[39mfrom_storage(pa_type, storage)\n\u001b[1;32m    221\u001b[0m     \u001b[38;5;66;03m# efficient np array to pyarrow array\u001b[39;00m\n",
+      "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/datasets/features/features.py:1591\u001b[0m, in \u001b[0;36mto_pyarrow_listarray\u001b[0;34m(data, pa_type)\u001b[0m\n\u001b[1;32m   1589\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m any_np_array_to_pyarrow_listarray(data, \u001b[38;5;28mtype\u001b[39m\u001b[38;5;241m=\u001b[39mpa_type\u001b[38;5;241m.\u001b[39mvalue_type)\n\u001b[1;32m   1590\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1591\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mpa\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43marray\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdata\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mpa_type\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstorage_dtype\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/pyarrow/array.pxi:372\u001b[0m, in \u001b[0;36mpyarrow.lib.array\u001b[0;34m()\u001b[0m\n",
+      "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/pyarrow/array.pxi:42\u001b[0m, in \u001b[0;36mpyarrow.lib._sequence_to_array\u001b[0;34m()\u001b[0m\n",
+      "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/pyarrow/error.pxi:155\u001b[0m, in \u001b[0;36mpyarrow.lib.pyarrow_internal_check_status\u001b[0;34m()\u001b[0m\n",
+      "File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/.venv/lib/python3.12/site-packages/pyarrow/error.pxi:92\u001b[0m, in \u001b[0;36mpyarrow.lib.check_status\u001b[0;34m()\u001b[0m\n",
+      "\u001b[0;31mArrowTypeError\u001b[0m: Could not convert tensor([[  48,   25,  -83,  ...,   31,   45,   41],\n        [ -76, -100,  -94,  ...,   26,  -84, -117],\n        [ -97,   26,   15,  ...,  -97, -107, -109],\n        ...,\n        [  44,  -94, -104,  ..., -110,   18,   27],\n        [ -77,   26,  -77,  ..., -100,   33,   43],\n        [ -98,   22, -111,  ..., -110,   14, -107]], dtype=torch.int8) with type Tensor: was not a sequence or recognized null for conversion to list type"
+     ]
+    }
+   ],
+   "source": [
+    "f2 = activation_store(ds, model, layers=layers, writer_batch_size=10, \n",
+    "                      schema=schema,\n",
+    "                        postprocess_result=float8_postprocess_result)\n",
+    "f2\n",
+    "ds_a2 = Dataset.from_parquet(str(f2)).with_format(\"torch\")\n",
+    "ds_a2.info"
    ]
   },
   {
@@ -272,7 +618,46 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "f2 = activation_store(ds, model, layers=layers, writer_batch_size=10, postprocess_result=float8_postprocess_result)\n",
+    "f2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import Dataset\n",
+    "# load\n",
+    "ds_a2 = Dataset.from_parquet(str(f2)).with_format(\"torch\")\n",
+    "for c in ds_a2.column_names[1:]:\n",
+    "    print(c)\n",
+    "    ds_a2[c] = int8_to_float32(ds_a2[0:-1][c])\n",
+    "# ds_a2 = int8_to_float32(ds_a)\n",
+    "ds_a2.info"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "d = ds_a2[:][c]\n",
+    "print(c)\n",
+    "d.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds_a2.info"
+   ]
   },
   {
    "cell_type": "code",
@@ -282,22 +667,6 @@
    "source": [
     "ds_a[0:2]['logits'].shape"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "ds_a[0:2]['model.layers.0.mlp.down_proj'].shape"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {