Use uv

.github/actions/build-macos/action.yml

@@ -4,14 +4,12 @@ description: 'Build and test MLX on macOS'
 runs:
   using: "composite"
   steps:
-    - name: Install dependencies
+    - name: Build
       env:
         DEBUG: 1
         CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
       shell: bash -l {0}
       run: |
-        pip install --upgrade pip
-        pip install cmake setuptools typing_extensions
         pip install -e . -v
 
     - name: Install tests dependencies

.github/actions/setup-macos/action.yml

@@ -18,7 +18,18 @@ runs:
       shell: bash
       run: xcodebuild -showComponent MetalToolchain
 
-    - uses: conda-incubator/setup-miniconda@v3
-      with:
-        miniconda-version: "latest"
-        python-version: ${{ inputs.python-version }}
+    - name: Install Python
+      shell: sh
+      run: |
+        curl -LsSf https://astral.sh/uv/install.sh | sh
+        $HOME/.local/bin/uv venv --python ${{ inputs.python-version }}
+        source .venv/bin/activate
+        echo PATH=$PATH >> $GITHUB_ENV
+        # Search python packages in .venv
+        echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV
+
+    - name: Install build dependencies
+      shell: sh
+      run: |
+        pip install --upgrade pip
+        pip install cmake setuptools typing_extensions

docs/src/usage/distributed.rst

@@ -584,7 +584,7 @@ the process.
 **MLX_JACCL_COORDINATOR** should contain the IP and port that rank 0 can listen
 to all the other ranks connect to in order to establish the RDMA connections.
 
-**MLX_JACCL_DEVICES** should contain the path to a json file that contains the
+**MLX_IBV_DEVICES** should contain the path to a json file that contains the
 ibverbs device names that connect each node to each other node, something like
 the following:
 

mlx/backend/cpu/unary_ops.h

