mlx-cluster 0.0.2__tar.gz → 0.0.3__tar.gz

{mlx_cluster-0.0.2 → mlx_cluster-0.0.3}/CMakeLists.txt

@@ -9,7 +9,6 @@ set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 option(BUILD_SHARED_LIBS "Build extensions as a shared library" ON)
 
 # ----- Dependencies required ----
-find_package(fmt REQUIRED)
 find_package(MLX CONFIG REQUIRED)
 find_package(Python 3.8 COMPONENTS Interpreter Development.Module REQUIRED)
 execute_process(

{mlx_cluster-0.0.2/mlx_cluster.egg-info → mlx_cluster-0.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mlx_cluster
-Version: 0.0.2
+Version: 0.0.3
 Summary: C++ and Metal extensions for MLX CTC Loss
 Author-email: Vinay Pandya <vinayharshadpandya27@gmail.com>
 Project-URL: Homepage, https://github.com/vinayhpandya/mlx_cluster
@@ -15,8 +15,9 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Provides-Extra: dev
 Provides-Extra: test
-Requires-Dist: torch_geometric; extra == "test"
-Requires-Dist: pytest; extra == "test"
+Requires-Dist: pytest==7.4.4; extra == "test"
+Requires-Dist: torch; extra == "test"
+Requires-Dist: mlx_graphs; extra == "test"
 
 # mlx_cluster
 

{mlx_cluster-0.0.2 → mlx_cluster-0.0.3/mlx_cluster.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mlx_cluster
-Version: 0.0.2
+Version: 0.0.3
 Summary: C++ and Metal extensions for MLX CTC Loss
 Author-email: Vinay Pandya <vinayharshadpandya27@gmail.com>
 Project-URL: Homepage, https://github.com/vinayhpandya/mlx_cluster
@@ -15,8 +15,9 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Provides-Extra: dev
 Provides-Extra: test
-Requires-Dist: torch_geometric; extra == "test"
-Requires-Dist: pytest; extra == "test"
+Requires-Dist: pytest==7.4.4; extra == "test"
+Requires-Dist: torch; extra == "test"
+Requires-Dist: mlx_graphs; extra == "test"
 
 # mlx_cluster
 

{mlx_cluster-0.0.2 → mlx_cluster-0.0.3}/mlx_cluster.egg-info/SOURCES.txt

@@ -7,6 +7,7 @@ pyproject.toml
 setup.py
 mlx_cluster/__init__.py
 mlx_cluster/_ext.cpython-311-darwin.so
+mlx_cluster/libmlx.dylib
 mlx_cluster/libmlx_cluster.dylib
 mlx_cluster/mlx_cluster.metallib
 mlx_cluster.egg-info/PKG-INFO

mlx_cluster-0.0.3/mlx_cluster.egg-info/requires.txt

@@ -0,0 +1,7 @@
+
+[dev]
+
+[test]
+pytest==7.4.4
+torch
+mlx_graphs

{mlx_cluster-0.0.2 → mlx_cluster-0.0.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mlx_cluster"
-version = "0.0.2"
+version = "0.0.3"
 authors = [
   { name = "Vinay Pandya", email = "vinayharshadpandya27@gmail.com" },
 ]
@@ -17,7 +17,11 @@ classifiers = [
 
 [project.optional-dependencies]
 dev = []
-
+test = [
+  "pytest==7.4.4",
+  "torch",
+  "mlx_graphs",
+]
 [project.urls]
 Homepage = "https://github.com/vinayhpandya/mlx_cluster"
 Issues = "https://github.com/vinayhpandya/mlx_cluster/Issues"
@@ -27,8 +31,9 @@ Issues = "https://github.com/vinayhpandya/mlx_cluster/Issues"
 requires = [
   "setuptools>=42",
   "cmake>=3.24",
-  "mlx==0.15.*",
+  "mlx>=0.17.0",
   "nanobind@git+https://github.com/wjakob/nanobind.git@2f04eac452a6d9142dedb957701bdb20125561e4",
-  "fmt"
 ]
+
+
 build-backend = "setuptools.build_meta"

{mlx_cluster-0.0.2 → mlx_cluster-0.0.3}/random_walks/RandomWalk.cpp

@@ -37,6 +37,7 @@ namespace mlx::core {
     auto* col_values = col.data<int64_t>();
     auto* rand_values = rand.data<float>();
 
+    std::cout<<"After evaluating outputs"<<std::endl;
     for (int64_t n = 0; n < numel; n++) {
         int64_t n_cur = start_values[n];
         n_out_ptr[n * (walk_length_ + 1)] = n_cur;

{mlx_cluster-0.0.2 → mlx_cluster-0.0.3}/setup.py

@@ -4,13 +4,20 @@ from mlx import extension
 if __name__ == "__main__":
     setup(
         name="mlx_cluster",
-        version="0.0.2",
+        version="0.0.3",
         description="Sample C++ and Metal extensions for MLX primitives.",
         ext_modules=[extension.CMakeExtension("mlx_cluster._ext")],
         cmdclass={"build_ext": extension.CMakeBuild},
         packages=["mlx_cluster"],
         package_data={"mlx_cluster": ["*.so", "*.dylib", "*.metallib"]},
-        extras_require={"dev": [], "test": ["torch_geometric", "pytest"]},
+        extras_require={
+            "dev": [],
+            "test": [
+                "mlx_graphs",
+                "torch",
+                "pytest",
+            ],
+        },
         zip_safe=False,
         python_requires=">=3.8",
     )

mlx_cluster-0.0.3/tests/test_random_walk.py

@@ -0,0 +1,35 @@
+import mlx.core as mx
+import numpy as np
+import time
+
+# Torch dataset
+from torch.utils.data import DataLoader
+
+loader = DataLoader(range(2708), batch_size=2000)
+start_indices = next(iter(loader))
+
+from mlx_graphs.datasets import PlanetoidDataset
+from mlx_graphs.utils.sorting import sort_edge_index
+from torch.utils.data import DataLoader
+from mlx_cluster import random_walk
+
+cora_dataset = PlanetoidDataset(name="cora", base_dir="~")
+# For some reason int_64t and int_32t are not compatible
+edge_index = cora_dataset.graphs[0].edge_index.astype(mx.int64)
+# Convert edge index into a CSR matrix
+sorted_edge_index = sort_edge_index(edge_index=edge_index)
+row_mlx = sorted_edge_index[0][0]
+col_mlx = sorted_edge_index[0][1]
+_, counts_mlx = np.unique(np.array(row_mlx, copy=False), return_counts=True)
+cum_sum_mlx = counts_mlx.cumsum()
+row_ptr_mlx = mx.concatenate([mx.array([0]), mx.array(cum_sum_mlx)])
+start_indices = mx.array(start_indices.numpy())
+
+rand_data = mx.random.uniform(shape=[start_indices.shape[0], 5])
+start_time = time.time()
+
+node_sequence = random_walk(
+    row_ptr_mlx, col_mlx, start_indices, rand_data, 5, stream=mx.cpu
+)
+print("Time taken to complete 1000 random walks : ", time.time() - start_time)
+print("MLX random walks are", node_sequence)

{mlx_cluster-0.0.2 → mlx_cluster-0.0.3}/tests/test_rejection_sampling.py

@@ -2,19 +2,10 @@ import mlx.core as mx
 import numpy as np
 import time
 
-# Torch dataset
-import torch
-import torch_geometric.datasets as pyg_datasets
-from torch_geometric.utils import sort_edge_index
-from torch_geometric.utils.num_nodes import maybe_num_nodes
-from torch_geometric.utils.sparse import index2ptr
+# Torch dataloader
 from torch.utils.data import DataLoader
 
-torch_planetoid = pyg_datasets.Planetoid(root="data/Cora", name="Cora")
-edge_index_torch = torch_planetoid.edge_index
-num_nodes = maybe_num_nodes(edge_index=edge_index_torch)
-row, col = sort_edge_index(edge_index=edge_index_torch, num_nodes=num_nodes)
-row_ptr, col = index2ptr(row, num_nodes), col
+
 loader = DataLoader(range(2708), batch_size=2000)
 start_indices = next(iter(loader))
 # random_walks = torch.ops.torch_cluster.random_walk(
@@ -23,7 +14,6 @@ start_indices = next(iter(loader))
 
 from mlx_graphs.datasets import PlanetoidDataset
 from mlx_graphs.utils.sorting import sort_edge_index
-from torch.utils.data import DataLoader
 from mlx_cluster import rejection_sampling
 
 cora_dataset = PlanetoidDataset(name="cora", base_dir="~")
@@ -35,15 +25,10 @@ _, counts_mlx = np.unique(np.array(row_mlx, copy=False), return_counts=True)
 cum_sum_mlx = counts_mlx.cumsum()
 row_ptr_mlx = mx.concatenate([mx.array([0]), mx.array(cum_sum_mlx)])
 start_indices = mx.array(start_indices.numpy())
-print("row pointer datatype", row_ptr_mlx.dtype)
-print("col datatype", col_mlx.dtype)
-print("start pointer datatype", start_indices.dtype)
-assert mx.array_equal(row_ptr_mlx, mx.array(row_ptr.numpy())), "Arrays not equal"
-assert mx.array_equal(col_mlx, mx.array(col.numpy())), "Col arrays are not equal"
 rand_data = mx.random.uniform(shape=[start_indices.shape[0], 5])
 start_time = time.time()
 node_sequence = rejection_sampling(
     row_ptr_mlx, col_mlx, start_indices, 5, 1.0, 3.0, stream=mx.cpu
 )
-# print("Time taken to complete 1000 random walks : ", time.time() - start_time)
+print("Time taken to complete random walks : ", time.time() - start_time)
 print(node_sequence)

mlx_cluster-0.0.2/mlx_cluster.egg-info/requires.txt

@@ -1,6 +0,0 @@
-
-[dev]
-
-[test]
-torch_geometric
-pytest

mlx_cluster-0.0.2/tests/test_random_walk.py

@@ -1,58 +0,0 @@
-import mlx.core as mx
-import numpy as np
-import time
-
-# Torch dataset
-import torch
-import torch_geometric.datasets as pyg_datasets
-from torch_geometric.utils import sort_edge_index
-from torch_geometric.utils.num_nodes import maybe_num_nodes
-from torch_geometric.utils.sparse import index2ptr
-from torch.utils.data import DataLoader
-
-torch_planetoid = pyg_datasets.Planetoid(root="data/Cora", name="Cora")
-edge_index_torch = torch_planetoid.edge_index
-num_nodes = maybe_num_nodes(edge_index=edge_index_torch)
-row, col = sort_edge_index(edge_index=edge_index_torch, num_nodes=num_nodes)
-row_ptr, col = index2ptr(row, num_nodes), col
-loader = DataLoader(range(2708), batch_size=2000)
-start_indices = next(iter(loader))
-print(edge_index_torch.dtype)
-print(row_ptr.dtype)
-print(col.dtype)
-print(start_indices.dtype)
-random_walks = torch.ops.torch_cluster.random_walk(
-    row_ptr, col, start_indices, 5, 1.0, 1.0
-)
-
-from mlx_graphs.datasets import PlanetoidDataset
-from mlx_graphs.utils.sorting import sort_edge_index
-from torch.utils.data import DataLoader
-from mlx_cluster import random_walk
-
-cora_dataset = PlanetoidDataset(name="cora", base_dir="~")
-edge_index = cora_dataset.graphs[0].edge_index.astype(mx.int64)
-sorted_edge_index = sort_edge_index(edge_index=edge_index)
-print(edge_index.dtype)
-row_mlx = sorted_edge_index[0][0]
-col_mlx = sorted_edge_index[0][1]
-_, counts_mlx = np.unique(np.array(row_mlx, copy=False), return_counts=True)
-cum_sum_mlx = counts_mlx.cumsum()
-row_ptr_mlx = mx.concatenate([mx.array([0]), mx.array(cum_sum_mlx)])
-start_indices = mx.array(start_indices.numpy())
-print("Start indices data type is ", start_indices.dtype)
-print("Col mlx data type is ", col_mlx.dtype)
-print("Row mlx data type is ", row_ptr_mlx.dtype)
-assert mx.array_equal(row_ptr_mlx, mx.array(row_ptr.numpy())), "Arrays not equal"
-assert mx.array_equal(col_mlx, mx.array(col.numpy())), "Col arrays are not equal"
-rand_data = mx.random.uniform(shape=[start_indices.shape[0], 5])
-start_time = time.time()
-print("Start indices data type is ", start_indices.dtype)
-print("Col mlx data type is ", col_mlx.dtype)
-print("Row mlx data type is ", row_ptr_mlx.dtype)
-node_sequence = random_walk(
-    row_ptr_mlx, col_mlx, start_indices, rand_data, 5, stream=mx.gpu
-)
-# print("Time taken to complete 1000 random walks : ", time.time() - start_time)
-print("Torch random walks are", random_walks[0])
-print("MLX random walks are", node_sequence)