mlx-cluster 0.0.1__tar.gz

mlx_cluster-0.0.1/CMakeLists.txt

@@ -0,0 +1,66 @@
+cmake_minimum_required(VERSION 3.27)
+project(_ext LANGUAGES CXX)
+
+# ----- Setup required -----
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+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(
+  COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
+  OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE NB_DIR)
+list(APPEND CMAKE_PREFIX_PATH "${NB_DIR}")
+find_package(nanobind CONFIG REQUIRED)
+
+# ------ Adding extensions to the library -----
+
+# Add library
+add_library(mlx_cluster)
+
+target_sources(mlx_cluster
+PUBLIC
+${CMAKE_CURRENT_LIST_DIR}/random_walks/RandomWalk.cpp
+${CMAKE_CURRENT_LIST_DIR}/random_walks/BiasedRandomWalk.cpp
+)
+
+target_include_directories(mlx_cluster
+PUBLIC
+${CMAKE_CURRENT_LIST_DIR})
+
+target_link_libraries(mlx_cluster PUBLIC mlx)
+
+
+if(MLX_BUILD_METAL)
+  mlx_build_metallib(
+    TARGET mlx_cluster_metallib
+    TITLE mlx_cluster
+    SOURCES ${CMAKE_CURRENT_LIST_DIR}/random_walks/random_walk.metal
+    INCLUDE_DIRS ${PROJECT_SOURCE_DIR} ${MLX_INCLUDE_DIRS}
+    OUTPUT_DIRECTORY ${CMAKE_LIBRARY_OUTPUT_DIRECTORY}
+  )
+
+  add_dependencies(
+    mlx_cluster
+    mlx_cluster_metallib
+  )
+
+endif()
+# ----- Nanobind module -----
+nanobind_add_module(
+  _ext
+  NB_STATIC STABLE_ABI LTO NOMINSIZE
+  NB_DOMAIN mlx 
+  ${CMAKE_CURRENT_LIST_DIR}/bindings.cpp
+)
+
+target_link_libraries(_ext PRIVATE mlx_cluster)
+
+if(BUILD_SHARED_LIBS)
+  target_link_options(_ext PRIVATE -Wl,-rpath,@loader_path)
+endif()

mlx_cluster-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 vinayhpandya
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

mlx_cluster-0.0.1/MANIFEST.in

