Does Marimo support serializing Numpy arrays to array buffers?

[flekschas]

I'd like to send data as an array buffer over the wire instead of number[]. In Jupyter I can do this using traittypes.Array and the following array_to_binary() and binary_to_array() methods on the Python side as well as a custom serializer and deserializer on the JS side.

Unfortunately, in Marimo I noticed that the data is not going through array_to_binary() but instead send as number[]. Is there something I'm missing or does Marimo not yet support traittypes.Array?

import anywidget
import numpy as np
import traittypes

def array_to_binary(ar, obj=None, force_contiguous=True):
    if ar is None:
        return None
    if ar.dtype.kind not in ['u', 'i', 'f']:  # ints and floats
        raise ValueError('unsupported dtype: %s' % (ar.dtype))
    if ar.dtype == np.float64:  # WebGL does not support float64, case it here
        ar = ar.astype(np.float32)
    if ar.dtype == np.int64:  # JS does not support int64
        ar = ar.astype(np.int32)
    if force_contiguous and not ar.flags['C_CONTIGUOUS']:  # make sure it's contiguous
        ar = np.ascontiguousarray(ar)
    return {'view': memoryview(ar), 'dtype': str(ar.dtype), 'shape': ar.shape}

def binary_to_array(value, obj=None):
    if value is None:
        return None
    return np.frombuffer(value['view'], dtype=value['dtype']).reshape(value['shape'])

ndarray_serialization = dict(to_json=array_to_binary, from_json=binary_to_array)

class Values(anywidget.AnyWidget):
    _esm = """
    const DTYPES = {
      uint8: Uint8Array,
      int8: Int8Array,
      uint16: Uint16Array,
      int16: Int16Array,
      uint32: Uint32Array,
      int32: Int32Array,
      float32: Float32Array,
      float64: Float64Array,
    };

    function Numpy1D(dtype) {
      if (!(dtype in DTYPES)) {
        throw new Error(`Dtype not supported, got ${JSON.stringify(dtype)}.`);
      }
      return {
        deserialize(data) {
          if (data == null) {
            return null;
          }
          return new DTYPES[dtype](data.view.buffer);
        },
        serialize(arr) {
          return {
            view: new DataView(arr.buffer),
            dtype,
            shape: [data.length],
          };
        },
      };
    }

    function modelWithSerializers(model, serializers) {
      return {
        get(key) {
          const value = model.get(key);
          const serializer = serializers[key];
          if (serializer) {
            // The actual issue: `value` is of type `number[]` and not `{ view, dtype, shape }`
            return serializer.deserialize(value);
          }
          return value;
        },
        set(key, value) {
          if (key in serializers) {
            model.set(key, serializers[key].serialize(value));
          } else {
            model.set(key, value);
          }
        },
        on: model.on.bind(model),
        save_changes: model.save_changes.bind(model),
        send: model.send.bind(model),
      };
    }

    function render({ model: originalModel, el }) {
      const model = modelWithSerializers(originalModel, {
        values: Numpy1D('uint32'),
      })
      const div = document.createElement("div");
      const update = () => {
        div.textContent = model.get('values').join(', ')
      }
      model.on("change", update);
      update();
      el.appendChild(div);
    }
    export default { render };
    """
    values = traittypes.Array(default_value=None).tag(sync=True, **ndarray_serialization)

Values(values=np.array([1, 3, 3, 7]))

[mscolnick]

It's likely we don't support this yet (just because we haven't hit this case yet), but there should be no reason we can't support this.

[flekschas]

This would be really cool as serializing data to large-scale to JSON is not super efficient.