I want to make mean on groups of embeddings vectors.
For examples:

pl.DataFrame({
    "id": [1,1 ,2,2],
    "values": [
        [1,1,1], [3, 3, 3],
        [1,1,1], [2, 2, 2]
        ]
    })
shape: (4, 2)
id  values
i64 list[i64]
1   [1, 1, 1]
1   [3, 3, 3]
2   [1, 1, 1]
2   [2, 2, 2]

expected results:

pl.DataFrame(
    {
        "id":[1,2],
        "values": np.array([[[1,1,1], [3, 3, 3]],[[1,1,1], [2, 2, 2]]]).mean(axis=1)
    }
)

shape: (2, 2)
id  values
i64 list[f64]
1   [2.0, 2.0, 2.0]
2   [1.5, 1.5, 1.5]