def test_tensor_spec():
a1 = TensorSpec(np.dtype("float64"), (-1, 3, 3), "a")
a2 = TensorSpec(np.dtype("float"), (-1, 3, 3),
"a") # float defaults to float64
a3 = TensorSpec(np.dtype("float"), [-1, 3, 3], "a")
a4 = TensorSpec(np.dtype("int"), (-1, 3, 3), "a")
assert a1 == a2
assert a1 == a3
assert a1 != a4
b1 = TensorSpec(np.dtype("float64"), (-1, 3, 3), "b")
assert b1 != a1
with pytest.raises(TypeError, match="Expected `type` to be instance"):
TensorSpec("Unsupported", (-1, 3, 3), "a")
with pytest.raises(TypeError, match="Expected `shape` to be instance"):
TensorSpec(np.dtype("float64"), np.array([-1, 2, 3]), "b")
with pytest.raises(
MlflowException,
match=
"MLflow does not support size information in flexible numpy data types",
):
TensorSpec(np.dtype("<U10"), (-1, ), "b")
a5 = TensorSpec.from_json_dict(**a1.to_dict())
assert a5 == a1
assert TensorSpec.from_json_dict(
**json.loads(json.dumps(a1.to_dict()))) == a1
a6 = TensorSpec(np.dtype("float64"), (-1, 3, 3))
a7 = TensorSpec(np.dtype("float64"), (-1, 3, 3), None)
assert a6 == a7
assert TensorSpec.from_json_dict(
**json.loads(json.dumps(a6.to_dict()))) == a6
def test_tensor_spec():
a1 = TensorSpec(np.dtype("float64"), (-1, 3, 3), "a")
a2 = TensorSpec(np.dtype("float"), (-1, 3, 3),
"a") # float defaults to float64
a3 = TensorSpec(np.dtype("float"), [-1, 3, 3], "a")
a4 = TensorSpec(np.dtype("int"), (-1, 3, 3), "a")
assert a1 == a2
assert a1 == a3
assert a1 != a4
b1 = TensorSpec(np.dtype("float64"), (-1, 3, 3), "b")
assert b1 != a1
with pytest.raises(TypeError) as ex1:
TensorSpec("Unsupported", (-1, 3, 3), "a")
assert "Expected `type` to be instance" in str(ex1.value)
with pytest.raises(TypeError) as ex2:
TensorSpec(np.dtype("float64"), np.array([-1, 2, 3]), "b")
assert "Expected `shape` to be instance" in str(ex2.value)
a5 = TensorSpec.from_json_dict(**a1.to_dict())
assert a5 == a1
assert TensorSpec.from_json_dict(
**json.loads(json.dumps(a1.to_dict()))) == a1
a6 = TensorSpec(np.dtype("float64"), (-1, 3, 3))
a7 = TensorSpec(np.dtype("float64"), (-1, 3, 3), None)
assert a6 == a7
assert TensorSpec.from_json_dict(
**json.loads(json.dumps(a6.to_dict()))) == a6
def test_dtype(nparray, dtype):
schema = _infer_schema(nparray)
assert schema == Schema([TensorSpec(np.dtype(dtype), (-1, ))])
spec = schema.inputs[0]
recreated_spec = TensorSpec.from_json_dict(**spec.to_dict())
assert spec == recreated_spec
enforced_array = _enforce_tensor_spec(nparray, spec)
assert isinstance(enforced_array, np.ndarray)