mlx 1.0.0

data/mlx/README.md

@@ -0,0 +1,121 @@
+# MLX
+
+[**Quickstart**](#quickstart) | [**Installation**](#installation) |
+[**Documentation**](https://ml-explore.github.io/mlx/build/html/index.html) |
+[**Examples**](#examples)
+
+[![CircleCI](https://circleci.com/gh/ml-explore/mlx.svg?style=svg)](https://circleci.com/gh/ml-explore/mlx)
+
+MLX is an array framework for machine learning on Apple silicon,
+brought to you by Apple machine learning research.
+
+Some key features of MLX include:
+
+- **Familiar APIs**: MLX has a Python API that closely follows NumPy. MLX
+   also has fully featured C++, [C](https://github.com/ml-explore/mlx-c), and
+   [Swift](https://github.com/ml-explore/mlx-swift/) APIs, which closely mirror
+   the Python API. MLX has higher-level packages like `mlx.nn` and
+   `mlx.optimizers` with APIs that closely follow PyTorch to simplify building
+   more complex models.
+
+- **Composable function transformations**: MLX supports composable function
+  transformations for automatic differentiation, automatic vectorization,
+  and computation graph optimization.
+
+- **Lazy computation**: Computations in MLX are lazy. Arrays are only
+  materialized when needed.
+
+- **Dynamic graph construction**: Computation graphs in MLX are constructed
+  dynamically. Changing the shapes of function arguments does not trigger
+  slow compilations, and debugging is simple and intuitive.
+
+- **Multi-device**: Operations can run on any of the supported devices
+  (currently the CPU and the GPU).
+
+- **Unified memory**: A notable difference from MLX and other frameworks
+  is the *unified memory model*. Arrays in MLX live in shared memory.
+  Operations on MLX arrays can be performed on any of the supported
+  device types without transferring data.
+
+MLX is designed by machine learning researchers for machine learning
+researchers. The framework is intended to be user-friendly, but still efficient
+to train and deploy models. The design of the framework itself is also
+conceptually simple. We intend to make it easy for researchers to extend and
+improve MLX with the goal of quickly exploring new ideas.
+
+The design of MLX is inspired by frameworks like
+[NumPy](https://numpy.org/doc/stable/index.html),
+[PyTorch](https://pytorch.org/), [Jax](https://github.com/google/jax), and
+[ArrayFire](https://arrayfire.org/).
+
+## Examples
+
+The [MLX examples repo](https://github.com/ml-explore/mlx-examples) has a
+variety of examples, including:
+
+- [Transformer language model](https://github.com/ml-explore/mlx-examples/tree/main/transformer_lm) training.
+- Large-scale text generation with
+  [LLaMA](https://github.com/ml-explore/mlx-examples/tree/main/llms/llama) and
+  finetuning with [LoRA](https://github.com/ml-explore/mlx-examples/tree/main/lora).
+- Generating images with [Stable Diffusion](https://github.com/ml-explore/mlx-examples/tree/main/stable_diffusion).
+- Speech recognition with [OpenAI's Whisper](https://github.com/ml-explore/mlx-examples/tree/main/whisper).
+
+## Quickstart
+
+See the [quick start
+guide](https://ml-explore.github.io/mlx/build/html/usage/quick_start.html)
+in the documentation.
+
+## Installation
+
+MLX is available on [PyPI](https://pypi.org/project/mlx/). To install MLX on
+macOS, run:
+
+```bash
+pip install mlx
+```
+
+To install the CUDA backend on Linux, run:
+
+```bash
+pip install mlx[cuda]
+```
+
+To install a CPU-only Linux package, run:
+
+```bash
+pip install mlx[cpu]
+```
+
+Checkout the
+[documentation](https://ml-explore.github.io/mlx/build/html/install.html#)
+for more information on building the C++ and Python APIs from source.
+
+## Contributing
+
+Check out the [contribution guidelines](https://github.com/ml-explore/mlx/tree/main/CONTRIBUTING.md) for more information
+on contributing to MLX. See the
+[docs](https://ml-explore.github.io/mlx/build/html/install.html) for more
+information on building from source, and running tests.
+
+We are grateful for all of [our
+contributors](https://github.com/ml-explore/mlx/tree/main/ACKNOWLEDGMENTS.md#Individual-Contributors). If you contribute
+to MLX and wish to be acknowledged, please add your name to the list in your
+pull request.
+
+## Citing MLX
+
+The MLX software suite was initially developed with equal contribution by Awni
+Hannun, Jagrit Digani, Angelos Katharopoulos, and Ronan Collobert. If you find
+MLX useful in your research and wish to cite it, please use the following
+BibTex entry:
+
+```text
+@software{mlx2023,
+  author = {Awni Hannun and Jagrit Digani and Angelos Katharopoulos and Ronan Collobert},
+  title = {{MLX}: Efficient and flexible machine learning on Apple silicon},
+  url = {https://github.com/ml-explore},
+  version = {0.0},
+  year = {2023},
+}
+```

data/mlx/benchmarks/cpp/CMakeLists.txt

@@ -0,0 +1,11 @@
+function(build_benchmark SRCFILE)
+  get_filename_component(src_name ${SRCFILE} NAME_WE)
+  set(target "${src_name}")
+  add_executable(${target} ${SRCFILE})
+  target_link_libraries(${target} PRIVATE mlx)
+endfunction(build_benchmark)
+
+build_benchmark(single_ops.cpp)
+build_benchmark(irregular_strides.cpp)
+build_benchmark(compare_devices.cpp)
+build_benchmark(autograd.cpp)

data/mlx/benchmarks/cpp/autograd.cpp

@@ -0,0 +1,39 @@
+// Copyright © 2023 Apple Inc.
+
+#include <iostream>
+
+#include "mlx/mlx.h"
+#include "time_utils.h"
+
+namespace mx = mlx::core;
+
+void time_value_and_grad() {
+  auto x = mx::ones({200, 1000});
+  mx::eval(x);
+  auto fn = [](mx::array x) {
+    for (int i = 0; i < 20; ++i) {
+      x = mx::log(mx::exp(x));
+    }
+    return mx::sum(x);
+  };
+
+  auto grad_fn = mx::grad(fn);
+  auto independent_value_and_grad = [&]() {
+    auto value = fn(x);
+    auto dfdx = grad_fn(x);
+    return std::vector<mx::array>{value, dfdx};
+  };
+  TIME(independent_value_and_grad);
+
+  auto value_and_grad_fn = mx::value_and_grad(fn);
+  auto combined_value_and_grad = [&]() {
+    auto [value, dfdx] = value_and_grad_fn(x);
+    return std::vector<mx::array>{value, dfdx};
+  };
+  TIME(combined_value_and_grad);
+}
+
+int main() {
+  std::cout << "Benchmarks for " << mx::default_device() << std::endl;
+  time_value_and_grad();
+}

data/mlx/benchmarks/cpp/compare_devices.cpp

@@ -0,0 +1,27 @@
+// Copyright © 2023 Apple Inc.
+
+#include <iostream>
+#include "mlx/mlx.h"
+#include "time_utils.h"
+
+namespace mx = mlx::core;
+
+void time_add_op() {
+  std::vector<int> sizes(1, 1);
+  for (int i = 0; i < 9; ++i) {
+    sizes.push_back(10 * sizes.back());
+  }
+  set_default_device(mx::Device::cpu);
+  for (auto size : sizes) {
+    auto a = mx::random::uniform({size});
+    auto b = mx::random::uniform({size});
+    mx::eval(a, b);
+    std::cout << "Size " << size << std::endl;
+    TIMEM("cpu", mx::add, a, b, mx::Device::cpu);
+    TIMEM("gpu", mx::add, a, b, mx::Device::gpu);
+  }
+}
+
+int main() {
+  time_add_op();
+}

data/mlx/benchmarks/cpp/irregular_strides.cpp

@@ -0,0 +1,201 @@
+// Copyright © 2023 Apple Inc.
+
+#include <cstring>
+#include <iostream>
+#include <sstream>
+
+#include "mlx/mlx.h"
+#include "time_utils.h"
+
+namespace mx = mlx::core;
+
+void time_irregular_binary_ops_1D() {
+  auto device = mx::default_device();
+  int size = 1000000;
+  int step = 2;
+  auto a = mx::random::uniform({size});
+  auto b = mx::random::uniform({size});
+  mx::eval(a, b);
+  a = slice(a, {0}, {size}, {step});
+  b = slice(b, {0}, {size}, {step});
+  TIMEM("1D strided", mx::add, a, b, device);
+}
+
+void time_irregular_binary_ops_2D() {
+  auto device = mx::default_device();
+  int size = 2048;
+  auto a = mx::random::uniform({size, size});
+  auto b = mx::random::uniform({size, size});
+  mx::eval(a, b);
+  TIMEM("2D regular", mx::add, a, b, device);
+
+  b = mx::transpose(b);
+  mx::eval(b);
+  TIMEM("2D mx::transpose", mx::add, a, b, device);
+
+  b = mx::random::uniform({size});
+  mx::eval(b);
+  TIMEM("2D broadcast dim 0", mx::add, a, b, device);
+
+  b = mx::reshape(b, {size, 1});
+  mx::eval(b);
+  TIMEM("2D broadcast dim 1", mx::add, a, b, device);
+}
+
+void time_irregular_binary_ops_3D() {
+  auto device = mx::default_device();
+  int d0 = 32;
+  int d1 = 512;
+  int d2 = 512;
+  auto a = mx::random::uniform({d0, d1, d2});
+  auto b = mx::random::uniform({d0, d1, d2});
+  TIMEM("3D regular", mx::add, a, b, device);
+
+  b = mx::transpose(b, {0, 2, 1});
+  TIMEM("3D mx::transpose", mx::add, a, b, device);
+
+  b = mx::random::uniform({d1, d2});
+  TIMEM("3D broadcast dim 0", mx::add, a, b, device);
+
+  b = mx::random::uniform({d0, 1, d2});
+  TIMEM("3D broadcast dim 1", mx::add, a, b, device);
+
+  b = mx::random::uniform({d0, d1, 1});
+  TIMEM("3D broadcast dim 2", mx::add, a, b, device);
+
+  b = mx::random::uniform({d2});
+  TIMEM("3D broadcast dims 0, 1", mx::add, a, b, device);
+
+  b = mx::random::uniform({d1, 1});
+  TIMEM("3D broadcast dims 0, 2", mx::add, a, b, device);
+
+  b = mx::random::uniform({d0, 1, 1});
+  TIMEM("3D broadcast dims 1, 2", mx::add, a, b, device);
+}
+
+void time_irregular_binary_ops_4D() {
+  auto device = mx::default_device();
+  mx::Shape shape = {8, 8, 512, 512};
+  auto a = mx::random::uniform(shape);
+  auto b = mx::random::uniform(shape);
+
+  TIMEM("4D regular", mx::add, a, b, device);
+
+  b = mx::transpose(b, {0, 1, 3, 2});
+  TIMEM("4D mx::transpose", mx::add, a, b, device);
+
+  std::string om = "4D broadcast dims ";
+  for (int i = 0; i < shape.size(); ++i) {
+    shape[i] = 1;
+    b = mx::random::uniform(shape);
+    std::ostringstream msg;
+    msg << om << i;
+    TIMEM(msg.str(), mx::add, a, b, device);
+
+    for (int j = i + 1; j < shape.size(); ++j) {
+      shape[j] = 1;
+      std::ostringstream msg;
+      msg << om << i << ", " << j;
+      b = mx::random::uniform(shape);
+      TIMEM(msg.str(), mx::add, a, b, device);
+      shape[j] = a.shape(j);
+
+      for (int k = j + 1; k < shape.size(); ++k) {
+        shape[k] = 1;
+        std::ostringstream msg;
+        msg << om << i << ", " << j << ", " << k;
+        b = mx::random::uniform(shape);
+        TIMEM(msg.str(), mx::add, a, b, device);
+        shape[k] = a.shape(k);
+      }
+    }
+    shape[i] = a.shape(i);
+  }
+}
+
+void time_irregular_reshape() {
+  auto device = mx::default_device();
+  mx::Shape shape;
+  auto reshape_fn = [&shape, device](const mx::array& a) {
+    return mx::reshape(a, shape, device);
+  };
+
+  int size = 64;
+  int d = 2 * size;
+
+  auto a = mx::random::uniform({d, d, d});
+
+  shape = {8 * size, size, size};
+  TIMEM("3D contiguous", reshape_fn, a);
+
+  a = mx::transpose(a);
+  shape = {8 * size, size, size};
+  TIMEM("3D mx::transpose", reshape_fn, a);
+
+  a = mx::transpose(a, {1, 2, 0});
+  shape = {8 * size, size, size};
+  TIMEM("3D mx::transpose dims 1 2", reshape_fn, a);
+
+  a = mx::broadcast_to(mx::random::uniform({d, d}), {d, d, d});
+  TIMEM("3D broadcast dim 0", reshape_fn, a);
+
+  a = mx::broadcast_to(mx::random::uniform({d, 1, d}), {d, d, d});
+  TIMEM("3D broadcast dim 1", reshape_fn, a);
+
+  a = mx::broadcast_to(mx::random::uniform({d, d, 1}), {d, d, d});
+  TIMEM("3D broadcast dim 2", reshape_fn, a);
+
+  a = mx::broadcast_to(mx::random::uniform({d}), {d, d, d});
+  TIMEM("3D broadcast dims 0, 1", reshape_fn, a);
+
+  a = mx::broadcast_to(mx::random::uniform({d, 1}), {d, d, d});
+  TIMEM("3D broadcast dims 0, 2", reshape_fn, a);
+
+  a = mx::broadcast_to(mx::random::uniform({d, 1, 1}), {d, d, d});
+  TIMEM("3D broadcast dims 1, 2", reshape_fn, a);
+
+  a = mx::broadcast_to(mx::random::uniform({1, 1, 1}), {d, d, d});
+  TIMEM("3D broadcast dims 1, 2, 3", reshape_fn, a);
+}
+
+void time_irregular_astype_1D() {
+  auto device = mx::default_device();
+  int size = 1000000;
+  int step = 2;
+  auto a = mx::random::uniform({size});
+  a = slice(a, {0}, {size}, {step});
+  TIMEM("1D strided", mx::astype, a, mx::int32, device);
+}
+
+void time_irregular_astype_2D() {
+  auto device = mx::default_device();
+  int size = 2048;
+  mx::Shape shape = {size, size};
+
+  auto a = mx::random::uniform(shape);
+  TIMEM("2D regular", mx::astype, a, mx::int32, device);
+
+  a = mx::transpose(a);
+  TIMEM("2D mx::transpose", mx::astype, a, mx::int32, device);
+
+  a = mx::broadcast_to(mx::random::uniform({size}), shape);
+  TIMEM("2D broadcast dim 0", mx::astype, a, mx::int32, device);
+
+  a = mx::broadcast_to(mx::random::uniform({size, 1}), shape);
+  TIMEM("2D broadcast dim 1", mx::astype, a, mx::int32, device);
+}
+
+int main(int argc, char** argv) {
+  if (argc > 1) {
+    bool use_gpu = !strcmp(argv[1], "gpu");
+    set_default_device(use_gpu ? mx::Device::gpu : mx::Device::cpu);
+  }
+  std::cout << "Benchmarks for " << mx::default_device() << std::endl;
+  time_irregular_binary_ops_1D();
+  time_irregular_binary_ops_2D();
+  time_irregular_binary_ops_3D();
+  time_irregular_binary_ops_4D();
+  time_irregular_reshape();
+  time_irregular_astype_1D();
+  time_irregular_astype_2D();
+}

data/mlx/benchmarks/cpp/single_ops.cpp

@@ -0,0 +1,288 @@
+// Copyright © 2023 Apple Inc.
+
+#include "mlx/mlx.h"
+#include "time_utils.h"
+
+namespace mx = mlx::core;
+
+void time_creation_ops() {
+  int M = 2000;
+  int N = 500;
+  auto shape = {M, N};
+  auto full_fp32 = [&]() { return mx::full(shape, 3.3f); };
+  TIME(full_fp32);
+  auto zeros_fp32 = [&]() { return mx::zeros(shape, mx::float32); };
+  TIME(zeros_fp32);
+  auto ones_fp32 = [&]() { return mx::ones(shape, mx::float32); };
+  TIME(ones_fp32);
+
+  auto arange_fp32 = [&]() { return mx::arange(0.0, 10.0, 1e-4); };
+  TIME(arange_fp32);
+}
+
+void time_type_conversions() {
+  int M = 2000;
+  int N = 500;
+  auto shape = {M, N};
+  auto device = mx::default_device();
+
+  auto a = mx::zeros(shape, mx::float32);
+  mx::eval(a);
+  TIMEM("mx::float32 to mx::int32", mx::astype, a, mx::int32, device);
+  TIMEM("mx::float32 to mx::uint32", mx::astype, a, mx::uint32, device);
+
+  a = mx::zeros(shape, mx::int32);
+  mx::eval(a);
+  TIMEM("mx::int32 to mx::float32", mx::astype, a, mx::float32, device);
+
+  a = mx::zeros(shape, mx::bool_);
+  mx::eval(a);
+  TIMEM("bool to mx::float32", mx::astype, a, mx::float32, device);
+  TIMEM("bool to mx::int32", mx::astype, a, mx::int32, device);
+  TIMEM("bool to mx::uint32", mx::astype, a, mx::uint32, device);
+}
+
+void time_random_generation() {
+  int M = 2000;
+  int N = 500;
+
+  auto uniform = [&]() { return mx::random::uniform({M, N}, mx::float32); };
+  TIME(uniform);
+  auto normal = [&]() { return mx::random::normal({M, N}, mx::float32); };
+  TIME(normal);
+}
+
+void time_unary_ops() {
+  int M = 2000;
+  int N = 500;
+  auto device = mx::default_device();
+
+  auto a = mx::random::normal({M, N});
+  mx::eval(a);
+  TIME(mlx::core::abs, a, device);
+  TIME(mx::negative, a, device);
+  TIME(mx::sign, a, device);
+  TIME(mx::square, a, device);
+  TIME(mlx::core::sqrt, a, device);
+  TIME(mx::rsqrt, a, device);
+  TIME(mlx::core::exp, a, device);
+
+  a = mx::random::uniform({M, N});
+  TIME(mlx::core::log, a, device);
+}
+
+void time_binary_ops() {
+  int M = 1000, N = 100, K = 10;
+  auto condition = mx::random::randint(0, 2, {M, N, K});
+  auto a = mx::random::uniform({M, N, K});
+  auto b = mx::random::uniform({M, N, K});
+  auto device = mx::default_device();
+  mx::eval(a, b);
+
+  TIME(mx::add, a, b, device);
+  TIME(mx::subtract, a, b, device);
+  TIME(mx::multiply, a, b, device);
+  TIME(mx::divide, a, b, device);
+  TIME(mx::maximum, a, b, device);
+  TIME(mx::minimum, a, b, device);
+  TIME(mx::where, condition, a, b, device);
+
+  condition = mx::array({true});
+  b = mx::random::uniform({1});
+  mx::eval(b);
+  TIMEM("scalar", mx::add, a, b, device);
+  TIMEM("vector-scalar", mx::subtract, a, b, device);
+  TIMEM("scalar-vector", mx::subtract, b, a, device);
+  TIMEM("scalar", mx::multiply, a, b, device);
+  TIMEM("vector-scalar", mx::divide, a, b, device);
+  TIMEM("scalar-vector", mx::divide, b, a, device);
+  TIMEM("scalar-vector", mx::where, condition, a, b, device);
+
+  condition = mx::broadcast_to(mx::array({true}), {1000, 100});
+  a = mx::broadcast_to(mx::random::uniform({1}), {1000, 100});
+  b = mx::broadcast_to(mx::random::uniform({1}), {1000, 100});
+  mx::eval(a, b);
+  TIMEM("scalar-scalar broadcast", mx::add, a, b, device);
+  TIMEM("scalar-scalar broadcast", mx::subtract, a, b, device);
+  TIMEM("scalar-scalar broadcast", mx::multiply, a, b, device);
+  TIMEM("scalar-scalar broadcast", mx::divide, a, b, device);
+  TIMEM("scalar-scalar broadcast", mx::where, condition, a, b, device);
+}
+
+void time_strided_ops() {
+  int M = 50, N = 50, O = 50, P = 50;
+  auto a = mx::random::uniform({M, N, O, P});
+  auto b = mx::random::uniform({M, N, O, P});
+  auto device = mx::default_device();
+  mx::eval(a, b);
+  TIMEM("non-strided", mx::add, a, b, device);
+  a = mx::transpose(a, {1, 0, 2, 3});
+  b = mx::transpose(b, {3, 2, 0, 1});
+  mx::eval(a, b);
+  TIMEM("strided", mx::add, a, b, device);
+}
+
+void time_comparisons() {
+  int M = 1000, N = 100, K = 10;
+  auto a = mx::random::uniform({M, N, K});
+  auto b = mx::random::uniform({M, N, K});
+  auto device = mx::default_device();
+  mx::eval(a, b);
+  TIME(mx::equal, a, b, device);
+  TIME(mx::greater, a, b, device);
+  TIME(mx::greater_equal, a, b, device);
+  TIME(mx::less, a, b, device);
+  TIME(mx::less_equal, a, b, device);
+}
+
+void time_matvec() {
+  int M = 2000, N = 200;
+  auto a = mx::random::uniform({M, N});
+  auto b = mx::random::uniform({N});
+  auto c = mx::random::uniform({M});
+  mx::eval(a, b, c);
+  auto matvec = [&]() { return mx::matmul(a, b); };
+  TIME(matvec);
+
+  auto matvec_transpose = [&]() { return mx::matmul(mx::transpose(a), c); };
+  TIME(matvec_transpose);
+}
+
+void time_matmul() {
+  int M = 1000, N = 1000, K = 1000;
+  auto a = mx::random::uniform({M, K});
+  auto b = mx::random::uniform({K, N});
+  auto device = mx::default_device();
+  mx::eval(a, b);
+  TIME(mx::matmul, a, b, device);
+
+  auto transpose_matmul = [&]() { return mx::matmul(mx::transpose(a), b); };
+  TIME(transpose_matmul);
+}
+
+void time_reductions() {
+  auto a = mx::random::normal({10000, 1000});
+  mx::eval(a);
+  auto sum_all = [&a]() { return mx::sum(a, false); };
+  TIME(sum_all);
+
+  auto sum_along_0 = [&a]() { return mx::sum(a, 0, false); };
+  TIME(sum_along_0);
+
+  auto sum_along_1 = [&a]() { return mx::sum(a, 1, false); };
+  TIME(sum_along_1);
+
+  auto prod_all = [&a]() { return mx::prod(a, false); };
+  TIME(prod_all);
+
+  auto all_true = [&a]() { return mx::all(a, false); };
+  TIME(all_true);
+
+  auto all_along_0 = [&a]() { return mx::all(a, 0, false); };
+  TIME(all_along_0);
+
+  auto all_along_1 = [&a]() { return mx::all(a, 1, false); };
+  TIME(all_along_1);
+
+  auto any_true = [&a]() { return mx::any(a, false); };
+  TIME(any_true);
+
+  auto argmin_along_0 = [&a]() { return mx::argmin(a, 0, false); };
+  TIME(argmin_along_0);
+
+  auto argmin_along_1 = [&a]() { return mx::argmin(a, 1, false); };
+  TIME(argmin_along_1);
+
+  auto indices = mx::array({1});
+  auto updates = mx::reshape(mx::array({NAN}), {1, 1, 1});
+  std::vector<int> axes{0};
+  auto b = scatter(a, {indices}, updates, axes);
+  mx::eval(b);
+
+  auto max_along_0 = [&b]() { return mx::max(b, 0, false); };
+  TIME(max_along_0);
+  auto max_along_1 = [&b]() { return mx::max(b, 1, false); };
+  TIME(max_along_1);
+
+  auto min_along_0 = [&b]() { return mx::min(b, 0, false); };
+  TIME(min_along_0);
+  auto min_along_1 = [&b]() { return mx::min(b, 1, false); };
+  TIME(min_along_1);
+}
+
+void time_gather_scatter() {
+  auto a = mx::random::normal({1000, 768});
+  mx::eval(a);
+  auto indices = mx::random::randint(0, 1000, {256});
+  mx::eval(indices);
+
+  auto embedding_lookup = [&a, &indices]() { return mx::take(a, indices, 0); };
+  TIME(embedding_lookup);
+
+  indices = mx::random::randint(0, 768 * 1000, {256 * 768});
+  mx::eval(indices);
+
+  auto single_element_lookup = [&a, &indices]() {
+    return mx::take(a, indices);
+  };
+  TIME(single_element_lookup);
+
+  indices = mx::random::randint(0, 1000, {256});
+  auto updates = mx::random::normal({256, 1, 768});
+  mx::eval(indices, updates);
+
+  auto embedding_update = [&a, &indices, &updates]() {
+    return scatter(a, indices, updates, 0);
+  };
+  TIME(embedding_update);
+
+  auto embedding_add = [&a, &indices, &updates]() {
+    return scatter_add(a, indices, updates, 0);
+  };
+  TIME(embedding_add);
+
+  a = mx::reshape(a, {-1});
+  indices = mx::random::randint(0, 768 * 1000, {768 * 256});
+  updates = mx::random::normal({256 * 768, 1});
+  mx::eval(a, indices, updates);
+
+  auto single_element_update = [&a, &indices, &updates]() {
+    return scatter(a, indices, updates, 0);
+  };
+  TIME(single_element_update);
+
+  auto single_element_add = [&a, &indices, &updates]() {
+    return scatter_add(a, indices, updates, 0);
+  };
+  TIME(single_element_add);
+}
+
+void time_divmod() {
+  auto a = mx::random::normal({1000});
+  auto b = mx::random::normal({1000});
+  mx::eval({a, b});
+
+  auto divmod_fused = [&a, &b]() { return mx::divmod(a, b); };
+  TIME(divmod_fused);
+
+  auto divmod_separate = [&a, &b]() {
+    return std::vector<mx::array>{mx::floor_divide(a, b), mx::remainder(a, b)};
+  };
+  TIME(divmod_separate);
+}
+
+int main() {
+  std::cout << "Benchmarks for " << mx::default_device() << std::endl;
+  time_creation_ops();
+  time_type_conversions();
+  time_unary_ops();
+  time_binary_ops();
+  time_strided_ops();
+  time_random_generation();
+  time_comparisons();
+  time_matvec();
+  time_matmul();
+  time_reductions();
+  time_gather_scatter();
+  time_divmod();
+}