@@ -0,0 +1 @@
+include CMakeLists.txt README.md bindings.cpp LICENSE random_walks/*.*

mlx_cluster-0.0.1/PKG-INFO

@@ -0,0 +1,74 @@
+Metadata-Version: 2.1
+Name: mlx_cluster
+Version: 0.0.1
+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
+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 :: OS Independent
+Requires-Python: >=3.8
+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"
+
+# mlx_cluster
+
+A C++ extension for generating ramdom walks for Homogeneous graphs using mlx
+
+## Installation
+
+To install the necessary dependencies:
+
+Clone the repositories:
+```bash
+git clone https://github.com/vinayhpandya/mlx_cluster.git
+```
+
+After cloning the repository install library using 
+
+```bash
+python setup.py build_ext -j8 --inplace
+```
+
+for testing purposes you need to have `mlx-graphs` installed
+
+## Usage
+
+
+```
+from mlx_graphs.utils.sorting import sort_edge_index
+from mlx_graphs.loaders import Dataloader
+from mlx_graphs_extension import random_walk
+
+
+cora_dataset = PlanetoidDataset(name="cora", base_dir="~")
+start = mx.arange(0, 1000)
+start_time = time.time()
+edge_index = cora_dataset.graphs[0].edge_index
+num_nodes = cora_dataset.graphs[0].num_nodes
+sorted_edge_index = sort_edge_index(edge_index=edge_index)
+row_mlx = sorted_edge_index[0][0]
+col_mlx = sorted_edge_index[0][1]
+unique_vals, counts_mlx = np.unique(np.array(row_mlx, copy=False), return_counts=True)
+cum_sum_mlx = counts_mlx.cumsum()
+rand = mx.random.uniform(shape=[start.shape[0], 100])
+row_ptr_mlx = mx.concatenate([mx.array([0]), mx.array(cum_sum_mlx)])
+random_walk(row_ptr_mlx, col_mlx, start, rand, 1000, stream = mx.gpu)
+```
+
+## TODO
+
+- [x] Add metal shaders to optimize the code
+- [ ] Benchmark random walk against different frameworks
+- [ ] Add more algorithms
+
+## Credits:
+
+torch_cluster random walk implementation : [random_walk](https://github.com/rusty1s/pytorch_cluster/blob/master/csrc/cpu/rw_cpu.cpp)

mlx_cluster-0.0.1/README.md

@@ -0,0 +1,54 @@
+# mlx_cluster
+
+A C++ extension for generating ramdom walks for Homogeneous graphs using mlx
+
+## Installation
+
+To install the necessary dependencies:
+
+Clone the repositories:
+```bash
+git clone https://github.com/vinayhpandya/mlx_cluster.git
+```
+
+After cloning the repository install library using 
+
+```bash
+python setup.py build_ext -j8 --inplace
+```
+
+for testing purposes you need to have `mlx-graphs` installed
+
+## Usage
+
+
+```
+from mlx_graphs.utils.sorting import sort_edge_index
+from mlx_graphs.loaders import Dataloader
+from mlx_graphs_extension import random_walk
+
+
+cora_dataset = PlanetoidDataset(name="cora", base_dir="~")
+start = mx.arange(0, 1000)
+start_time = time.time()
+edge_index = cora_dataset.graphs[0].edge_index
+num_nodes = cora_dataset.graphs[0].num_nodes
+sorted_edge_index = sort_edge_index(edge_index=edge_index)
+row_mlx = sorted_edge_index[0][0]
+col_mlx = sorted_edge_index[0][1]
+unique_vals, counts_mlx = np.unique(np.array(row_mlx, copy=False), return_counts=True)
+cum_sum_mlx = counts_mlx.cumsum()
+rand = mx.random.uniform(shape=[start.shape[0], 100])
+row_ptr_mlx = mx.concatenate([mx.array([0]), mx.array(cum_sum_mlx)])
+random_walk(row_ptr_mlx, col_mlx, start, rand, 1000, stream = mx.gpu)
+```
+
+## TODO
+
+- [x] Add metal shaders to optimize the code
+- [ ] Benchmark random walk against different frameworks
+- [ ] Add more algorithms
+
+## Credits:
+
+torch_cluster random walk implementation : [random_walk](https://github.com/rusty1s/pytorch_cluster/blob/master/csrc/cpu/rw_cpu.cpp)

mlx_cluster-0.0.1/bindings.cpp

@@ -0,0 +1,65 @@
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/variant.h>
+#include <random_walks/RandomWalk.h>
+#include <random_walks/BiasedRandomWalk.h>
+
+namespace nb = nanobind;
+using namespace nb::literals;
+using namespace mlx::core;
+
+NB_MODULE(_ext, m){
+
+      m.def(
+        "random_walk",
+        &random_walk,
+        "rowptr"_a,
+        "col"_a,
+        "start"_a,
+        "rand"_a,
+        "walk_length"_a,
+        nb::kw_only(),
+      "stream"_a = nb::none(),
+      R"(
+        uniformly sample a graph
+
+
+        Args:
+            rowptr (array): rowptr of graph in csr format.
+            col (array): edges in csr format.
+            walk_length (int) : walk length of random graph
+
+        Returns:
+            array: consisting of nodes visited on random walk
+      )");
+
+      m.def(
+        "rejection_sampling",
+        &rejection_sampling,
+        "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 (array): rowptr of graph in csr format.
+            col (array): edges in csr format.
+            start (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:
+            array: consisting of nodes visited on random walk
+      )");
+}
+

mlx_cluster-0.0.1/mlx_cluster/__init__.py

@@ -0,0 +1,4 @@
+import mlx.core as mx
+
+from ._ext import random_walk
+from ._ext import rejection_sampling

mlx_cluster-0.0.1/mlx_cluster.egg-info/PKG-INFO

@@ -0,0 +1,74 @@
+Metadata-Version: 2.1
+Name: mlx_cluster
+Version: 0.0.1
+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
+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 :: OS Independent
+Requires-Python: >=3.8
+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"
+
+# mlx_cluster
+
+A C++ extension for generating ramdom walks for Homogeneous graphs using mlx
+
+## Installation
+
+To install the necessary dependencies:
+
+Clone the repositories:
+```bash
+git clone https://github.com/vinayhpandya/mlx_cluster.git
+```
+
+After cloning the repository install library using 
+
+```bash
+python setup.py build_ext -j8 --inplace
+```
+
+for testing purposes you need to have `mlx-graphs` installed
+
+## Usage
+
+
+```
+from mlx_graphs.utils.sorting import sort_edge_index
+from mlx_graphs.loaders import Dataloader
+from mlx_graphs_extension import random_walk
+
+
+cora_dataset = PlanetoidDataset(name="cora", base_dir="~")
+start = mx.arange(0, 1000)
+start_time = time.time()
+edge_index = cora_dataset.graphs[0].edge_index
+num_nodes = cora_dataset.graphs[0].num_nodes
+sorted_edge_index = sort_edge_index(edge_index=edge_index)
+row_mlx = sorted_edge_index[0][0]
+col_mlx = sorted_edge_index[0][1]
+unique_vals, counts_mlx = np.unique(np.array(row_mlx, copy=False), return_counts=True)
+cum_sum_mlx = counts_mlx.cumsum()
+rand = mx.random.uniform(shape=[start.shape[0], 100])
+row_ptr_mlx = mx.concatenate([mx.array([0]), mx.array(cum_sum_mlx)])
+random_walk(row_ptr_mlx, col_mlx, start, rand, 1000, stream = mx.gpu)
+```
+
+## TODO
+
+- [x] Add metal shaders to optimize the code
+- [ ] Benchmark random walk against different frameworks
+- [ ] Add more algorithms
+
+## Credits:
+
+torch_cluster random walk implementation : [random_walk](https://github.com/rusty1s/pytorch_cluster/blob/master/csrc/cpu/rw_cpu.cpp)

mlx_cluster-0.0.1/mlx_cluster.egg-info/SOURCES.txt

@@ -0,0 +1,24 @@
+CMakeLists.txt
+LICENSE
+MANIFEST.in
+README.md
+bindings.cpp
+pyproject.toml
+setup.py
+mlx_cluster/__init__.py
+mlx_cluster/_ext.cpython-311-darwin.so
+mlx_cluster/libmlx_cluster.dylib
+mlx_cluster/mlx_cluster.metallib
+mlx_cluster.egg-info/PKG-INFO
+mlx_cluster.egg-info/SOURCES.txt
+mlx_cluster.egg-info/dependency_links.txt
+mlx_cluster.egg-info/not-zip-safe
+mlx_cluster.egg-info/requires.txt
+mlx_cluster.egg-info/top_level.txt
+random_walks/BiasedRandomWalk.cpp
+random_walks/BiasedRandomWalk.h
+random_walks/RandomWalk.cpp
+random_walks/RandomWalk.h
+random_walks/random_walk.metal
+tests/test_random_walk.py
+tests/test_rejection_sampling.py

mlx_cluster-0.0.1/mlx_cluster.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

mlx_cluster-0.0.1/mlx_cluster.egg-info/not-zip-safe

@@ -0,0 +1 @@
+

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

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

mlx_cluster-0.0.1/mlx_cluster.egg-info/top_level.txt

@@ -0,0 +1 @@
+mlx_cluster

mlx_cluster-0.0.1/pyproject.toml

@@ -0,0 +1,33 @@
+[project]
+name = "mlx_cluster"
+version = "0.0.1"
+authors = [
+  { name = "Vinay Pandya", email = "vinayharshadpandya27@gmail.com" },
+]
+description = "C++ and Metal extensions for MLX CTC Loss"
+readme = "README.md"
+requires-python = ">=3.8"
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: C++",
+  "License :: OSI Approved :: MIT License",
+  "Operating System :: OS Independent",
+]
+
+[project.optional-dependencies]
+dev = []
+
+[project.urls]
+Homepage = "https://github.com/vinayhpandya/mlx_cluster"
+Issues = "https://github.com/vinayhpandya/mlx_cluster/Issues"
+
+
+[build-system]
+requires = [
+  "setuptools>=42",
+  "cmake>=3.24",
+  "mlx>=0.9.0",
+  "nanobind@git+https://github.com/wjakob/nanobind.git@2f04eac452a6d9142dedb957701bdb20125561e4",
+]
+build-backend = "setuptools.build_meta"

mlx_cluster-0.0.1/random_walks/BiasedRandomWalk.cpp

@@ -0,0 +1,194 @@
+#include <cassert>
+#include <iostream>
+#include <sstream>
+#include <random>
+
+#include "mlx/backend/common/copy.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/utils.h"
+#include "mlx/random.h"
+#include "mlx/ops.h"
+#include "mlx/array.h"
+#ifdef _METAL_
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/utils.h"
+#endif
+#include "random_walks/BiasedRandomWalk.h"
+
+namespace mlx::core {
+
+    bool inline is_neighbor(const int64_t *rowptr, const int64_t *col, int64_t v,
+                        int64_t w) {
+        int64_t row_start = rowptr[v], row_end = rowptr[v + 1];
+        for (auto i = row_start; i < row_end; i++) {
+            if (col[i] == w)
+            return true;
+            }
+        return false;
+    }
+
+    void BiasedRandomWalk::eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs) {
+        auto& rowptr = inputs[0];
+        auto& col = inputs[1];
+        auto& start = inputs[2];
+        auto& rand = inputs[3];
+        int numel = start.size();
+
+        // Initialize outputs
+        assert(outputs.size() == 2);
+        // Allocate memory for outputs if not already allocated
+        outputs[0].set_data(allocator::malloc_or_wait(numel*(walk_length_+1)*sizeof(int64_t)));
+        outputs[1].set_data(allocator::malloc_or_wait(numel*walk_length_*sizeof(int64_t)));
+        auto& n_out = outputs[0];
+        auto& e_out = outputs[1];
+
+        auto* n_out_ptr = n_out.data<int64_t>();
+        auto* e_out_ptr = e_out.data<int64_t>();
+        auto* start_values = start.data<int64_t>();
+        auto* row_ptr = rowptr.data<int64_t>();
+        auto* col_values = col.data<int64_t>();
+        auto* rand_values = rand.data<float>();
+
+        double max_prob = fmax(fmax(1. / p_, 1.), 1. / q_);
+        double prob_0 = 1. / p_ / max_prob;
+        double prob_1 = 1. / max_prob;
+        double prob_2 = 1. / q_ / max_prob;
+
+        for (int64_t n = 0; n < numel; n++) {
+            int64_t t = start_values[n], v, x, e_cur, row_start, row_end;
+            n_out_ptr[n * (walk_length_ + 1)] = t;
+            row_start = row_ptr[t], row_end = row_ptr[t + 1];
+            if (row_end - row_start == 0) {
+                e_cur = -1;
+                v = t;
+            } else {
+                e_cur = row_start + (std::rand() % (row_end - row_start));
+                v = col_values[e_cur];
+            }
+            n_out_ptr[n * (walk_length_ + 1) + 1] = v;
+            e_out_ptr[n * walk_length_] = e_cur;
+            for (auto l = 1; l < walk_length_; l++) {
+                row_start = row_ptr[v], row_end = row_ptr[v + 1];
+
+                if (row_end - row_start == 0) {
+                e_cur = -1;
+                x = v;
+                } else if (row_end - row_start == 1) {
+                e_cur = row_start;
+                x = col_values[e_cur];
+                } else {
+                while (true) {
+                    e_cur = row_start + (std::rand() % (row_end - row_start));
+                    x = col_values[e_cur];
+
+                    auto r = ((double)std::rand() / (RAND_MAX)); // [0, 1)
+
+                    if (x == t && r < prob_0)
+                    break;
+                    else if (is_neighbor(row_ptr, col_values, x, t) && r < prob_1)
+                    break;
+                    else if (r < prob_2)
+                    break;
+                    }
+                }
+
+                n_out_ptr[n * (walk_length_ + 1) + (l + 1)] = x;
+                e_out_ptr[n * walk_length_ + l] = e_cur;
+                t = v;
+                v = x;
+            }
+        }
+   
+    };
+
+    std::vector<array> BiasedRandomWalk::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums)
+{
+    // Random walk is not differentiable, so we return zero tangents
+    throw std::runtime_error("Biased random walk has no jvp implementation.");
+}
+// #ifdef _METAL_
+// void BiasedRandomWalk::eval_gpu(
+//         const std::vector<array>& inputs,
+//         std::vector<array>& outputs
+//     ){
+//        auto& rowptr = inputs[0];
+//         auto& col = inputs[1];
+//         auto& start = inputs[2];
+//         auto& rand = inputs[3];
+//         int numel = start.size();
+        
+//         assert(outputs.size() == 2);
+//         outputs[0].set_data(allocator::malloc_or_wait(numel * (walk_length_ + 1) * sizeof(int64_t)));
+//         outputs[1].set_data(allocator::malloc_or_wait(numel * walk_length_ * sizeof(int64_t)));
+//         std::cout<<"after setting data"<<std::endl;
+//         auto& s = stream();
+//         auto& d = metal::device(s.device);
+
+//         d.register_library("mlx_cluster", metal::get_colocated_mtllib_path);
+//         auto kernel = d.get_kernel("random_walk", "mlx_cluster");
+
+//         auto& compute_encoder = d.get_command_encoder(s.index);
+//         compute_encoder->setComputePipelineState(kernel);
+
+//         compute_encoder.set_input_array(rowptr, 0);
+//         compute_encoder.set_input_array(col, 1);
+//         compute_encoder.set_input_array(start, 2);
+//         compute_encoder.set_input_array(rand, 3);
+//         compute_encoder.set_output_array(outputs[0], 4);
+//         compute_encoder.set_output_array(outputs[1], 5);
+//         compute_encoder->setBytes(&walk_length_, sizeof(int32), 6);
+
+//         MTL::Size grid_size = MTL::Size(numel, 1, 1);
+//         MTL::Size thread_group_size = MTL::Size(kernel->maxTotalThreadsPerThreadgroup(), 1, 1);
+
+//         compute_encoder.dispatchThreads(grid_size, thread_group_size);
+//     }
+// #endif
+void BiasedRandomWalk::eval_gpu(
+       const std::vector<array>& inputs, std::vector<array>& outputs
+  )
+{
+    throw std::runtime_error("Random walk has no GPU implementation.");
+}
+std::vector<array> BiasedRandomWalk::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs)
+{
+    // Random walk is not differentiable, so we return zero gradients
+    throw std::runtime_error("Random walk has no JVP implementation.");
+}
+
+std::pair<std::vector<array>, std::vector<int>> BiasedRandomWalk::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes)
+{
+    throw std::runtime_error("vmap not implemented for biasedRandomWalk");
+}
+
+bool BiasedRandomWalk::is_equivalent(const Primitive& other) const
+{
+    throw std::runtime_error("biased Random walk has no GPU implementation.");
+}
+
+std::vector<std::vector<int>> BiasedRandomWalk::output_shapes(const std::vector<array>& inputs)
+{
+    throw std::runtime_error("biased Random walk has no GPU implementation.");
+}
+
+array rejection_sampling(const array& rowptr, const array& col, const array& start, int walk_length, const double p, 
+       const double q, StreamOrDevice s)
+{   
+    int nodes = start.size();
+    auto primitive = std::make_shared<BiasedRandomWalk>(to_stream(s), walk_length, p, q);
+    return array::make_arrays({{nodes,walk_length+1},{nodes, walk_length}},
+     {rowptr.dtype(), rowptr.dtype()},
+     primitive,
+     {rowptr, col, start}
+    )[0];
+}
+}

mlx_cluster-0.0.1/random_walks/BiasedRandomWalk.h

@@ -0,0 +1,66 @@
+#pragma once
+
+#include <mlx/array.h>
+#include <mlx/ops.h>
+#include <mlx/primitives.h>
+
+namespace mlx::core{
+
+    class BiasedRandomWalk : public Primitive {
+        public:
+            BiasedRandomWalk(Stream stream, int walk_length, double p, double q)
+            : Primitive(stream), walk_length_(walk_length), p_(p), q_(q) {}
+            void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+            override;
+            void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+            override;
+
+            /** The Jacobian-vector product. */
+            std::vector<array> jvp(
+                const std::vector<array>& primals,
+                const std::vector<array>& tangents,
+                const std::vector<int>& argnums) override;
+
+            /** The vector-Jacobian product. */
+            std::vector<array> vjp(
+                const std::vector<array>& primals,
+                const std::vector<array>& cotangents,
+                const std::vector<int>& argnums,
+                const std::vector<array>& outputs) override;
+
+            /**
+             * The primitive must know how to vectorize itself across
+             * the given axes. The output is a pair containing the array
+             * representing the vectorized computation and the axis which
+             * corresponds to the output vectorized dimension.
+             */
+            std::pair<std::vector<array>, std::vector<int>> vmap(
+                const std::vector<array>& inputs,
+                const std::vector<int>& axes) override;
+
+            /** Print the primitive. */
+            void print(std::ostream& os) override {
+                os << "biased random walk implementation";
+            }
+
+            /** Equivalence check **/
+            bool is_equivalent(const Primitive& other) const override;
+
+            std::vector<std::vector<int>> output_shapes(const std::vector<array>& inputs) override;
+        
+        private:
+            int walk_length_;
+            double p_;
+            double q_;
+
+    };
+
+    array rejection_sampling(const array& rowptr,
+     const array& col,
+    const array& start,
+       int walk_length,
+       const double p, 
+       const double q,
+        StreamOrDevice s = {});
+
+};

