diff --git a/recipes/bridge/README.md b/recipes/bridge/README.md index 59787c7f..07a57222 100644 --- a/recipes/bridge/README.md +++ b/recipes/bridge/README.md @@ -6,24 +6,33 @@ This package acts as a bridge between Mojo native types and those from the Python world. -Example usage, to convert a NumPy array to a Mojo tensor: +Example usage, to convert a NumPy array to a Mojo `LayoutTensor` or `Tensor`: ```mojo from python import Python -from bridge.numpy import ndarray_to_tensor +from bridge.numpy import ndarray_to_layouttensor, ndarray_to_tensor var np = Python.import_module("numpy") np_array = np.arange(6.0).reshape(2,3) -mojo_tensor = ndarray_to_tensor[DType.float64](np_array) + +# Convert to new-style `LayoutTensor`. +mojo_tensor = ndarray_to_layouttensor[order=2](np_array) + +# Convert to old-style `Tensor`. +mojo_oldtensor = ndarray_to_tensor[DType.float64](np_array) ``` Or to achieve the reverse: ```mojo from tensor import Tensor -from bridge.numpy import tensor_to_ndarray +from bridge.numpy import layouttensor_to_ndarray, tensor_to_ndarray + +# Convert new-style `LayoutTensor` to numpy array. +np_array = layouttensor_to_ndarray(mojo_tensor) +# Convert old-style `Tensor` to numpy array. values = List[Float64](0.0, 1.0, 2.0, 3.0, 4.0, 5.0) -mojo_tensor = Tensor[DType.float64](shape=(2, 3), list=values) -np_array = tensor_to_ndarray(mojo_tensor) +mojo_oldtensor = Tensor[DType.float64](shape=(2, 3), list=values) +np_array = tensor_to_ndarray(mojo_oldtensor) ``` # Installation diff --git a/recipes/bridge/recipe.yaml b/recipes/bridge/recipe.yaml index 28b41f2f..1d970802 100644 --- a/recipes/bridge/recipe.yaml +++ b/recipes/bridge/recipe.yaml @@ -1,5 +1,5 @@ context: - version: "0.1.0" + version: "0.2.0" package: name: "bridge" @@ -7,7 +7,7 @@ package: source: - git: https://gitlab.com/hylkedonker/bridge.git - rev: dc89432cb79b6eeef42e1755593ea3837c48d7dd + rev: 87f10144b004ff4ef8e73d6b8b52e3f81effd40f build: number: 0 @@ -15,7 +15,7 @@ build: - mojo package src/bridge -o ${{ PREFIX }}/lib/mojo/bridge.mojopkg requirements: host: - - max=25.1 + - max=25.2 run: - ${{ pin_compatible('max') }} diff --git "a/recipes/bridge/test_numpy.\360\237\224\245" "b/recipes/bridge/test_numpy.\360\237\224\245" index adde832b..5b690d06 100644 --- "a/recipes/bridge/test_numpy.\360\237\224\245" +++ "b/recipes/bridge/test_numpy.\360\237\224\245" @@ -1,8 +1,14 @@ from python import Python -from tensor import Tensor, rand -from testing import assert_true, assert_equal +from layout import Layout, LayoutTensor +from tensor import Tensor +from testing import assert_equal, assert_raises, assert_true -from bridge.numpy import ndarray_to_tensor, tensor_to_ndarray +from bridge.numpy import ( + layouttensor_to_ndarray, + ndarray_to_layouttensor, + ndarray_to_tensor, + tensor_to_ndarray, +) def test_tensor_identity_transformation(): @@ -14,10 +20,104 @@ def test_tensor_identity_transformation(): assert_equal(in_matrix, out_matrix) -def test_numpy_identity_transformation(): +def test_numpy_tensor_identity_transformation(): """Test that `tensor_to_ndarray` is inverse of `ndarray_to_tensor`.""" var np = Python.import_module("numpy") var in_array = np.arange(6).reshape(3, 2) var tensor = ndarray_to_tensor[DType.float64](in_array) var out_array = tensor_to_ndarray(tensor) assert_equal(in_array, out_array) + + +def test_ndarray_to_layouttensor(): + """Test numpy array conversion to layouttensor for various tensor shapes.""" + var np = Python.import_module("numpy") + + # 1) Test vectors + var in_vector = np.arange(4.0) + var out_vector = ndarray_to_layouttensor[order=1](in_vector) + assert_equal(out_vector[1], 1.0) + assert_equal(out_vector[3], 3.0) + + # 2) Test matrices + var in_matrix = np.arange(4.0 * 3.0).reshape(3, 4) + var out_matrix = ndarray_to_layouttensor[order=2](in_matrix) + assert_equal(out_matrix[0, 0], 0.0) + assert_equal(out_matrix[1, 1], 5.0) + assert_equal(out_matrix[1, 3], 7.0) + assert_equal(out_matrix[2, 1], 9.0) + + # Check that non-contiguous arrays raise exceptions. + with assert_raises(): + var in_matrix_col_major = np.asfortranarray(in_matrix) + _ = ndarray_to_layouttensor[order=2](in_matrix_col_major) + + # 3) Test three-index tensors + var in_tensor = np.arange(4.0 * 3.0).reshape(3, 1, 4) + var out_tensor = ndarray_to_layouttensor[order=3](in_tensor) + assert_equal(out_tensor[0, 0, 0], 0.0) + assert_equal(out_tensor[1, 0, 1], 5.0) + assert_equal(out_tensor[1, 0, 3], 7.0) + assert_equal(out_tensor[2, 0, 1], 9.0) + + # 3) Test four-index tensors + var in_4tensor = np.arange(4.0 * 3.0).reshape(2, 3, 1, 2) + var out_4tensor = ndarray_to_layouttensor[order=4](in_4tensor) + assert_equal(out_4tensor[0, 0, 0, 0], 0.0) + assert_equal(out_4tensor[0, 1, 0, 1], 3.0) + assert_equal(out_4tensor[1, 0, 0, 1], 7.0) + assert_equal(out_4tensor[0, 2, 0, 0], 4.0) + + +def test_memory_leaks(): + """Test that we can safely remove the reference to the numpy array.""" + var np = Python.import_module("numpy") + var np_array = np.arange(6.0).reshape(3, 2) + var tensor = ndarray_to_layouttensor[order=2](np_array) + np_array.__del__() + assert_equal(tensor[1, 0], 2.0) + assert_equal(tensor[1, 1], 3.0) + assert_equal(tensor[2, 1], 5.0) + + +# def test_layouttensor_numpy_identity_transformation(): +# """Test that `layouttensor_to_ndarray` is inverse of `ndarray_to_layouttensor`. +# """ +# values = List[Float64](0.0, 1.0, 2.0, 3.0, 4.0, 5.0) +# tensor = Tensor[DType.float64](shape=(2, 3), list=values) +# var ptr = tensor.unsafe_ptr() +# layouttensor = LayoutTensor[ +# mut=False, +# DType.float64, +# Layout.row_major(2, 3), +# __origin_of(ptr[]), +# ](ptr=ptr) +# np_array = layouttensor_to_ndarray(layouttensor) +# out_layouttensor = ndarray_to_layouttensor[order=2](in_array) + + +def test_numpy_layouttensor_identity_transformation(): + """Test that `ndarray_to_layouttensor` is the inverse of `layouttensor_to_ndarray`. + """ + var np = Python.import_module("numpy") + + # 1) Test vectors + # TODO: Add support for vectors! + + # 2) Test matrices + var in_matrix = np.arange(4.0 * 3.0).reshape(3, 4) + var layout_matrix = ndarray_to_layouttensor[order=2](in_matrix) + var out_matrix = layouttensor_to_ndarray(layout_matrix) + np.testing.assert_array_equal(in_matrix, out_matrix) + + # 3) Test three-index tensors + var in_tensor = np.arange(4.0 * 3.0).reshape(3, 1, 4) + var layout_tensor = ndarray_to_layouttensor[order=3](in_tensor) + var out_tensor = layouttensor_to_ndarray(layout_tensor) + np.testing.assert_array_equal(in_tensor, out_tensor) + + # 3) Test four-index tensors + var in_4tensor = np.arange(4.0 * 3.0).reshape(2, 3, 1, 2) + var layout_4tensor = ndarray_to_layouttensor[order=4](in_4tensor) + var out_4tensor = layouttensor_to_ndarray(layout_4tensor) + np.testing.assert_array_equal(in_4tensor, out_4tensor)