import pytest
import pickle
import numpy as np
from tvm import ffi as tvm_ffi

float32 = tvm_ffi.dtype("float32")
assert float32.__repr__() == "dtype('float32')"
assert type(float32) == tvm_ffi.dtype
x = np.array([1, 2, 3], dtype=float32)
assert x.dtype == float32

for dtype_str in ["int32", "float32x4", "float8_e5m2x4", "uint8"]:
    dtype = tvm_ffi.dtype(dtype_str)
    print(f"{dtype_str}: {dtype.itemsize}")
    # pickle
    dtype_pickled = pickle.loads(pickle.dumps(dtype))
    assert dtype_pickled.type_code == dtype.type_code
    assert dtype_pickled.bits == dtype.bits
    assert dtype_pickled.lanes == dtype.lanes

dtype = tvm_ffi.dtype("float32")
dtype_with_lanes = dtype.with_lanes(4)
assert dtype_with_lanes.type_code == dtype.type_code
assert dtype_with_lanes.bits == dtype.bits
assert dtype_with_lanes.lanes == 4