def test_get_tensor_shape(dict_of_ndarrays):
assert all(-1 == _get_tensor_shape(tensor)[0]
for tensor in dict_of_ndarrays.values())
data = dict_of_ndarrays["4D"]
# Specify variable dimension
for i in range(-4, 4):
assert _get_tensor_shape(data, i)[i] == -1
# Specify None
assert all([_get_tensor_shape(data, None) != -1])
# Out of bounds
with pytest.raises(
MlflowException,
match="The specified variable_dimension 10 is out of bounds"):
_get_tensor_shape(data, 10)
with pytest.raises(
MlflowException,
match="The specified variable_dimension -10 is out of bounds"):
_get_tensor_shape(data, -10)
with pytest.raises(
TypeError,
match="Data in the dictionary must be of type numpy.ndarray"):
_infer_schema({"x": 1})
def test_schema_inference_on_dictionary(dict_of_ndarrays):
# test dictionary
schema = _infer_schema(dict_of_ndarrays)
assert schema == Schema([
TensorSpec(tensor.dtype, _get_tensor_shape(tensor), name)
for name, tensor in dict_of_ndarrays.items()
])
# test exception is raised if non-numpy data in dictionary
with pytest.raises(TypeError):
_infer_schema({"x": 1})
with pytest.raises(TypeError):
_infer_schema({"x": [1]})
def test_get_tensor_shape(dict_of_ndarrays):
assert all([-1 == _get_tensor_shape(tensor)[0] for tensor in dict_of_ndarrays.values()])
data = dict_of_ndarrays["4D"]
# Specify variable dimension
for i in range(-4, 4):
assert _get_tensor_shape(data, i)[i] == -1
# Specify None
assert all([_get_tensor_shape(data, None) != -1])
# Out of bounds
with pytest.raises(MlflowException):
_get_tensor_shape(data, 10)
with pytest.raises(MlflowException):
_get_tensor_shape(data, -10)
with pytest.raises(TypeError):
_infer_schema({"x": 1})
def test_get_sparse_matrix_data_type_and_shape(dict_of_sparse_matrix):
for sparse_matrix in dict_of_sparse_matrix.values():
schema = _infer_schema(sparse_matrix)
assert schema.numpy_types() == ["float64"]
assert _get_tensor_shape(sparse_matrix) == (-1, 8)
