mlx.data.Buffer.shape

Buffer.shape(self: mlx.data._c.Buffer, key: str, output_key: str, dim: Optional[int] = None) → mlx.data._c.Buffer

Extracts the shape of an array in the sample.

If a dimension is provided then only the size of that dimension is extracted.

Parameters:

• key (str) – The sample key that contains the array we are operating on.

• output_key (str) – The key to write the output at. It is required on this operation as it is very unlikely that we will want to replace the original key.

• dim (int, optional) – The dimension to report the size for. If not provided then the full size of the array is returned. (default: None)