mlx-cluster 0.0.4__tar.gz → 0.0.6__tar.gz

{mlx_cluster-0.0.4 → mlx_cluster-0.0.6}/CMakeLists.txt

@@ -17,6 +17,15 @@ execute_process(
 list(APPEND CMAKE_PREFIX_PATH "${NB_DIR}")
 find_package(nanobind CONFIG REQUIRED)
 
+include(FetchContent)
+
+FetchContent_Declare(
+  parallel-hashmap
+  GIT_REPOSITORY https://github.com/greg7mdp/parallel-hashmap.git
+  GIT_TAG        v1.4.1  # Use latest stable version
+)
+
+FetchContent_MakeAvailable(parallel-hashmap)
 # ------ Adding extensions to the library -----
 
 # Add library
@@ -26,13 +35,14 @@ target_sources(mlx_cluster
 PUBLIC
 ${CMAKE_CURRENT_LIST_DIR}/random_walks/RandomWalk.cpp
 ${CMAKE_CURRENT_LIST_DIR}/random_walks/BiasedRandomWalk.cpp
+${CMAKE_CURRENT_LIST_DIR}/random_walks/NeighborSample.cpp
 )
 
 target_include_directories(mlx_cluster
 PUBLIC
 ${CMAKE_CURRENT_LIST_DIR})
 
-target_link_libraries(mlx_cluster PUBLIC mlx)
+target_link_libraries(mlx_cluster PUBLIC mlx phmap)
 
 
 if(MLX_BUILD_METAL)

mlx_cluster-0.0.6/PKG-INFO

@@ -0,0 +1,243 @@
+Metadata-Version: 2.4
+Name: mlx_cluster
+Version: 0.0.6
+Summary: C++ extension for generating random graphs
+Author-email: Vinay Pandya <vinayharshadpandya27@gmail.com>
+Project-URL: Homepage, https://github.com/vinayhpandya/mlx_cluster
+Project-URL: Issues, https://github.com/vinayhpandya/mlx_cluster/Issues
+Classifier: Development Status :: 3 - Alpha
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: C++
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: MacOS
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Provides-Extra: dev
+Provides-Extra: docs
+Requires-Dist: mlx>=0.27.1; extra == "docs"
+Requires-Dist: mlx-graphs>=0.0.8; extra == "docs"
+Requires-Dist: ipython==8.21.0; extra == "docs"
+Requires-Dist: sphinx>=7.2.6; extra == "docs"
+Requires-Dist: sphinx-book-theme==1.1.0; extra == "docs"
+Requires-Dist: sphinx-autodoc-typehints==1.25.2; extra == "docs"
+Requires-Dist: nbsphinx==0.9.3; extra == "docs"
+Requires-Dist: sphinx-gallery==0.15.0; extra == "docs"
+Provides-Extra: test
+Requires-Dist: mlx-graphs>=0.0.8; extra == "test"
+Requires-Dist: torch>=2.2.0; extra == "test"
+Requires-Dist: mlx>=0.26.0; extra == "test"
+Requires-Dist: pytest==7.4.4; extra == "test"
+Requires-Dist: scipy>=1.13.0; extra == "test"
+Requires-Dist: requests==2.31.0; extra == "test"
+Requires-Dist: fsspec[http]==2024.2.0; extra == "test"
+Requires-Dist: tqdm==4.66.1; extra == "test"
+Dynamic: license-file
+Dynamic: requires-python
+
+# MLX-Cluster
+
+High-performance graph algorithms optimized for Apple's MLX framework, featuring random walks, biased random walks, and neighbor sampling.
+
+[![PyPI version](https://badge.fury.io/py/mlx-cluster.svg)](https://badge.fury.io/py/mlx-cluster)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://www.python.org/downloads/)
+
+**[Documentation](https://vinayhpandya.github.io/mlx_cluster/)** | **[Quickstart](https://vinayhpandya.github.io/mlx_cluster/)** |
+
+## 🚀 Features
+
+- **🔥 MLX Optimized**: Built specifically for Apple's MLX framework with GPU acceleration
+- **⚡ High Performance**: Optimized C++ implementations with Metal shaders for Apple Silicon
+- **🎯 Graph Algorithms**: 
+  - Uniform random walks
+  - Biased random walks (Node2Vec style with p/q parameters)
+  - Multi-hop neighbor sampling (GraphSAGE style)
+
+## 📦 Installation
+
+### From PyPI (Recommended)
+
+```bash
+pip install mlx-cluster
+```
+
+### From Source
+
+```bash
+git clone https://github.com/vinayhpandya/mlx_cluster.git
+cd mlx_cluster
+pip install -e .
+```
+
+### Development Installation
+
+```bash
+git clone https://github.com/vinayhpandya/mlx_cluster.git
+cd mlx_cluster
+pip install -e . --verbose
+```
+
+### Dependencies
+
+Required:
+- Python 3.8+
+- MLX framework
+- NumPy
+
+Optional (for examples and testing):
+- MLX-Graphs
+- PyTorch (for dataset utilities)
+- pytest
+
+## 🔧 Quick Start
+
+### Random Walks
+
+```python
+import mlx.core as mx
+import numpy as np
+from mlx_cluster import random_walk
+from mlx_graphs.datasets import PlanetoidDataset
+from mlx_graphs.utils.sorting import sort_edge_index
+
+# Load dataset
+cora = PlanetoidDataset(name="cora")
+edge_index = cora.graphs[0].edge_index.astype(mx.int64)
+
+# Convert to CSR format
+sorted_edge_index = sort_edge_index(edge_index=edge_index)
+row = sorted_edge_index[0][0]
+col = sorted_edge_index[0][1]
+_, counts = np.unique(np.array(row, copy=False), return_counts=True)
+row_ptr = mx.concatenate([mx.array([0]), mx.array(counts.cumsum())])
+
+# Generate random walks
+num_walks = 1000
+walk_length = 10
+start_nodes = mx.array(np.random.randint(0, cora.graphs[0].num_nodes, num_walks))
+rand_values = mx.random.uniform(shape=[num_walks, walk_length])
+
+mx.eval(rowptr,col, start_nodes, rand_values)
+# Perform walks
+node_sequences, edge_sequences = random_walk(
+    row_ptr, col, start_nodes, rand_values, walk_length, stream=mx.gpu
+)
+
+print(f"Generated {num_walks} walks of length {walk_length + 1}")
+print(f"Shape: {node_sequences.shape}")
+```
+
+### Biased Random Walks (Node2Vec)
+
+```python
+from mlx_cluster import rejection_sampling
+
+# Biased random walks with p/q parameters
+node_sequences, edge_sequences = rejection_sampling(
+    row_ptr, col, start_nodes, walk_length,
+    p=1.0,  # Return parameter
+    q=2.0,  # In-out parameter
+    stream=mx.gpu
+)
+```
+
+### Neighbor Sampling
+
+```python
+from mlx_cluster import neighbor_sample
+
+# Convert to CSC format (required for neighbor sampling)
+def create_csc_format(edge_index, num_nodes):
+    sources, targets = edge_index[0].tolist(), edge_index[1].tolist()
+    edges = sorted(zip(sources, targets), key=lambda x: x[1])
+    
+    colptr = np.zeros(num_nodes + 1, dtype=np.int64)
+    for _, target in edges:
+        colptr[target + 1] += 1
+    colptr = np.cumsum(colptr)
+    
+    sorted_sources = [source for source, _ in edges]
+    return mx.array(colptr), mx.array(sorted_sources, dtype=mx.int64)
+
+colptr, row = create_csc_format(edge_index, cora.graphs[0].num_nodes)
+
+# Multi-hop neighbor sampling
+input_nodes = mx.array([0, 1, 2], dtype=mx.int64)
+num_neighbors = [10, 5]  # 10 neighbors in first hop, 5 in second
+mx.eval(colptr, row, input_nodes)
+samples, rows, cols, edges = neighbor_sample(
+    colptr, row, input_nodes, num_neighbors,
+    replace=True, directed=True
+)
+
+print(f"Sampled {len(samples)} nodes and {len(edges)} edges")
+```
+
+## 📚 Documentation
+
+For comprehensive documentation, examples, and API reference, visit:
+[Documentation]()
+
+## 🧪 Testing
+
+Run the test suite:
+
+```bash
+# Install test dependencies
+pip install pytest mlx-graphs torch
+
+# Run tests
+pytest -s -v
+```
+
+## ⚡ Performance
+
+MLX-Cluster is optimized for Apple Silicon and shows significant speedups:
+
+- **Apple M1/M2/M3**: 2-5x faster than CPU-only implementations
+- **GPU Acceleration**: Automatic optimization for Metal Performance Shaders
+- **Memory Efficient**: Optimized sparse graph representations
+- **Batch Processing**: Efficient handling of thousands of concurrent walks
+
+## 🤝 Contributing
+
+We welcome contributions!
+1. Fork the repository
+2. Create your feature branch (`git checkout -b feature/new-feature`)
+3. Commit your changes (`git commit -m 'Add new algorithm'`)
+4. Push to the branch (`git push origin feature/new-feature`)
+5. Open a Pull Request
+For installation and test intructions please visit the documentation
+
+## 📄 License
+
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
+
+## 🙏 Acknowledgments
+
+- [PyTorch Cluster](https://github.com/rusty1s/pytorch_cluster) for everything
+- [MLX](https://github.com/ml-explore/mlx) for the foundational framework
+- [MLX-Graphs](https://github.com/mlx-graphs/mlx-graphs) for graph utilities and datasets
+
+## 📊 Citation
+
+If you use MLX-Cluster in your research, please cite:
+
+```bibtex
+@software{mlx_cluster,
+  author = {Vinay Pandya},
+  title = {MLX-Cluster: High-Performance Graph Algorithms for Apple MLX},
+  url = {https://github.com/vinayhpandya/mlx_cluster},
+  version = {0.0.6},
+  year = {2025}
+}
+```
+
+## 🔗 Related Projects
+
+- [MLX](https://github.com/ml-explore/mlx) - Apple's machine learning framework
+- [MLX-Graphs](https://github.com/mlx-graphs/mlx-graphs) - Graph neural networks for MLX
+- [PyTorch Geometric](https://github.com/pyg-team/pytorch_geometric) - Graph deep learning library
+
+---

mlx_cluster-0.0.6/README.md

@@ -0,0 +1,206 @@
+# MLX-Cluster
+
+High-performance graph algorithms optimized for Apple's MLX framework, featuring random walks, biased random walks, and neighbor sampling.
+
+[![PyPI version](https://badge.fury.io/py/mlx-cluster.svg)](https://badge.fury.io/py/mlx-cluster)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://www.python.org/downloads/)
+
+**[Documentation](https://vinayhpandya.github.io/mlx_cluster/)** | **[Quickstart](https://vinayhpandya.github.io/mlx_cluster/)** |
+
+## 🚀 Features
+
+- **🔥 MLX Optimized**: Built specifically for Apple's MLX framework with GPU acceleration
+- **⚡ High Performance**: Optimized C++ implementations with Metal shaders for Apple Silicon
+- **🎯 Graph Algorithms**: 
+  - Uniform random walks
+  - Biased random walks (Node2Vec style with p/q parameters)
+  - Multi-hop neighbor sampling (GraphSAGE style)
+
+## 📦 Installation
+
+### From PyPI (Recommended)
+
+```bash
+pip install mlx-cluster
+```
+
+### From Source
+
+```bash
+git clone https://github.com/vinayhpandya/mlx_cluster.git
+cd mlx_cluster
+pip install -e .
+```
+
+### Development Installation
+
+```bash
+git clone https://github.com/vinayhpandya/mlx_cluster.git
+cd mlx_cluster
+pip install -e . --verbose
+```
+
+### Dependencies
+
+Required:
+- Python 3.8+
+- MLX framework
+- NumPy
+
+Optional (for examples and testing):
+- MLX-Graphs
+- PyTorch (for dataset utilities)
+- pytest
+
+## 🔧 Quick Start
+
+### Random Walks
+
+```python
+import mlx.core as mx
+import numpy as np
+from mlx_cluster import random_walk
+from mlx_graphs.datasets import PlanetoidDataset
+from mlx_graphs.utils.sorting import sort_edge_index
+
+# Load dataset
+cora = PlanetoidDataset(name="cora")
+edge_index = cora.graphs[0].edge_index.astype(mx.int64)
+
+# Convert to CSR format
+sorted_edge_index = sort_edge_index(edge_index=edge_index)
+row = sorted_edge_index[0][0]
+col = sorted_edge_index[0][1]
+_, counts = np.unique(np.array(row, copy=False), return_counts=True)
+row_ptr = mx.concatenate([mx.array([0]), mx.array(counts.cumsum())])
+
+# Generate random walks
+num_walks = 1000
+walk_length = 10
+start_nodes = mx.array(np.random.randint(0, cora.graphs[0].num_nodes, num_walks))
+rand_values = mx.random.uniform(shape=[num_walks, walk_length])
+
+mx.eval(rowptr,col, start_nodes, rand_values)
+# Perform walks
+node_sequences, edge_sequences = random_walk(
+    row_ptr, col, start_nodes, rand_values, walk_length, stream=mx.gpu
+)
+
+print(f"Generated {num_walks} walks of length {walk_length + 1}")
+print(f"Shape: {node_sequences.shape}")
+```
+
+### Biased Random Walks (Node2Vec)
+
+```python
+from mlx_cluster import rejection_sampling
+
+# Biased random walks with p/q parameters
+node_sequences, edge_sequences = rejection_sampling(
+    row_ptr, col, start_nodes, walk_length,
+    p=1.0,  # Return parameter
+    q=2.0,  # In-out parameter
+    stream=mx.gpu
+)
+```
+
+### Neighbor Sampling
+
+```python
+from mlx_cluster import neighbor_sample
+
+# Convert to CSC format (required for neighbor sampling)
+def create_csc_format(edge_index, num_nodes):
+    sources, targets = edge_index[0].tolist(), edge_index[1].tolist()
+    edges = sorted(zip(sources, targets), key=lambda x: x[1])
+    
+    colptr = np.zeros(num_nodes + 1, dtype=np.int64)
+    for _, target in edges:
+        colptr[target + 1] += 1
+    colptr = np.cumsum(colptr)
+    
+    sorted_sources = [source for source, _ in edges]
+    return mx.array(colptr), mx.array(sorted_sources, dtype=mx.int64)
+
+colptr, row = create_csc_format(edge_index, cora.graphs[0].num_nodes)
+
+# Multi-hop neighbor sampling
+input_nodes = mx.array([0, 1, 2], dtype=mx.int64)
+num_neighbors = [10, 5]  # 10 neighbors in first hop, 5 in second
+mx.eval(colptr, row, input_nodes)
+samples, rows, cols, edges = neighbor_sample(
+    colptr, row, input_nodes, num_neighbors,
+    replace=True, directed=True
+)
+
+print(f"Sampled {len(samples)} nodes and {len(edges)} edges")
+```
+
+## 📚 Documentation
+
+For comprehensive documentation, examples, and API reference, visit:
+[Documentation]()
+
+## 🧪 Testing
+
+Run the test suite:
+
+```bash
+# Install test dependencies
+pip install pytest mlx-graphs torch
+
+# Run tests
+pytest -s -v
+```
+
+## ⚡ Performance
+
+MLX-Cluster is optimized for Apple Silicon and shows significant speedups:
+
+- **Apple M1/M2/M3**: 2-5x faster than CPU-only implementations
+- **GPU Acceleration**: Automatic optimization for Metal Performance Shaders
+- **Memory Efficient**: Optimized sparse graph representations
+- **Batch Processing**: Efficient handling of thousands of concurrent walks
+
+## 🤝 Contributing
+
+We welcome contributions!
+1. Fork the repository
+2. Create your feature branch (`git checkout -b feature/new-feature`)
+3. Commit your changes (`git commit -m 'Add new algorithm'`)
+4. Push to the branch (`git push origin feature/new-feature`)
+5. Open a Pull Request
+For installation and test intructions please visit the documentation
+
+## 📄 License
+
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
+
+## 🙏 Acknowledgments
+
+- [PyTorch Cluster](https://github.com/rusty1s/pytorch_cluster) for everything
+- [MLX](https://github.com/ml-explore/mlx) for the foundational framework
+- [MLX-Graphs](https://github.com/mlx-graphs/mlx-graphs) for graph utilities and datasets
+
+## 📊 Citation
+
+If you use MLX-Cluster in your research, please cite:
+
+```bibtex
+@software{mlx_cluster,
+  author = {Vinay Pandya},
+  title = {MLX-Cluster: High-Performance Graph Algorithms for Apple MLX},
+  url = {https://github.com/vinayhpandya/mlx_cluster},
+  version = {0.0.6},
+  year = {2025}
+}
+```
+
+## 🔗 Related Projects
+
+- [MLX](https://github.com/ml-explore/mlx) - Apple's machine learning framework
+- [MLX-Graphs](https://github.com/mlx-graphs/mlx-graphs) - Graph neural networks for MLX
+- [PyTorch Geometric](https://github.com/pyg-team/pytorch_geometric) - Graph deep learning library
+
+---

mlx_cluster-0.0.6/bindings.cpp

@@ -0,0 +1,133 @@
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/variant.h>
+#include <random_walks/RandomWalk.h>
+#include <random_walks/BiasedRandomWalk.h>
+#include <nanobind/stl/vector.h>
+#include <random_walks/NeighborSample.h>
+namespace nb = nanobind;
+using namespace nb::literals;
+using namespace mlx::core;
+
+NB_MODULE(_ext, m){
+
+    m.def(
+      "random_walk",
+      [](const mx::array& rowptr,
+        const mx::array& col,
+        const mx::array& start,
+        const mx::array& rand,
+        int  walk_length,
+        nb::object stream = nb::none()) {
+
+          // call the real C++ implementation
+          auto outs = mlx_random_walk::random_walk(
+              rowptr, col, start, rand, walk_length,
+              stream.is_none() ? mx::StreamOrDevice{}
+                              : nb::cast<mx::StreamOrDevice>(stream));
+
+          // vector -> tuple (move avoids a copy)
+          return nb::make_tuple(std::move(outs[0]), std::move(outs[1]));
+      },
+      "rowptr"_a, "col"_a, "start"_a, "rand"_a, "walk_length"_a,
+      nb::kw_only(), "stream"_a = nb::none(),
+      R"(
+          Uniform random walks.
+
+          Args:
+            rowptr (mlx.core.array): rowptr of graph in csr format.
+            col (mlx.core.array): edges(col) in csr format.
+            start_indices (mlx.core.array): starting nodes of graph from which 
+                            sampling will be performed.
+            random_values (mlx.corearray): random values (between 0 to 1)
+            walk_length (int) : walk length of random graph
+
+          Returns:
+              (nodes, edges) tuple of arrays
+      )",
+      nb::rv_policy::move);
+
+      m.def(
+        "rejection_sampling",
+        [](const mx::array& rowptr,
+          const mx::array& col,
+          const mx::array& start,
+          int walk_length,
+          float p,
+          float q,
+          nb::object stream = nb::none()
+        ){
+          auto outs = mlx_biased_random_walk::rejection_sampling(
+            rowptr, col, start, walk_length, p, q,
+            stream.is_none() ? mx::StreamOrDevice{}
+                              : nb::cast<mx::StreamOrDevice>(stream));
+          return nb::make_tuple(std::move(outs[0]), std::move(outs[1]));
+        },
+        "rowptr"_a,
+        "col"_a,
+        "start"_a,
+        "walk_length"_a,
+        "p"_a,
+        "q"_a,
+        nb::kw_only(), "stream"_a = nb::none(),
+      R"(
+        Sample nodes from the graph by sampling neighbors based
+        on probablity p and q
+
+        Args:
+            rowptr (mlx.core.array): rowptr of graph in csr format.
+            col (mlx.core.array): edges in csr format.
+            start (mlx.core.array): starting node of graph from which 
+                            biased sampling will be performed.
+            walk_length (int) : walk length of random graph
+            p : Likelihood of immediately revisiting a node in the walk.
+            q : Control parameter to interpolate between
+                breadth-first strategy and depth-first strategy
+
+        Returns:
+            (nodes, edges) tuple of arrays
+      )",
+      nb::rv_policy::move);
+
+      m.def(
+        "neighbor_sample",
+        [](const mx::array& colptr,
+           const mx::array& row,
+           const mx::array& input_node,
+           const std::vector<int64_t>& num_neighbors,
+           bool replace = false,
+           bool directed = true) {
+            
+            // Call your C++ function
+            auto result = neighbor_sample(colptr, row, input_node, num_neighbors, replace, directed);
+            
+            // Convert std::tuple to nanobind tuple with move semantics
+            return nb::make_tuple(
+                std::move(std::get<0>(result)),  // samples
+                std::move(std::get<1>(result)),  // rows
+                std::move(std::get<2>(result)),  // cols
+                std::move(std::get<3>(result))   // edges
+            );
+        },
+        "colptr"_a,
+        "row"_a, 
+        "input_node"_a,
+        "num_neighbors"_a,
+        "replace"_a = false,
+        "directed"_a = true,
+        R"(
+            Simple neighbor sampling without primitives.
+            
+            Args:
+                colptr: Column pointers (CSC format)
+                row: Row indices (CSC format)  
+                input_node: Input nodes to sample from
+                num_neighbors: Number of neighbors per hop
+                replace: Sample with replacement
+                directed: Directed graph
+                
+            Returns:
+                tuple: (samples, rows, cols, edges)
+        )",
+        nb::rv_policy::move  // Add this return value policy
+    );
+  }

{mlx_cluster-0.0.4 → mlx_cluster-0.0.6}/mlx_cluster/__init__.py

@@ -2,3 +2,4 @@ import mlx.core as mx
 
 from ._ext import random_walk
 from ._ext import rejection_sampling
+from ._ext import neighbor_sample