mlx_cluster-0.0.1/random_walks/RandomWalk.cpp

@@ -0,0 +1,147 @@
+#include <cassert>
+#include <iostream>
+#include <sstream>
+
+#include "mlx/backend/common/copy.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/utils.h"
+#include "mlx/random.h"
+#include "mlx/ops.h"
+#include "mlx/array.h"
+#ifdef _METAL_
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/utils.h"
+#endif
+#include "random_walks/RandomWalk.h"
+
+namespace mlx::core {
+    void RandomWalk::eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs) {
+    auto& rowptr = inputs[0];
+    auto& col = inputs[1];
+    auto& start = inputs[2];
+    auto& rand = inputs[3];
+    int numel = start.size();
+
+    // Initialize outputs
+    assert(outputs.size() == 2);
+    // Allocate memory for outputs if not already allocated
+    outputs[0].set_data(allocator::malloc_or_wait(numel*(walk_length_+1)*sizeof(int64_t)));
+    outputs[1].set_data(allocator::malloc_or_wait(numel*walk_length_*sizeof(int64_t)));
+    auto& n_out = outputs[0];
+    auto& e_out = outputs[1];
+
+    auto* n_out_ptr = n_out.data<int64_t>();
+    auto* e_out_ptr = e_out.data<int64_t>();
+    auto* start_values = start.data<int64_t>();
+    auto* row_ptr = rowptr.data<int64_t>();
+    auto* col_values = col.data<int64_t>();
+    auto* rand_values = rand.data<float>();
+
+    for (int64_t n = 0; n < numel; n++) {
+        int64_t n_cur = start_values[n];
+        n_out_ptr[n * (walk_length_ + 1)] = n_cur;
+        for (int l = 0; l < walk_length_; l++) {
+            int64_t row_start = row_ptr[n_cur];
+            int64_t row_end = row_ptr[n_cur+1];
+            int64_t e_cur;
+            if (row_end - row_start == 0) {
+                e_cur = -1;
+            } else {
+                float r = rand_values[n*walk_length_+l];
+                int64_t idx = static_cast<int64_t>(r * (row_end - row_start));
+                e_cur = row_start + idx;
+                n_cur = col_values[e_cur];
+            }
+
+            n_out_ptr[n * (walk_length_ + 1) + (l + 1)] = n_cur;
+            e_out_ptr[n * walk_length_ + l] = e_cur;
+        }
+    }
+   
+    };
+
+    std::vector<array> RandomWalk::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums)
+{
+    // Random walk is not differentiable, so we return zero tangents
+    throw std::runtime_error("Random walk has no GPU implementation.");
+}
+#ifdef _METAL_
+void RandomWalk::eval_gpu(
+        const std::vector<array>& inputs,
+        std::vector<array>& outputs
+    ){
+       auto& rowptr = inputs[0];
+        auto& col = inputs[1];
+        auto& start = inputs[2];
+        auto& rand = inputs[3];
+        int numel = start.size();
+        
+        assert(outputs.size() == 2);
+        outputs[0].set_data(allocator::malloc_or_wait(numel * (walk_length_ + 1) * sizeof(int64_t)));
+        outputs[1].set_data(allocator::malloc_or_wait(numel * walk_length_ * sizeof(int64_t)));
+        std::cout<<"after setting data"<<std::endl;
+        auto& s = stream();
+        auto& d = metal::device(s.device);
+
+        d.register_library("mlx_cluster");
+        auto kernel = d.get_kernel("random_walk", "mlx_cluster");
+
+        auto& compute_encoder = d.get_command_encoder(s.index);
+        compute_encoder->setComputePipelineState(kernel);
+
+        compute_encoder.set_input_array(rowptr, 0);
+        compute_encoder.set_input_array(col, 1);
+        compute_encoder.set_input_array(start, 2);
+        compute_encoder.set_input_array(rand, 3);
+        compute_encoder.set_output_array(outputs[0], 4);
+        compute_encoder.set_output_array(outputs[1], 5);
+        compute_encoder->setBytes(&walk_length_, sizeof(int32), 6);
+
+        MTL::Size grid_size = MTL::Size(numel, 1, 1);
+        MTL::Size thread_group_size = MTL::Size(kernel->maxTotalThreadsPerThreadgroup(), 1, 1);
+
+        compute_encoder.dispatchThreads(grid_size, thread_group_size);
+    }
+#endif
+
+std::vector<array> RandomWalk::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs)
+{
+    // Random walk is not differentiable, so we return zero gradients
+    throw std::runtime_error("Random walk has no GPU implementation.");
+}
+
+std::pair<std::vector<array>, std::vector<int>> RandomWalk::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes)
+{
+    throw std::runtime_error("vmap not implemented for RandomWalk");
+}
+
+bool RandomWalk::is_equivalent(const Primitive& other) const
+{
+    throw std::runtime_error("Random walk has no GPU implementation.");
+}
+
+std::vector<std::vector<int>> RandomWalk::output_shapes(const std::vector<array>& inputs)
+{
+    throw std::runtime_error("Random walk has no GPU implementation.");
+}
+
+array random_walk(const array& rowptr, const array& col, const array& start, const array& rand, int walk_length, StreamOrDevice s)
+{   
+    int nodes = start.size();
+    auto primitive = std::make_shared<RandomWalk>(walk_length, to_stream(s));
+    return array::make_arrays({{nodes,walk_length+1},{nodes, walk_length}},
+     {start.dtype(), start.dtype()},
+     primitive,
+     {rowptr, col, start, rand}
+    )[0];
+}
+}