@@ -154,12 +154,24 @@ struct ToFP8 {
 struct FromFP8 {
   template <int N>
   Simd<float, N> operator()(Simd<uint8_t, N> x) {
-    auto v = Simd<uint16_t, N>(x & 127) << 7;
-    auto converted = *(Simd<float16_t, N>*)(&v);
-    converted = converted * 256.0;
-    auto sign = Simd<bool, N>(x & 128);
-    Simd<float, N> out = select(sign, -converted, converted);
-    return out;
+    auto w = Simd<uint32_t, N>(x) << 24;
+    auto sign = w & 0x80000000;
+    auto nonsign = w & 0x7FFFFFFF;
+
+    auto renorm_shift = clz(nonsign);
+    renorm_shift = simd::select(
+        renorm_shift > Simd<uint32_t, N>{4},
+        renorm_shift - Simd<uint32_t, N>{4},
+        Simd<uint32_t, N>{0});
+
+    Simd<int32_t, N> inf_nan_mask =
+        (Simd<int32_t, N>(nonsign + 0x01000000) >> 8) & 0x7F800000;
+    auto zero_mask = Simd<int32_t, N>(nonsign - 1) >> 31;
+    auto result = sign |
+        ((((nonsign << renorm_shift >> 4) + ((0x78 - renorm_shift) << 23)) |
+          inf_nan_mask) &
+         ~zero_mask);
+    return fp32_from_bits(result);
   }
   float operator()(uint8_t x) {
     return (*this)(Simd<uint8_t, 1>(x)).value;

mlx/backend/cuda/quantized/affine_quantize.cu

@@ -23,7 +23,8 @@ affine_quantize(const T* w, uint8_t* out, T* scales, T* biases, size_t size) {
   auto tidx = block_idx.x * block_size.x + idx_in_block.x;
   auto tidy = block_idx.y * block_size.y + idx_in_block.y;
 
-  auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x;
+  auto grid_dim_x =
+      cg::this_grid().dim_blocks().x * cg::this_grid().block_index().x;
   constexpr float eps = 1e-7;
   constexpr int simd_size = WARP_SIZE;
   constexpr float n_bins = (1 << bits) - 1;
@@ -140,7 +141,8 @@ __global__ void affine_dequantize(
   auto tidx = block_idx.x * block_size.x + idx_in_block.x;
   auto tidy = block_idx.y * block_size.y + idx_in_block.y;
 
-  auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x;
+  auto grid_dim_x =
+      cg::this_grid().dim_blocks().x * cg::this_grid().block_index().x;
 
   constexpr int pack_factor = get_pack_factor<bits, 8>();
   constexpr int bytes_per_pack = get_bytes_per_pack<bits>();

mlx/backend/cuda/quantized/qqmm.cpp

@@ -46,7 +46,7 @@ inline array ensure_row_contiguous_matrix(
   return x_copy;
 }
 
-array pad_and_swizzle_scales(
+array pad_and_repack_scales(
     const array& scale,
     cu::CommandEncoder& encoder,
     const Stream& s) {
@@ -64,12 +64,14 @@ array pad_and_swizzle_scales(
       cu::malloc_async(pad_outer * pad_inner, encoder),
       Shape{pad_outer, pad_inner},
       scale.dtype());
-  swizzle_scales(scale, scale_tiled, encoder, s);
+  repack_scales(scale, scale_tiled, encoder, s);
 
   encoder.add_temporary(scale_tiled);
   return scale_tiled;
 }
+} // namespace
 
+namespace {
 void qqmm_impl(
     cu::CommandEncoder& encoder,
     int M,
@@ -174,8 +176,8 @@ void QQMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
   int K = K_packed * (32 / bits_);
 
   // Repack scales from linear to tiled layout for tensor cores
-  array scale_x = pad_and_swizzle_scales(scale_x_pre, encoder, s);
-  array scale_w = pad_and_swizzle_scales(scale_w_pre, encoder, s);
+  array scale_x = pad_and_repack_scales(scale_x_pre, encoder, s);
+  array scale_w = pad_and_repack_scales(scale_w_pre, encoder, s);
 
   bool x_transposed = false;
   bool w_transposed = true; // always transposed

mlx/backend/cuda/quantized/qqmm_utils.cu

@@ -10,11 +10,6 @@ namespace mlx::core {
 
 namespace cg = cooperative_groups;
 
-constexpr int TILE_ROWS = 128;
-constexpr int TILE_COLS = 4;
-constexpr int TILES_PER_LANE = 1;
-constexpr int LANES_PER_BLOCK = 32;
-
 // To pass scales to tensor cores, they need to be repacked into a tiled layout
 // https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout
 // Tiled layout for scale factors is very well described in CUTLASS
@@ -48,171 +43,118 @@ constexpr int LANES_PER_BLOCK = 32;
 //          [252, 253, 254, 255],
 //          [380, 381, 382, 383],
 //          [508, 509, 510, 511]]]]],
+__device__ size_t
+scale_tiled_offset(size_t scale_index, size_t num_rows, size_t num_scale_cols) {
+  // Compute the tiled layout offset for scale factors used in tensor cores
+  // This function maps from a linear scale index to the tiled layout expected
+  // by tensor cores (and cublaslt).
+  //
+  // Input: linear scale index (e.g., for a matrix M x K with group_size,
+  //        scale_index ranges from 0 to (M * K/group_size - 1))
+  //
+  // The tiled layout organizes scales into tiles of 128 rows x 4 columns,
+  // where each tile is subdivided into 4 sub-blocks of 32 rows x 4 columns.
+  size_t row = scale_index / num_scale_cols;
+  size_t col = scale_index % num_scale_cols;
 
-inline std::tuple<dim3, dim3> get_swizzle_launch_args(
-    size_t M_swizzled,
-    size_t K_swizzled) {
-  constexpr int tiles_per_block = LANES_PER_BLOCK * TILES_PER_LANE;
-  constexpr int warps_per_block = TILE_ROWS / 4; // 128 / 4 = 32
+  constexpr size_t rows_per_tile = 128;
+  constexpr size_t rows_per_sub_block = 32;
+  constexpr size_t cols_per_sub_block = 4;
+  constexpr size_t sub_blocks_per_tile = 4; // Vertically stacked
 
-  const int num_tiles_k = K_swizzled / TILE_COLS;
-  const int num_tiles_m = M_swizzled / TILE_ROWS;
+  // Decompose row position
+  size_t tile_row = row / rows_per_tile; // Which tile row
+  size_t row_in_tile = row % rows_per_tile; // Row within tile
+  size_t sub_block_row =
+      row_in_tile / rows_per_sub_block; // Sub-block within tile
+  size_t row_in_sub_block =
+      row_in_tile % rows_per_sub_block; // Row in sub-block
 
-  dim3 grid;
-  grid.x = cuda::ceil_div(num_tiles_k, tiles_per_block);
-  grid.y = num_tiles_m;
-  grid.z = 1;
-  // Block is always (32, 32) = 1024 threads
-  dim3 block(LANES_PER_BLOCK, warps_per_block, 1);
+  // Decompose column position
+  size_t col_tile = col / cols_per_sub_block; // Which column tile
+  size_t col_in_sub_block = col % cols_per_sub_block; // Column within sub-block
 
-  return std::make_tuple(grid, block);
+  // Compute tile index and offset within tile
+  size_t num_col_tiles = cuda::ceil_div(num_scale_cols, cols_per_sub_block);
+  size_t tile_idx = tile_row * num_col_tiles + col_tile;
+
+  size_t offset_in_tile =
+      (row_in_sub_block * sub_blocks_per_tile * cols_per_sub_block) +
+      (sub_block_row * cols_per_sub_block) + col_in_sub_block;
+
+  constexpr size_t tile_size = rows_per_tile * cols_per_sub_block;
+  return tile_idx * tile_size + offset_in_tile;
 }
 
 namespace cu {
 
-__global__ void swizzle_scales(
+__global__ void repack_scales(
     const uint8_t* scales_linear,
-    uint8_t* scales_swizzled,
-    const size_t M,
-    const size_t K,
-    const size_t M_swizzled,
-    const size_t K_swizzled) {
-  constexpr int tile_size = TILE_ROWS * TILE_COLS;
-  constexpr int num_tile_rows_per_thread = 4;
-  constexpr int max_tiles_per_block = LANES_PER_BLOCK * TILES_PER_LANE;
-
-  constexpr int tile_stride = tile_size / 16; // 32 int4s per tile
-
-  // Each thread loads 16 scales from 4 rows (stride 32) and packs them into
-  // int4. For example: thread (0, 0) loads scales at rows 0,32,64,96 of tile 0,
-  // thread (1, 0) loads rows 0,32,64,96 of of tile 1, etc.
-  // The store is strided within a warp (stride 32 int4s), so we first
-  // write to shared memory, then do a coalesced store from shared to global
+    uint8_t* scales_tiled,
+    size_t input_rows,
+    size_t input_cols,
+    size_t output_rows,
+    size_t output_cols) {
   auto block_size = cg::this_thread_block().dim_threads();
   auto block_idx = cg::this_thread_block().group_index();
   auto idx_in_block = cg::this_thread_block().thread_index();
 
-  auto tidx = idx_in_block.x;
-  auto tidy = idx_in_block.y;
-  auto linear_tid = tidy * block_size.x + tidx;
+  auto tidx = block_idx.x * block_size.x + idx_in_block.x;
+  auto tidy = block_idx.y * block_size.y + idx_in_block.y;
 
-  const int bid_x = block_idx.x;
-  const int bid_y = block_idx.y;
+  auto grid_dim_x =
+      cg::this_grid().dim_blocks().x * cg::this_grid().block_index().x;
 
-  const int K_int = K_swizzled / 4;
+  size_t output_index = tidx + grid_dim_x * size_t(tidy);
+  size_t output_size = output_rows * output_cols;
 
-  const size_t output_offset = static_cast<size_t>(bid_y) * TILE_ROWS * K_int +
-      static_cast<size_t>(bid_x) * max_tiles_per_block * tile_size / 4;
-  int* output_block = reinterpret_cast<int*>(scales_swizzled) + output_offset;
-
-  const int grid_dim_x = cg::this_grid().dim_blocks().x;
-  const int grid_dim_y = cg::this_grid().dim_blocks().y;
-
-  int remaining = K_int - bid_x * max_tiles_per_block;
-  int tiles_in_block = min(remaining, max_tiles_per_block);
-  bool valid_tile = tidx * TILES_PER_LANE < tiles_in_block;
-
-  __shared__ int4 strided_scales_thread[max_tiles_per_block * tile_stride];
-
-  // Initialize to zero for padding
-  int thread_tile_rows[num_tile_rows_per_thread] = {0};
-
-  if (valid_tile) {
-    const size_t col_base =
-        static_cast<size_t>(bid_x) * max_tiles_per_block * TILE_COLS +
-        tidx * TILE_COLS;
-
-    const bool aligned_k = (K % 4 == 0);
-
-    if (aligned_k) {
-      // fast path: K is aligned, use vectorized loads with stride K/4
-      const int K_stride = K / 4;
-      const size_t block_offset =
-          static_cast<size_t>(bid_y) * TILE_ROWS * K_stride +
-          static_cast<size_t>(bid_x) * max_tiles_per_block;
-      const int* input_block =
-          reinterpret_cast<const int*>(scales_linear) + block_offset;
-// load
-#pragma unroll
-      for (int i = 0; i < num_tile_rows_per_thread; i++) {
-        const size_t row =
-            static_cast<size_t>(bid_y) * TILE_ROWS + i * block_size.x + tidy;
-        const int thread_offset =
-            (i * block_size.x + tidy) * K_stride + tidx * TILES_PER_LANE;
-        if (row < M && col_base + TILE_COLS <= K) {
-          thread_tile_rows[i] = __ldg(input_block + thread_offset);
-        } else if (row < M) {
-// partial tile at K boundary: load byte-by-byte
-#pragma unroll
-          for (int c = 0; c < TILE_COLS; c++) {
-            if (col_base + c < K) {
-              reinterpret_cast<uint8_t*>(&thread_tile_rows[i])[c] =
-                  scales_linear[row * K + col_base + c];
-            }
-          }
-        }
-      }
-    } else {
-#pragma unroll
-      for (int i = 0; i < num_tile_rows_per_thread; i++) {
-        const size_t row =
-            static_cast<size_t>(bid_y) * TILE_ROWS + i * block_size.x + tidy;
-        if (row < M) {
-          const size_t row_start = row * K;
-#pragma unroll
-          for (int c = 0; c < TILE_COLS; c++) {
-            if (col_base + c < K) {
-              reinterpret_cast<uint8_t*>(&thread_tile_rows[i])[c] =
-                  scales_linear[row_start + col_base + c];
-            }
-          }
-        }
-      }
-    }
-    // store to shared with XOR swizzle to avoid bank conflicts
-    int base_idx = tidx * tile_stride + tidy;
-    int xor_bits = (tidy >> 3) & 0x3;
-    int swizzled_idx = base_idx ^ xor_bits;
-    strided_scales_thread[swizzled_idx] =
-        *reinterpret_cast<int4*>(thread_tile_rows);
+  if (output_index >= output_size) {
+    return;
   }
 
-  cg::thread_block block = cg::this_thread_block();
-  cg::sync(block);
+  size_t tiled_offset =
+      scale_tiled_offset(output_index, output_rows, output_cols);
 
-  const int total_int4s = tiles_in_block * tile_stride;
-#pragma unroll
-  for (int i = linear_tid; i < total_int4s; i += block_size.x * block_size.y) {
-    int tile_idx = i / tile_stride;
-    int row_idx = i % tile_stride;
-    int base_idx = tile_idx * tile_stride + row_idx;
-    int xor_bits = (row_idx >> 3) & 0x3;
-    int swizzled_idx = base_idx ^ xor_bits;
-    reinterpret_cast<int4*>(output_block)[i] =
-        strided_scales_thread[swizzled_idx];
+  size_t row = output_index / output_cols;
+  size_t col = output_index % output_cols;
+
+  // Probably this can be done better with 2 separated paths for valid and
+  // padding
+  if (row < input_rows && col < input_cols) {
+    size_t input_index = row * input_cols + col;
+    scales_tiled[tiled_offset] = scales_linear[input_index];
+  } else {
+    // Zero-fill padding region
+    scales_tiled[tiled_offset] = 0;
   }
 }
+
 } // namespace cu
 
-void swizzle_scales(
+void repack_scales(
     const array& scales,
     array& scales_tiled,
     cu::CommandEncoder& enc,
     const Stream& s) {
   enc.set_input_array(scales);
   enc.set_output_array(scales_tiled);
+
   // Note: scales_tiled is padded to full tiles so if num_rows or num_cols
-  // are not multiples of tile sizes
+  // are not multiples of tile sizes, the extra space is filled with zeros
+
   size_t input_rows = scales.shape(-2);
   size_t input_cols = scales.shape(-1);
 
   size_t output_rows = scales_tiled.shape(-2);
   size_t output_cols = scales_tiled.shape(-1);
+  size_t output_size = output_rows * output_cols;
+
+  bool large = output_size > UINT_MAX;
+  auto [num_blocks, block_dims] = get_launch_args(
+      output_size, scales_tiled.shape(), scales_tiled.strides(), large);
 
-  auto [num_blocks, block_dims] =
-      get_swizzle_launch_args(output_rows, output_cols);
   enc.add_kernel_node(
-      cu::swizzle_scales,
+      cu::repack_scales,
       num_blocks,
       block_dims,
       0,

mlx/backend/cuda/quantized/qqmm_utils.h

@@ -21,7 +21,7 @@ inline std::pair<int, int> get_padded_scale_dims(int num_rows, int num_cols) {
   return {padded_rows, padded_cols};
 }
 
-void swizzle_scales(
+void repack_scales(
     const array& scales,
     array& scales_tiled,
     cu::CommandEncoder& enc,

mlx/backend/metal/device.h

@@ -266,9 +266,6 @@ Device& device(mlx::core::Device);
 std::unique_ptr<void, std::function<void(void*)>> new_scoped_memory_pool();
 
 inline bool is_nax_available() {
-#ifdef MLX_METAL_NO_NAX
-  return false;
-#else
   auto _check_nax = []() {
     bool can_use_nax = false;
     if (__builtin_available(
@@ -281,7 +278,6 @@ inline bool is_nax_available() {
   };
   static bool is_nax_available_ = _check_nax();
   return is_nax_available_;
-#endif
 }
 
 } // namespace mlx::core::metal

mlx/backend/metal/kernels/CMakeLists.txt

@@ -164,8 +164,6 @@ if(NOT MLX_METAL_JIT)
     build_kernel(steel/attn/kernels/steel_attention_nax
                  ${STEEL_NAX_ATTN_HEADERS})
 
-  else()
-    target_compile_definitions(mlx PRIVATE MLX_METAL_NO_NAX)
   endif()
 
 endif()

mlx/distributed/jaccl/jaccl.cpp

@@ -1159,7 +1159,7 @@ bool is_available() {
 }
 
 std::shared_ptr<GroupImpl> init(bool strict /* = false */) {
-  const char* dev_file = std::getenv("MLX_JACCL_DEVICES");
+  const char* dev_file = std::getenv("MLX_IBV_DEVICES");
   const char* coordinator = std::getenv("MLX_JACCL_COORDINATOR");
   const char* rank_str = std::getenv("MLX_RANK");
 
@@ -1167,9 +1167,9 @@ std::shared_ptr<GroupImpl> init(bool strict /* = false */) {
     if (strict) {
       std::ostringstream msg;
       msg << "[jaccl] You need to provide via environment variables a rank (MLX_RANK), "
-          << "a device file (MLX_JACCL_DEVICES) and a coordinator ip/port (MLX_JACCL_COORDINATOR) "
+          << "a device file (MLX_IBV_DEVICES) and a coordinator ip/port (MLX_JACCL_COORDINATOR) "
           << "but provided MLX_RANK=\"" << ((rank_str) ? rank_str : "")
-          << "\", MLX_JACCL_DEVICES=\"" << ((dev_file) ? dev_file : "")
+          << "\", MLX_IBV_DEVICES=\"" << ((dev_file) ? dev_file : "")
           << "\" and MLX_JACCL_COORDINATOR=\""
           << ((coordinator) ? coordinator : "");
       throw std::runtime_error(msg.str());

mlx/io/safetensors.cpp

@@ -95,6 +95,7 @@ Dtype dtype_from_safetensor_str(std::string_view str) {
   } else if (str == ST_C64) {
     return complex64;
   } else if (str == ST_F8_E4M3) {
+    // We convert this manually later
     return uint8;
   } else {
     throw std::runtime_error(
@@ -147,14 +148,16 @@ SafetensorsLoad load_safetensors(
     const Shape& shape = item.value().at("shape");
     const std::vector<size_t>& data_offsets = item.value().at("data_offsets");
     Dtype type = dtype_from_safetensor_str(dtype);
-    res.insert(
-        {item.key(),
-         array(
-             shape,
-             type,
-             std::make_shared<Load>(
-                 stream, in_stream, offset + data_offsets.at(0), false),
-             std::vector<array>{})});
+    auto loaded_array = array(
+        shape,
+        type,
+        std::make_shared<Load>(
+            stream, in_stream, offset + data_offsets.at(0), false),
+        std::vector<array>{});
+    if (dtype == ST_F8_E4M3) {
+      loaded_array = from_fp8(loaded_array, bfloat16, s);
+    }
+    res.insert({item.key(), loaded_array});
   }
   return {res, metadata_map};
 }

mlx/ops.cpp

@@ -4335,7 +4335,7 @@ std::pair<int, int> extract_qqmm_dims(
 }
 
 array qqmm(
-    array in_x,
+    array x,
     array w,
     std::optional<array> scales_w,
     std::optional<int> group_size_ /* = std::nullopt */,
@@ -4360,16 +4360,6 @@ array qqmm(
   // 2. w is not quantized, scales is not provided
   auto [group_size, bits] =
       quantization_params_from_mode(qmode, group_size_, bits_);
-
-  // Allow gemv
-  auto x = in_x;
-  if (x.ndim() == 1) {
-    // Insert a singleton dim in the beginning
-    x = expand_dims(x, 0, s);
-  } else if (w.ndim() == 2 && x.ndim() > 2) {
-    x = flatten(x, 0, -2, s);
-  }
-
   // validate inputs
   validate_qqmm_inputs(x, w, scales_w, group_size, bits);
   // validate and extract shapes
@@ -4384,19 +4374,11 @@ array qqmm(
   }
   auto out_shape = inputs[0].shape();
   out_shape.back() = w_outer_dims;
-  auto out = array(
+  return array(
       std::move(out_shape),
       x.dtype(), // output dtype is the same as x dtype
       std::make_shared<QQMatmul>(stream, group_size, bits, qmode),
       std::move(inputs));
-  if (in_x.ndim() > 2) {
-    auto orig_shape = in_x.shape();
-    orig_shape.pop_back();
-    out = unflatten(out, 0, std::move(orig_shape), s);
-  } else if (in_x.ndim() == 1) {
-    out = squeeze(out, 0, s);
-  }
-  return out;
 }
 
 array pack_and_quantize(

mlx/primitives.cpp

@@ -3717,7 +3717,7 @@ std::pair<std::vector<array>, std::vector<int>> RandomBits::vmap(
 
 bool RandomBits::is_equivalent(const Primitive& other) const {
   const RandomBits& r_other = static_cast<const RandomBits&>(other);
-  return shape_ == r_other.shape_ && width_ == r_other.width_;
+  return shape_ == r_other.shape_;
 }
 
 std::vector<array> Real::vjp(

python/mlx/_distributed_utils/launch.py

@@ -376,7 +376,7 @@ def launch_jaccl(parser, hosts, args, command):
     env = args.env
     cwd = args.cwd
     env.append(f"MLX_JACCL_COORDINATOR={coordinator}:{args.starting_port}")
-    files = {"MLX_JACCL_DEVICES": json.dumps([h.rdma for h in hosts])}
+    files = {"MLX_IBV_DEVICES": json.dumps([h.rdma for h in hosts])}
 
     log(args.verbose, "Running", shlex.join(command))
 

python/src/ops.cpp

@@ -1668,25 +1668,6 @@ void init_ops(nb::module_& m) {
         Returns:
             array: The array of zeros with the specified shape.
       )pbdoc");
-  m.def(
-      "asarray",
-      [](const ArrayInitType& a, std::optional<mx::Dtype> dtype) {
-        return create_array(a, dtype);
-      },
-      nb::arg(),
-      "dtype"_a = nb::none(),
-      nb::sig("def asarray(a: Union[scalar, array, Sequence], dtype: "
-              "Optional[Dtype] = None) -> array"),
-      R"pbdoc(
-        Convert the input to an array.
-
-        Args:
-            a: Input data.
-            dtype (Dtype, optional): The desired data-type for the array.
-
-        Returns:
-            array: An array interpretation of the input.
-      )pbdoc");
   m.def(
       "zeros_like",
       &mx::zeros_like,
@@ -4300,7 +4281,7 @@ void init_ops(nb::module_& m) {
             ======  ======================   ==========================  =============  =====
             affine  32, 64\ :sup:`*`, 128    2, 3, 4\ :sup:`*`, 5, 6, 8  same as input  yes
             mxfp4   32\ :sup:`*`             4\ :sup:`*`                 e8m0           no
-            mxfp8   32\ :sup:`*`             8\ :sup:`*`                 e8m0           no
+            mxfp8   32\ :sup:`*`             4\ :sup:`*`                 e8m0           no
             nvfp4   16\ :sup:`*`             4\ :sup:`*`                 e4m3           no
             ======  ======================   ==========================  =============  =====
 
@@ -4335,7 +4316,7 @@ void init_ops(nb::module_& m) {
           size must be 16. The elements are quantized to 4-bit or 8-bit
           precision floating-point values: E2M1 for ``"fp4"`` and E4M3 for
           ``"fp8"``. There is a shared 8-bit scale per group. The ``"mx"``
-          modes use an E8M0 scale and the ``"nv"`` mode uses an E4M3 scale.
+          modes us an E8M0 scale and the ``"nv"`` mode uses an E4M3 scale.
           Unlike ``affine`` quantization, these modes does not have a bias
           value.
 
@@ -5479,10 +5460,10 @@ void init_ops(nb::module_& m) {
         If ``w`` is expected to receive gradients, it must be provided in
         non-quantized form.
 
-        If ``x`` and `w`` are not quantized, their data types must be ``float32``,
+        If ``x`` and `w`` are not quantized, their data types must be ``float32``, 
         ``float16``, or ``bfloat16``.
         If ``w`` is quantized, it must be packed in unsigned integers.
-
+        
       Args:
         x (array): Input array.
         w (array): Weight matrix. If quantized, it is packed in unsigned integers.
@@ -5502,40 +5483,4 @@ void init_ops(nb::module_& m) {
         array: The result of the multiplication of quantized ``x`` with quantized ``w``.
         needed).
   )pbdoc");
-  m.def(
-      "from_fp8",
-      &mx::from_fp8,
-      nb::arg(),
-      "dtype"_a = mx::bfloat16,
-      nb::kw_only(),
-      "stream"_a = nb::none(),
-      nb::sig(
-          "def from_fp8(x: array, dtype: Dtype = bfloat16, *, stream: Union[None, Stream, Device] = None) -> array"),
-      R"pbdoc(
-      Convert the array from fp8 (e4m3) to another floating-point type.
-
-      Args:
-        x (array): The input fp8 array with type ``uint8``.
-        dtype (Dtype): The data type to convert to. Default: ``bfloat16``.
-
-      Returns:
-        array: The array converted from fp8.
-  )pbdoc");
-  m.def(
-      "to_fp8",
-      &mx::to_fp8,
-      nb::arg(),
-      nb::kw_only(),
-      "stream"_a = nb::none(),
-      nb::sig(
-          "def to_fp8(x: array, *, stream: Union[None, Stream, Device] = None) -> array"),
-      R"pbdoc(
-      Convert the array to fp8 (e4m3) from another floating-point type.
-
-      Args:
-        x (array): The input array.
-
-      Returns:
-        array: The array converted to fp8 with type ``uint8``.
-  )pbdoc");
 }

python/tests/cuda_skip.py

@@ -1,4 +1,5 @@
 cuda_skip = {
+    "TestLoad.test_load_f8_e4m3",
     "TestLayers.test_quantized_embedding",
     # Block masked matmul NYI
     "TestBlas.test_block_masked_matmul",

python/tests/test_array.py

@@ -1973,58 +1973,6 @@ class TestArray(mlx_tests.MLXTestCase):
         self.assertTrue(hasattr(api, "array"))
         self.assertTrue(hasattr(api, "add"))
 
-    def test_array_namespace_asarray(self):
-        xp = mx.array(1.0).__array_namespace__()
-        self.assertTrue(hasattr(xp, "asarray"))
-
-        arr = xp.asarray([1, 2, 3])
-        self.assertEqual(arr.tolist(), [1, 2, 3])
-
-        arr_f32 = xp.asarray([1, 2, 3], dtype=mx.float32)
-        self.assertEqual(arr_f32.dtype, mx.float32)
-
-        existing = mx.array([4, 5, 6])
-        arr_pass = xp.asarray(existing)
-        self.assertEqual(arr_pass.tolist(), [4, 5, 6])
-
-    def test_asarray(self):
-        # List inputs
-        self.assertEqual(mx.asarray([1, 2, 3]).tolist(), [1, 2, 3])
-        self.assertEqual(mx.asarray([[1, 2], [3, 4]]).tolist(), [[1, 2], [3, 4]])
-
-        # Tuple inputs
-        self.assertEqual(mx.asarray((1, 2, 3)).tolist(), [1, 2, 3])
-        self.assertEqual(mx.asarray(((1, 2), (3, 4))).tolist(), [[1, 2], [3, 4]])
-
-        # Mixed nesting
-        self.assertEqual(mx.asarray([(1, 2), (3, 4)]).tolist(), [[1, 2], [3, 4]])
-        self.assertEqual(mx.asarray(([1, 2], [3, 4])).tolist(), [[1, 2], [3, 4]])
-
-        # Scalar inputs
-        self.assertEqual(mx.asarray(42).item(), 42)
-        self.assertEqual(mx.asarray(3.14).item(), 3.140000104904175)
-        self.assertEqual(mx.asarray(True).item(), True)
-        self.assertEqual(mx.asarray(1 + 2j).item(), (1 + 2j))
-
-        # MLX array inputs
-        arr = mx.array([1, 2, 3])
-        self.assertEqual(mx.asarray(arr).tolist(), [1, 2, 3])
-
-        arr_int = mx.array([1, 2, 3], dtype=mx.int32)
-        arr_float = mx.asarray(arr_int, dtype=mx.float32)
-        self.assertEqual(arr_float.dtype, mx.float32)
-        self.assertEqual(arr_float.tolist(), [1.0, 2.0, 3.0])
-
-        # NumPy array inputs
-        np_arr = np.array([1.0, 2.0, 3.0], dtype=np.float32)
-        mx_arr = mx.asarray(np_arr)
-        self.assertEqual(mx_arr.tolist(), [1.0, 2.0, 3.0])
-        self.assertEqual(mx_arr.dtype, mx.float32)
-
-        # dtype parameter
-        self.assertEqual(mx.asarray([1, 2, 3], dtype=mx.float32).dtype, mx.float32)
-        self.assertEqual(mx.asarray(42, dtype=mx.float16).dtype, mx.float16)
-
     def test_to_scalar(self):
         a = mx.array(1)
         self.assertEqual(int(a), 1)

python/tests/test_blas.py

@@ -297,7 +297,6 @@ class TestBlas(mlx_tests.MLXTestCase):
             self.assertTrue(np.allclose(out_mlx, out_npy, atol=1e-5))
 
     def test_matrix_vector(self):
-        mx.random.seed(0)
         for dtype in self.dtypes:
             with self.subTest(dtype=dtype):
                 np_dtype = getattr(np, dtype)

python/tests/test_load.py

@@ -168,8 +168,8 @@ class TestLoad(mlx_tests.MLXTestCase):
 
         expected = [
             0,
-            448,
-            -448,
+            mx.nan,
+            mx.nan,
             -0.875,
             0.4375,
             -0.005859,
@@ -179,12 +179,12 @@ class TestLoad(mlx_tests.MLXTestCase):
             -0.0039,
         ]
         expected = mx.array(expected, dtype=mx.bfloat16)
-        contents = b'H\x00\x00\x00\x00\x00\x00\x00{"tensor":{"dtype":"F8_E4M3","shape":[10],"data_offsets":[0,10]}}       \x00~\xfe\xb6.\x83\xba\xba\xbc\x82'
+        contents = b'H\x00\x00\x00\x00\x00\x00\x00{"tensor":{"dtype":"F8_E4M3","shape":[10],"data_offsets":[0,10]}}       \x00\x7f\xff\xb6.\x83\xba\xba\xbc\x82'
         with tempfile.NamedTemporaryFile(suffix=".safetensors") as f:
             f.write(contents)
             f.seek(0)
             out = mx.load(f)["tensor"]
-        self.assertTrue(mx.allclose(mx.from_fp8(out), expected))
+        self.assertTrue(mx.allclose(out[0], expected[0], equal_nan=True))
 
     def test_save_and_load_gguf_metadata_basic(self):
         if not os.path.isdir(self.test_dir):

python/tests/test_ops.py

@@ -3197,6 +3197,8 @@ class TestOps(mlx_tests.MLXTestCase):
             )
         )
 
+
+class TestBroadcast(mlx_tests.MLXTestCase):
     def test_broadcast_shapes(self):
         # Basic broadcasting
         self.assertEqual(mx.broadcast_shapes((1, 2, 3), (3,)), (1, 2, 3))
@@ -3241,13 +3243,6 @@ class TestOps(mlx_tests.MLXTestCase):
         self.assertTrue(mx.array_equal(mx.sort(x), expected, equal_nan=True))
         x = mx.array([3.0, mx.nan, 2.0, 0.0]) + 1j * mx.array([1.0] * 4)
 
-    def test_to_from_fp8(self):
-        vals = mx.array(
-            [448, 256, 192, 128, 96, 64, 48, 32, 24, 16, 12, 8, 6, 4, 3, 2, 0.015625]
-        )
-        self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(vals)), vals))
-        self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(-vals)), -vals))
-
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()

tests/compile_tests.cpp

@@ -803,16 +803,3 @@ TEST_CASE("test compile with no-ops") {
   auto out = compile(fun)({in})[0];
   CHECK_EQ(out.inputs()[0].id(), in.id());
 }
-
-TEST_CASE("test compile random bits") {
-  auto fun = [](const std::vector<array>& inputs) {
-    auto key = inputs[0];
-    auto a = random::bits({32, 32}, 4, key);
-    auto b = random::bits({32, 32}, 2, key);
-    return std::vector<array>{a + b};
-  };
-  auto in = random::key(0);
-  auto expected = fun({in})[0];
-  auto out = compile(fun)({in})[0];
-  CHECK(array_equal(out, expected).item<bool>());
-}