Merge branch 'main' into add-larecs

Author: carolinefrasca

recipes/ExtraMojo/recipe.yaml

@@ -1,13 +1,13 @@
 context:
-  version: "0.12.0"
+  version: "0.13.0"
 
 package:
   name: "extramojo"
   version: ${{ version }}
 
 source:
   - git: https://github.com/ExtraMojo/ExtraMojo.git
-    rev: 25b6a6e7089138ab93f99d79401c57c8174c38ee
+    rev: 97c9151d968206581cfd8a1d626c2e5d519046e5
 
 build:
   number: 0

recipes/ExtraMojo/tests/test_bstr.mojo

@@ -5,7 +5,6 @@ from ExtraMojo.bstr.bstr import (
     to_ascii_uppercase,
 )
 from ExtraMojo.bstr.memchr import memchr, memchr_wide
-from memory import Span
 from testing import *
 
 
@@ -28,30 +27,36 @@ fn s(bytes: Span[UInt8]) -> String:
 
 
 fn test_memchr() raises:
-    var cases = List[(StringLiteral, Int)](
-        (
-            "enlivened,unleavened,Arnulfo's,Unilever's,unloved|Anouilh,analogue,analogy",
-            49,
-        ),
-        (
-            "enlivened,unleavened,Arnulfo's,Unilever's,unloved,Anouilh,analogue,analogy,enlivened,unleavened,Arnulfo's,Unilever's,unloved|Anouilh,analogue,analogy",
-            124,
-        ),
-    )
-
-    for kase in cases:
-        var index = memchr(kase[][0].as_bytes(), ord("|"))
-        assert_equal(
-            index,
-            kase[][1],
-            "Expected "
-            + String(kase[][1])
-            + " Found "
-            + String(index)
-            + " in "
-            + kase[][0],
+    alias check = InlineArray[Bool, 2](True, False)
+
+    @parameter
+    for do_alignment in range(0, len(check)):
+        var cases = List[(StringLiteral, Int)](
+            (
+                "enlivened,unleavened,Arnulfo's,Unilever's,unloved|Anouilh,analogue,analogy",
+                49,
+            ),
+            (
+                "enlivened,unleavened,Arnulfo's,Unilever's,unloved,Anouilh,analogue,analogy,enlivened,unleavened,Arnulfo's,Unilever's,unloved|Anouilh,analogue,analogy",
+                124,
+            ),
         )
 
+        for kase in cases:
+            var index = memchr[do_alignment = check[do_alignment]](
+                kase[][0].as_bytes(), ord("|")
+            )
+            assert_equal(
+                index,
+                kase[][1],
+                "Expected "
+                + String(kase[][1])
+                + " Found "
+                + String(index)
+                + " in "
+                + kase[][0],
+            )
+
 
 fn test_memchr_wide() raises:
     var cases = List[(StringLiteral, Int)](

recipes/ExtraMojo/tests/test_file.mojo

@@ -1,6 +1,5 @@
-from collections import Optional, Dict
+from collections import Dict
 from utils import StringSlice
-from memory import Span
 from pathlib import Path
 from python import Python
 from tensor import Tensor

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

recipes/bridge/recipe.yaml

@@ -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') }}
 

recipes/bridge/test_numpy.🔥

@@ -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)