Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 15 additions & 6 deletions recipes/bridge/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 3 additions & 3 deletions recipes/bridge/recipe.yaml
Original file line number Diff line number Diff line change
@@ -1,21 +1,21 @@
context:
version: "0.1.0"
version: "0.2.0"

package:
name: "bridge"
version: ${{ version }}

source:
- git: https://gitlab.com/hylkedonker/bridge.git
rev: dc89432cb79b6eeef42e1755593ea3837c48d7dd
rev: 87f10144b004ff4ef8e73d6b8b52e3f81effd40f

build:
number: 0
script:
- mojo package src/bridge -o ${{ PREFIX }}/lib/mojo/bridge.mojopkg
requirements:
host:
- max=25.1
- max=25.2
run:
- ${{ pin_compatible('max') }}

Expand Down
108 changes: 104 additions & 4 deletions recipes/bridge/test_numpy.🔥
Original file line number Diff line number Diff line change
@@ -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():
Expand All @@ -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)