Function NB_MODULE

python/src/mlx.cpp:28–53

void init_print(nb::module_&);

NB_MODULE(core, m) {
  m.doc() = "mlx: A framework for machine learning on Apple silicon.";

  auto reprlib_fix = nb::module_::import_("mlx._reprlib_fix");
  nb::set_leak_warnings(false);

  init_mlx_func(m);
  init_device(m);
  init_stream(m);
  init_array(m);
  init_metal(m);
  init_cuda(m);
  init_memory(m);
  init_ops(m);
  init_transforms(m);
  init_random(m);
  init_fft(m);
  init_linalg(m);
  init_constants(m);
  init_fast(m);
  init_distributed(m);
  init_export(m);
  init_print(m);

  m.attr("__version__") = mx::version();
}