mlx_cluster-0.0.1/random_walks/RandomWalk.h

@@ -0,0 +1,63 @@
+#pragma once
+
+#include <mlx/array.h>
+#include <mlx/ops.h>
+#include <mlx/primitives.h>
+
+namespace mlx::core{
+
+    class RandomWalk : public Primitive {
+        public:
+            explicit RandomWalk(int walk_length, Stream stream):
+                Primitive(stream), walk_length_(walk_length) {};
+            void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+            override;
+            void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+            override;
+
+            /** The Jacobian-vector product. */
+            std::vector<array> jvp(
+                const std::vector<array>& primals,
+                const std::vector<array>& tangents,
+                const std::vector<int>& argnums) override;
+
+            /** The vector-Jacobian product. */
+            std::vector<array> vjp(
+                const std::vector<array>& primals,
+                const std::vector<array>& cotangents,
+                const std::vector<int>& argnums,
+                const std::vector<array>& outputs) override;
+
+            /**
+             * The primitive must know how to vectorize itself across
+             * the given axes. The output is a pair containing the array
+             * representing the vectorized computation and the axis which
+             * corresponds to the output vectorized dimension.
+             */
+            std::pair<std::vector<array>, std::vector<int>> vmap(
+                const std::vector<array>& inputs,
+                const std::vector<int>& axes) override;
+
+            /** Print the primitive. */
+            void print(std::ostream& os) override {
+                os << "Random walk implementation";
+            }
+
+            /** Equivalence check **/
+            bool is_equivalent(const Primitive& other) const override;
+
+            std::vector<std::vector<int>> output_shapes(const std::vector<array>& inputs) override;
+        
+        private:
+            int walk_length_;
+
+    };
+
+    array random_walk(const array& rowptr,
+     const array& col,
+    const array& start,
+    const array& rand,
+       int walk_length,
+        StreamOrDevice s = {});
+
+};

mlx_cluster-0.0.1/random_walks/random_walk.metal

@@ -0,0 +1,35 @@
+#include <metal_stdlib>
+#include "mlx/backend/metal/kernels/utils.h"
+using namespace metal;
+
+kernel void random_walk(
+    const device int64_t* rowptr [[buffer(0)]],
+    const device int64_t* col [[buffer(1)]],
+    const device int64_t* start [[buffer(2)]],
+    const device float* rand [[buffer(3)]],
+    device int64_t* n_out [[buffer(4)]],
+    device int64_t* e_out [[buffer(5)]],
+    constant int& walk_length [[buffer(6)]],
+    uint tid [[thread_position_in_grid]]
+) {
+    int64_t n_cur = start[tid];
+    n_out[tid * (walk_length + 1)] = n_cur;
+    
+    for (int l = 0; l < walk_length; l++) {
+        int64_t row_start = rowptr[n_cur];
+        int64_t row_end = rowptr[n_cur + 1];
+        int64_t e_cur;
+        
+        if (row_end - row_start == 0) {
+            e_cur = -1;
+        } else {
+            float r = rand[tid * walk_length + l];
+            int64_t idx = static_cast<int64_t>(r * (row_end - row_start));
+            e_cur = row_start + idx;
+            n_cur = col[e_cur];
+        }
+        
+        n_out[tid * (walk_length + 1) + (l + 1)] = n_cur;
+        e_out[tid * walk_length + l] = e_cur;
+    }
+}

mlx_cluster-0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

mlx_cluster-0.0.1/setup.py

@@ -0,0 +1,16 @@
+from setuptools import setup
+from mlx import extension
+
+if __name__ == "__main__":
+    setup(
+        name="mlx_cluster",
+        version="0.0.1",
+        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"]},
+        zip_safe=False,
+        python_requires=">=3.8",
+    )

mlx_cluster-0.0.1/tests/test_random_walk.py

@@ -0,0 +1,58 @@
+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)

mlx_cluster-0.0.1/tests/test_rejection_sampling.py

@@ -0,0 +1,49 @@
+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))
+# random_walks = torch.ops.torch_cluster.random_walk(
+#     row_ptr, col, start_indices, 5, 1.0, 3.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 rejection_sampling
+
+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)
+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("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(node_sequence)