L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update pre-commit hooks and versions for clang-format, black, and isort (#3059)

Browse Source
This commit is contained in:
Nripesh Niketan
2026-01-26 14:57:04 +00:00
committed by GitHub
parent 5bd99dd5ec
commit b6aa03e5b8
37 changed files with 376 additions and 335 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.pre-commit-config.yaml
+3 -3
View File
@@ -6,17 +6,17 @@ repos:
# - id: end-of-file-fixer
# - id: trailing-whitespace
- repo: https://github.com/pre-commit/mirrors-clang-format
rev: v19.1.7
rev: v21.1.8
hooks:
- id: clang-format
# Using this mirror lets us use mypyc-compiled black, which is about 2x faster
- repo: https://github.com/psf/black-pre-commit-mirror
rev: 25.1.0
rev: 26.1.0
hooks:
- id: black
- repo: https://github.com/pycqa/isort
rev: 6.0.0
rev: 7.0.0
hooks:
- id: isort
args:
mlx/array.cpp
+14 -11
View File
@@ -21,11 +21,12 @@ array::array(
Dtype dtype,
std::shared_ptr<Primitive> primitive,
std::vector<array> inputs)
: array_desc_(std::make_shared<ArrayDesc>(
std::move(shape),
dtype,
std::move(primitive),
std::move(inputs))) {
: array_desc_(
std::make_shared<ArrayDesc>(
std::move(shape),
dtype,
std::move(primitive),
std::move(inputs))) {
if (has_primitive() && this->primitive().stream().device == Device::gpu) {
for (auto& in : this->inputs()) {
if (in.dtype() == float64) {
@@ -69,16 +70,18 @@ array array::unsafe_weak_copy(const array& other) {
}
array::array(std::initializer_list<float> data)
: array_desc_(std::make_shared<ArrayDesc>(
Shape{static_cast<ShapeElem>(data.size())},
float32)) {
: array_desc_(
std::make_shared<ArrayDesc>(
Shape{static_cast<ShapeElem>(data.size())},
float32)) {
init(data.begin());
}
array::array(std::initializer_list<int> data, Dtype dtype)
: array_desc_(std::make_shared<ArrayDesc>(
Shape{static_cast<ShapeElem>(data.size())},
dtype)) {
: array_desc_(
std::make_shared<ArrayDesc>(
Shape{static_cast<ShapeElem>(data.size())},
dtype)) {
init(data.begin());
}
mlx/array.h
+4 -3
View File
@@ -542,9 +542,10 @@ template <typename T>
array::array(
std::initializer_list<T> data,
Dtype dtype /* = TypeToDtype<T>() */)
: array_desc_(std::make_shared<ArrayDesc>(
Shape{static_cast<ShapeElem>(data.size())},
dtype)) {
: array_desc_(
std::make_shared<ArrayDesc>(
Shape{static_cast<ShapeElem>(data.size())},
dtype)) {
init(data.begin());
}
mlx/backend/cpu/compiled.cpp
+8 -6
View File
@@ -119,13 +119,15 @@ void* compile(
source_file.close();
try {
JitCompiler::exec(JitCompiler::build_command(
output_dir, source_file_name, shared_lib_name));
JitCompiler::exec(
JitCompiler::build_command(
output_dir, source_file_name, shared_lib_name));
} catch (const std::exception& error) {
throw std::runtime_error(fmt::format(
"[Compile::eval_cpu] Failed to compile function {0}: {1}",
kernel_name,
error.what()));
throw std::runtime_error(
fmt::format(
"[Compile::eval_cpu] Failed to compile function {0}: {1}",
kernel_name,
error.what()));
}
}
mlx/backend/cpu/jit_compiler.cpp
+17 -14
View File
@@ -36,10 +36,11 @@ struct VisualStudioInfo {
// Get path of Visual Studio.
// Use -latest to get only the most recent installation when multiple
// versions are installed, avoiding path concatenation issues.
std::string vs_path = JitCompiler::exec(fmt::format(
"\"{0}\\Microsoft Visual Studio\\Installer\\vswhere.exe\""
" -latest -property installationPath",
std::getenv("ProgramFiles(x86)")));
std::string vs_path = JitCompiler::exec(
fmt::format(
"\"{0}\\Microsoft Visual Studio\\Installer\\vswhere.exe\""
" -latest -property installationPath",
std::getenv("ProgramFiles(x86)")));
if (vs_path.empty()) {
throw std::runtime_error("Can not find Visual Studio.");
}
@@ -52,10 +53,11 @@ struct VisualStudioInfo {
.base(),
vs_path.end());
// Read the envs from vcvarsall.
std::string envs = JitCompiler::exec(fmt::format(
"\"{0}\\VC\\Auxiliary\\Build\\vcvarsall.bat\" {1} >NUL && set",
vs_path,
arch));
std::string envs = JitCompiler::exec(
fmt::format(
"\"{0}\\VC\\Auxiliary\\Build\\vcvarsall.bat\" {1} >NUL && set",
vs_path,
arch));
for (const std::string& line : str_split(envs, '\n')) {
// Each line is in the format "ENV_NAME=values".
auto pos = line.find_first_of('=');
@@ -150,12 +152,13 @@ std::string JitCompiler::exec(const std::string& cmd) {
int code = WEXITSTATUS(status);
#endif
if (code != 0) {
throw std::runtime_error(fmt::format(
"Failed to execute command with return code {0}: \"{1}\", "
"the output is: {2}",
code,
cmd,
ret));
throw std::runtime_error(
fmt::format(
"Failed to execute command with return code {0}: \"{1}\", "
"the output is: {2}",
code,
cmd,
ret));
}
return ret;
}
mlx/backend/cuda/binary/binary.cuh
+6 -5
View File
@@ -346,11 +346,12 @@ void binary_op_gpu_inplace(
});
}
} else {
throw std::runtime_error(fmt::format(
"Can not do binary op {} on inputs of {} with result of {}.",
op,
dtype_to_string(a.dtype()),
dtype_to_string(out.dtype())));
throw std::runtime_error(
fmt::format(
"Can not do binary op {} on inputs of {} with result of {}.",
op,
dtype_to_string(a.dtype()),
dtype_to_string(out.dtype())));
}
});
});
mlx/backend/cuda/binary_two.cu
+6 -5
View File
@@ -376,11 +376,12 @@ void binary_two_op_gpu_inplace(
});
}
} else {
throw std::runtime_error(fmt::format(
"Can not do binary op {} on inputs of {} with result of {}.",
op,
dtype_to_string(a.dtype()),
dtype_to_string(out_a.dtype())));
throw std::runtime_error(
fmt::format(
"Can not do binary op {} on inputs of {} with result of {}.",
op,
dtype_to_string(a.dtype()),
dtype_to_string(out_a.dtype())));
}
});
});
mlx/backend/cuda/compiled.cpp
+29 -17
View File
@@ -36,14 +36,16 @@ struct FusedKernelBuilder {
params.push_back(
fmt::format("const {}* {}", dtype_to_cuda_type(x.dtype()), xname));
if (!is_scalar(x) && !contiguous) {
params.push_back(fmt::format(
"const __grid_constant__ cuda::std::array<int64_t, NDIM> {}_strides",
xname));
params.push_back(
fmt::format(
"const __grid_constant__ cuda::std::array<int64_t, NDIM> {}_strides",
xname));
}
}
for (const auto& x : outputs) {
params.push_back(fmt::format(
"{}* {}", dtype_to_cuda_type(x.dtype()), namer.get_name(x)));
params.push_back(
fmt::format(
"{}* {}", dtype_to_cuda_type(x.dtype()), namer.get_name(x)));
}
if (!contiguous) {
params.push_back(
@@ -250,20 +252,30 @@ void Compiled::eval_gpu(
builder.os += "\n} // namespace mlx::core::cu\n";
// Build kernel names.
std::vector<std::string> kernel_names;
kernel_names.push_back(fmt::format(
"mlx::core::cu::{}_contiguous<uint32_t, {}>",
lib_name(),
work_per_thread));
kernel_names.push_back(fmt::format(
"mlx::core::cu::{}_contiguous<int64_t, {}>",
lib_name(),
work_per_thread));
kernel_names.push_back(
fmt::format(
"mlx::core::cu::{}_contiguous<uint32_t, {}>",
lib_name(),
work_per_thread));
kernel_names.push_back(
fmt::format(
"mlx::core::cu::{}_contiguous<int64_t, {}>",
lib_name(),
work_per_thread));
for (int wpt : {1, work_per_thread}) {
for (int i = 1; i <= MAX_NDIM; ++i) {
kernel_names.push_back(fmt::format(
"mlx::core::cu::{}_strided<{}, uint32_t, {}>", lib_name(), i, wpt));
kernel_names.push_back(fmt::format(
"mlx::core::cu::{}_strided<{}, int64_t, {}>", lib_name(), i, wpt));
kernel_names.push_back(
fmt::format(
"mlx::core::cu::{}_strided<{}, uint32_t, {}>",
lib_name(),
i,
wpt));
kernel_names.push_back(
fmt::format(
"mlx::core::cu::{}_strided<{}, int64_t, {}>",
lib_name(),
i,
wpt));
}
}
mlx/backend/cuda/cublas_utils.h
+3 -2
View File
@@ -34,8 +34,9 @@ inline cudaDataType_t dtype_to_cublas_type(Dtype dtype, std::string_view tag) {
case complex64:
return CUDA_C_32F;
default:
throw std::runtime_error(fmt::format(
"Unsupported dtype in {}: {}.", tag, dtype_to_string(dtype)));
throw std::runtime_error(
fmt::format(
"Unsupported dtype in {}: {}.", tag, dtype_to_string(dtype)));
}
}
mlx/backend/cuda/cudnn_utils.h
+11 -9
View File
@@ -62,8 +62,9 @@ inline fe::DataType_t dtype_to_cudnn_type(Dtype dtype) {
case float64:
return fe::DataType_t::DOUBLE;
default:
throw std::runtime_error(fmt::format(
"Unsupported dtype in cuDNN: {}.", dtype_to_string(dtype)));
throw std::runtime_error(
fmt::format(
"Unsupported dtype in cuDNN: {}.", dtype_to_string(dtype)));
}
}
@@ -124,13 +125,14 @@ class DnnGraph : public fe::graph::Graph {
// Create a cuDNN tensor for scalar.
auto scalar(const char* name, int64_t uid, Dtype dtype) {
return Graph::tensor(fe::graph::Tensor_attributes()
.set_name(name)
.set_uid(uid)
.set_dim({1, 1, 1, 1})
.set_stride({1, 1, 1, 1})
.set_is_pass_by_value(true)
.set_data_type(dtype_to_cudnn_type(dtype)));
return Graph::tensor(
fe::graph::Tensor_attributes()
.set_name(name)
.set_uid(uid)
.set_dim({1, 1, 1, 1})
.set_stride({1, 1, 1, 1})
.set_is_pass_by_value(true)
.set_data_type(dtype_to_cudnn_type(dtype)));
}
// Call this before setting notes.
mlx/backend/cuda/cutlass_utils.cuh
+5 -4
View File
@@ -13,10 +13,11 @@ namespace mlx::core {
// Throw exception if the cutlass API does not succeed.
inline void check_cutlass_error(const char* name, cutlass::Status status) {
if (status != cutlass::Status::kSuccess) {
throw std::runtime_error(fmt::format(
"{} failed with code: {}.",
name,
cutlass::cutlassGetStatusString(status)));
throw std::runtime_error(
fmt::format(
"{} failed with code: {}.",
name,
cutlass::cutlassGetStatusString(status)));
}
}
mlx/backend/cuda/device.cpp
+4 -3
View File
@@ -42,9 +42,10 @@ Device::Device(int device) : device_(device) {
CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
&attr, cudaDevAttrConcurrentManagedAccess, device_));
if (attr != 1) {
throw std::runtime_error(fmt::format(
"Device {} does not support synchronization in managed memory.",
device_));
throw std::runtime_error(
fmt::format(
"Device {} does not support synchronization in managed memory.",
device_));
}
// The cublasLt handle is used by matmul.
mlx/backend/cuda/event.cu
+4 -3
View File
@@ -119,9 +119,10 @@ void CudaEvent::init_pool() {
class CopyableCudaEvent {
public:
explicit CopyableCudaEvent(Device& d)
: event_(std::make_shared<CudaEvent>(
d,
cudaEventDisableTiming | cudaEventBlockingSync)) {}
: event_(
std::make_shared<CudaEvent>(
d,
cudaEventDisableTiming | cudaEventBlockingSync)) {}
void wait() {
event_->wait();
mlx/backend/cuda/gemms/cublas_gemm.cpp
+3 -2
View File
@@ -27,8 +27,9 @@ cublasComputeType_t dtype_to_compute_type(Dtype dtype) {
return mlx::core::env::enable_tf32() ? CUBLAS_COMPUTE_32F_FAST_TF32
: CUBLAS_COMPUTE_32F;
default:
throw std::runtime_error(fmt::format(
"Unsupported dtype in CublasGemm: {}.", dtype_to_string(dtype)));
throw std::runtime_error(
fmt::format(
"Unsupported dtype in CublasGemm: {}.", dtype_to_string(dtype)));
}
}
mlx/backend/cuda/indexing.cpp
+36 -32
View File
@@ -86,13 +86,14 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
std::vector<std::string> kernel_names;
for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) {
for (int large = 0; large <= 1; ++large) {
kernel_names.push_back(fmt::format(
"mlx::core::cu::gather<{}, {}, {}, {}, {}>",
dtype_to_cuda_type(out.dtype()),
dtype_to_cuda_type(idx_dtype),
nidx,
ndim,
large ? "int64_t" : "int32_t"));
kernel_names.push_back(
fmt::format(
"mlx::core::cu::gather<{}, {}, {}, {}, {}>",
dtype_to_cuda_type(out.dtype()),
dtype_to_cuda_type(idx_dtype),
nidx,
ndim,
large ? "int64_t" : "int32_t"));
}
}
return std::make_tuple(false, jit_source_gather, std::move(kernel_names));
@@ -179,14 +180,15 @@ void Scatter::eval_gpu(const std::vector<array>& inputs, array& out) {
std::vector<std::string> kernel_names;
for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) {
for (int large = 0; large <= 1; ++large) {
kernel_names.push_back(fmt::format(
"mlx::core::cu::scatter<{}, {}, mlx::core::cu::Scatter{}, {}, {}, {}>",
dtype_to_cuda_type(out.dtype()),
dtype_to_cuda_type(idx_dtype),
op,
nidx,
ndim,
large ? "int64_t" : "int32_t"));
kernel_names.push_back(
fmt::format(
"mlx::core::cu::scatter<{}, {}, mlx::core::cu::Scatter{}, {}, {}, {}>",
dtype_to_cuda_type(out.dtype()),
dtype_to_cuda_type(idx_dtype),
op,
nidx,
ndim,
large ? "int64_t" : "int32_t"));
}
}
return std::make_tuple(false, jit_source_scatter, std::move(kernel_names));
@@ -258,14 +260,15 @@ void GatherAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) {
for (int contiguous = 0; contiguous < 4; ++contiguous) {
for (int large = 0; large <= 1; ++large) {
kernel_names.push_back(fmt::format(
"mlx::core::cu::gather_axis<{}, {}, {}, {}, {}, {}>",
dtype_to_cuda_type(out.dtype()),
dtype_to_cuda_type(idx.dtype()),
ndim,
contiguous & 1 ? true : false,
contiguous & 2 ? true : false,
large ? "int64_t" : "int32_t"));
kernel_names.push_back(
fmt::format(
"mlx::core::cu::gather_axis<{}, {}, {}, {}, {}, {}>",
dtype_to_cuda_type(out.dtype()),
dtype_to_cuda_type(idx.dtype()),
ndim,
contiguous & 1 ? true : false,
contiguous & 2 ? true : false,
large ? "int64_t" : "int32_t"));
}
}
}
@@ -360,15 +363,16 @@ void ScatterAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) {
for (int contiguous = 0; contiguous < 4; ++contiguous) {
for (int large = 0; large <= 1; ++large) {
kernel_names.push_back(fmt::format(
"mlx::core::cu::scatter_axis<{}, {}, mlx::core::cu::Scatter{}, {}, {}, {}, {}>",
dtype_to_cuda_type(out.dtype()),
dtype_to_cuda_type(idx.dtype()),
op,
ndim,
contiguous & 1 ? true : false,
contiguous & 2 ? true : false,
large ? "int64_t" : "int32_t"));
kernel_names.push_back(
fmt::format(
"mlx::core::cu::scatter_axis<{}, {}, mlx::core::cu::Scatter{}, {}, {}, {}, {}>",
dtype_to_cuda_type(out.dtype()),
dtype_to_cuda_type(idx.dtype()),
op,
ndim,
contiguous & 1 ? true : false,
contiguous & 2 ? true : false,
large ? "int64_t" : "int32_t"));
}
}
}
mlx/backend/cuda/jit_module.cpp
+3 -2
View File
@@ -330,8 +330,9 @@ void load_module(
CUresult jit_result = cuModuleLoadDataEx(
&module_, ptx.data(), std::size(options), options, values);
if (jit_result != CUDA_SUCCESS) {
throw std::runtime_error(fmt::format(
"Failed to load compiled {} kernel: {}.", module_name, jit_log));
throw std::runtime_error(
fmt::format(
"Failed to load compiled {} kernel: {}.", module_name, jit_log));
}
// Load kernels.
mlx/backend/cuda/lru_cache.h
+6 -5
View File
@@ -89,11 +89,12 @@ class LRUCache {
}
if (env_name_ && ++cache_misses_ > 2 * capacity_) {
throw std::runtime_error(fmt::format(
"Cache thrashing is happening, please set the environment variable "
"{} to a larger value than {} to fix degraded performance.",
env_name_,
capacity_));
throw std::runtime_error(
fmt::format(
"Cache thrashing is happening, please set the environment variable "
"{} to a larger value than {} to fix degraded performance.",
env_name_,
capacity_));
}
vlist_.emplace_front(key, std::forward<U>(value));
mlx/backend/cuda/scan.cu
+6 -5
View File
@@ -454,11 +454,12 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
});
});
} else {
throw std::runtime_error(fmt::format(
"Can not do scan op {} on inputs of {} with result of {}.",
op_to_string<Op>(),
dtype_to_string(in.dtype()),
dtype_to_string(out.dtype())));
throw std::runtime_error(
fmt::format(
"Can not do scan op {} on inputs of {} with result of {}.",
op_to_string<Op>(),
dtype_to_string(in.dtype()),
dtype_to_string(out.dtype())));
}
});
});
mlx/backend/cuda/sort.cu
+1 -1
View File
@@ -1073,4 +1073,4 @@ void Partition::eval_gpu(const std::vector<array>& inputs, array& out) {
gpu_sort(stream(), inputs[0], out, axis_, false);
}
} // namespace mlx::core
} // namespace mlx::core
mlx/backend/cuda/unary/unary.cuh
+6 -5
View File
@@ -191,11 +191,12 @@ void unary_op_gpu_inplace(
}
});
} else {
throw std::runtime_error(fmt::format(
"Can not do unary op {} on input of {} with output of {}.",
op,
dtype_to_string(in.dtype()),
dtype_to_string(out.dtype())));
throw std::runtime_error(
fmt::format(
"Can not do unary op {} on input of {} with output of {}.",
op,
dtype_to_string(in.dtype()),
dtype_to_string(out.dtype())));
}
});
});
mlx/backend/metal/kernels/conv.metal
+34 -34
View File
@@ -445,20 +445,20 @@ winograd_conv_2d_weight_transform(
}
}
#define instantiate_winograd_conv_2d_weight_transform_base(name, itype, bc) \
template [[host_name("winograd_conv_2d_weight_transform_" #name \
"_bc" #bc)]] [[kernel]] void \
winograd_conv_2d_weight_transform<itype, bc>( \
const device itype* wt_in [[buffer(0)]], \
device itype* wt_out [[buffer(1)]], \
const constant int& C [[buffer(2)]], \
const constant int& O [[buffer(3)]], \
uint tid [[threadgroup_position_in_grid]], \
uint simd_group_id [[simdgroup_index_in_threadgroup]], \
#define instantiate_winograd_conv_2d_weight_transform_base(name, itype, bc) \
template [[host_name( \
"winograd_conv_2d_weight_transform_" #name "_bc" #bc)]] [[kernel]] void \
winograd_conv_2d_weight_transform<itype, bc>( \
const device itype* wt_in [[buffer(0)]], \
device itype* wt_out [[buffer(1)]], \
const constant int& C [[buffer(2)]], \
const constant int& O [[buffer(3)]], \
uint tid [[threadgroup_position_in_grid]], \
uint simd_group_id [[simdgroup_index_in_threadgroup]], \
uint simd_lane_id [[thread_index_in_simdgroup]]);
template <typename T, int BC, int WM, int WN, int M = 6, int R = 3>
[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void
[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void
winograd_conv_2d_input_transform(
const device T* inp_in [[buffer(0)]],
device T* inp_out [[buffer(1)]],
@@ -555,21 +555,21 @@ winograd_conv_2d_input_transform(
}
}
#define instantiate_winograd_conv_2d_input_transform(name, itype, bc) \
template [[host_name("winograd_conv_2d_input_transform_" #name \
"_bc" #bc)]] [[kernel]] void \
winograd_conv_2d_input_transform<itype, bc, 2, 2>( \
const device itype* inp_in [[buffer(0)]], \
device itype* inp_out [[buffer(1)]], \
const constant MLXConvParams<2>& params [[buffer(2)]], \
uint3 tid [[threadgroup_position_in_grid]], \
uint3 lid [[thread_position_in_threadgroup]], \
uint3 tgp_per_grid [[threadgroups_per_grid]], \
uint simd_group_id [[simdgroup_index_in_threadgroup]], \
#define instantiate_winograd_conv_2d_input_transform(name, itype, bc) \
template [[host_name( \
"winograd_conv_2d_input_transform_" #name "_bc" #bc)]] [[kernel]] void \
winograd_conv_2d_input_transform<itype, bc, 2, 2>( \
const device itype* inp_in [[buffer(0)]], \
device itype* inp_out [[buffer(1)]], \
const constant MLXConvParams<2>& params [[buffer(2)]], \
uint3 tid [[threadgroup_position_in_grid]], \
uint3 lid [[thread_position_in_threadgroup]], \
uint3 tgp_per_grid [[threadgroups_per_grid]], \
uint simd_group_id [[simdgroup_index_in_threadgroup]], \
uint simd_lane_id [[thread_index_in_simdgroup]]);
template <typename T, int BO, int WM, int WN, int M = 6, int R = 3>
[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void
[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void
winograd_conv_2d_output_transform(
const device T* out_in [[buffer(0)]],
device T* out_out [[buffer(1)]],
@@ -676,17 +676,17 @@ winograd_conv_2d_output_transform(
}
}
#define instantiate_winograd_conv_2d_output_transform(name, itype, bo) \
template [[host_name("winograd_conv_2d_output_transform_" #name \
"_bo" #bo)]] [[kernel]] void \
winograd_conv_2d_output_transform<itype, bo, 2, 2>( \
const device itype* out_in [[buffer(0)]], \
device itype* out_out [[buffer(1)]], \
const constant MLXConvParams<2>& params [[buffer(2)]], \
uint3 tid [[threadgroup_position_in_grid]], \
uint3 lid [[thread_position_in_threadgroup]], \
uint3 tgp_per_grid [[threadgroups_per_grid]], \
uint simd_group_id [[simdgroup_index_in_threadgroup]], \
#define instantiate_winograd_conv_2d_output_transform(name, itype, bo) \
template [[host_name( \
"winograd_conv_2d_output_transform_" #name "_bo" #bo)]] [[kernel]] void \
winograd_conv_2d_output_transform<itype, bo, 2, 2>( \
const device itype* out_in [[buffer(0)]], \
device itype* out_out [[buffer(1)]], \
const constant MLXConvParams<2>& params [[buffer(2)]], \
uint3 tid [[threadgroup_position_in_grid]], \
uint3 lid [[thread_position_in_threadgroup]], \
uint3 tgp_per_grid [[threadgroups_per_grid]], \
uint simd_group_id [[simdgroup_index_in_threadgroup]], \
uint simd_lane_id [[thread_index_in_simdgroup]]);
// clang-format off
mlx/backend/metal/kernels/gemv.metal
+4 -4
View File
@@ -445,7 +445,7 @@ template <
const int TN, /* Thread cols (in elements) */
const bool kDoNCBatch, /* Batch ndim > 1 */
const bool kDoAxpby> /* Do out = alpha * out + beta * bias */
[[kernel, max_total_threads_per_threadgroup(BM* BN * 32)]] void gemv(
[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv(
const device T* mat [[buffer(0)]],
const device T* in_vec [[buffer(1)]],
const device T* bias [[buffer(2)]],
@@ -553,7 +553,7 @@ template <
const int SN, /* Simdgroup cols (in threads) */
const int TM, /* Thread rows (in elements) */
const int TN> /* Thread cols (in elements) */
[[kernel, max_total_threads_per_threadgroup(BM* BN * 32)]] void gemv_gather(
[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv_gather(
const device T* mat [[buffer(0)]],
const device T* in_vec [[buffer(1)]],
const device T* bias [[buffer(2)]],
@@ -666,7 +666,7 @@ template <
const int TN, /* Thread cols (in elements) */
const bool kDoNCBatch, /* Batch ndim > 1 */
const bool kDoAxpby> /* Do out = alpha * out + beta * bias */
[[kernel, max_total_threads_per_threadgroup(BM* BN * 32)]] void gemv_t(
[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv_t(
const device T* mat [[buffer(0)]],
const device T* in_vec [[buffer(1)]],
const device T* bias [[buffer(2)]],
@@ -764,7 +764,7 @@ template <
const int SN, /* Simdgroup cols (in threads) */
const int TM, /* Thread rows (in elements) */
const int TN> /* Thread cols (in elements) */
[[kernel, max_total_threads_per_threadgroup(BM* BN * 32)]] void gemv_t_gather(
[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv_t_gather(
const device T* mat [[buffer(0)]],
const device T* in_vec [[buffer(1)]],
const device T* bias [[buffer(2)]],
mlx/backend/metal/kernels/gemv_masked.h
+2 -2
View File
@@ -641,7 +641,7 @@ template <
const int TM, /* Thread rows (in elements) */
const int TN, /* Thread cols (in elements) */
const bool kDoNCBatch> /* Batch ndim > 1 */
[[kernel, max_total_threads_per_threadgroup(BM* BN * 32)]] void gemv_masked(
[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv_masked(
const device T* mat [[buffer(0)]],
const device T* in_vec [[buffer(1)]],
device T* out_vec [[buffer(3)]],
@@ -741,7 +741,7 @@ template <
const int TM, /* Thread rows (in elements) */
const int TN, /* Thread cols (in elements) */
const bool kDoNCBatch> /* Batch ndim > 1 */
[[kernel, max_total_threads_per_threadgroup(BM* BN * 32)]] void gemv_t_masked(
[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv_t_masked(
const device T* mat [[buffer(0)]],
const device T* in_vec [[buffer(1)]],
device T* out_vec [[buffer(3)]],
mlx/backend/metal/kernels/steel/attn/mma.h
+2 -2
View File
@@ -68,8 +68,8 @@ struct BaseMMAFrag<T, 8, 8> {
template <typename U>
using dtype_frag_t = typename metal::vec<U, kElemsPerFrag>;
METAL_FUNC static constexpr short2 get_coord(ushort simd_lane_id
[[thread_index_in_simdgroup]]) {
METAL_FUNC static constexpr short2 get_coord(
ushort simd_lane_id [[thread_index_in_simdgroup]]) {
const short qid = simd_lane_id / 4;
const short fm = (qid & 4) + ((simd_lane_id / 2) % 4);
const short fn = (qid & 2) * 2 + (simd_lane_id % 2) * 2;
mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.h
+1 -1
View File
@@ -13,7 +13,7 @@ template <
int WN,
int N_CHANNELS = 0,
bool SMALL_FILTER = false>
[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void
[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void
implicit_gemm_conv_2d(
const device T* A [[buffer(0)]],
const device T* B [[buffer(1)]],
mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h
+1 -1
View File
@@ -13,7 +13,7 @@ template <
int WN,
typename AccumType = float,
typename Epilogue = TransformNone<T, AccumType>>
[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void
[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void
implicit_gemm_conv_2d_general(
const device T* A [[buffer(0)]],
const device T* B [[buffer(1)]],
mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather.h
+2 -2
View File
@@ -17,7 +17,7 @@ template <
bool transpose_a,
bool transpose_b,
typename AccumType = float>
[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void gather_mm_rhs(
[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void gather_mm_rhs(
const device T* A [[buffer(0)]],
const device T* B [[buffer(1)]],
const device uint32_t* rhs_indices [[buffer(2)]],
@@ -248,7 +248,7 @@ template <
bool transpose_a,
bool transpose_b,
typename AccumType = float>
[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void gather_mm(
[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void gather_mm(
const device T* A [[buffer(0)]],
const device T* B [[buffer(1)]],
const device uint32_t* lhs_indices [[buffer(2)]],
mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.h
+1 -1
View File
@@ -16,7 +16,7 @@ template <
bool transpose_a,
bool transpose_b,
typename AccumType = float>
[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void
[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void
gather_mm_rhs_nax(
const device T* A [[buffer(0)]],
const device T* B [[buffer(1)]],
mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.h
+2 -2
View File
@@ -49,7 +49,7 @@ template <
bool transpose_b,
bool MN_aligned,
bool K_aligned>
[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void
[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void
block_masked_gemm(
const device T* A [[buffer(0)]],
const device T* B [[buffer(1)]],
@@ -435,7 +435,7 @@ template <
bool MN_aligned,
bool K_aligned,
bool has_operand_mask = false>
[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void
[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void
block_masked_gemm(
const device T* A [[buffer(0)]],
const device T* B [[buffer(1)]],
mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented.h
+1 -1
View File
@@ -16,7 +16,7 @@ template <
bool transpose_a,
bool transpose_b,
typename AccumType = float>
[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void segmented_mm(
[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void segmented_mm(
const device T* A [[buffer(0)]],
const device T* B [[buffer(1)]],
const device uint32_t* segments [[buffer(2)]],
mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk.h
+1 -1
View File
@@ -18,7 +18,7 @@ template <
bool transpose_b,
bool MN_aligned,
bool K_aligned>
[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void gemm_splitk(
[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void gemm_splitk(
const device T* A [[buffer(0)]],
const device T* B [[buffer(1)]],
device U* C [[buffer(2)]],
mlx/backend/metal/kernels/steel/gemm/mma.h
+2 -2
View File
@@ -46,8 +46,8 @@ struct BaseMMAFrag<T, 8, 8> {
typedef metal::simdgroup_matrix<T, kFragRows, kFragCols> mat_type;
typedef metal::vec<T, kElemsPerFrag> frag_type;
METAL_FUNC static constexpr short2 get_coord(ushort simd_lane_id
[[thread_index_in_simdgroup]]) {
METAL_FUNC static constexpr short2 get_coord(
ushort simd_lane_id [[thread_index_in_simdgroup]]) {
const short qid = simd_lane_id / 4;
const short fm = (qid & 4) + ((simd_lane_id / 2) % 4);
const short fn = (qid & 2) * 2 + (simd_lane_id % 2) * 2;
mlx/backend/metal/matmul.cpp
+82 -89
View File
@@ -270,21 +270,20 @@ void steel_matmul_regular_axpby_nax(
int swizzle_log = tm <= 3 ? 0 : 1;
// Prepare steel matmul params
GEMMParams params{
/* const int M = */ M,
/* const int N = */ N,
/* const int K = */ K,
/* const int lda = */ lda,
/* const int ldb = */ ldb,
/* const int ldd = */ ldd,
/* const int tiles_n = */ tn,
/* const int tiles_m = */ tm,
/* const int64_t batch_stride_a = */ A_batch_stride,
/* const int64_t batch_stride_b = */ B_batch_stride,
/* const int64_t batch_stride_d = */ matrix_stride_out,
/* const int swizzle_log = */ swizzle_log,
/* const int gemm_k_iterations_aligned = */ (K / bk),
/* const int batch_ndim = */ int(batch_shape.size())};
GEMMParams params{/* const int M = */ M,
/* const int N = */ N,
/* const int K = */ K,
/* const int lda = */ lda,
/* const int ldb = */ ldb,
/* const int ldd = */ ldd,
/* const int tiles_n = */ tn,
/* const int tiles_m = */ tm,
/* const int64_t batch_stride_a = */ A_batch_stride,
/* const int64_t batch_stride_b = */ B_batch_stride,
/* const int64_t batch_stride_d = */ matrix_stride_out,
/* const int swizzle_log = */ swizzle_log,
/* const int gemm_k_iterations_aligned = */ (K / bk),
/* const int batch_ndim = */ int(batch_shape.size())};
// Prepare launch grid params
int tile = 1 << swizzle_log;
@@ -310,12 +309,11 @@ void steel_matmul_regular_axpby_nax(
int ldc = c.strides()[c.ndim() - 2];
int fdc = c.strides()[c.ndim() - 1];
GEMMAddMMParams params{
/* const int ldc = */ ldc,
/* const int fdc = */ fdc,
/* const int64_t batch_stride_c = */ C_batch_stride,
/* const float alpha = */ alpha,
/* const float beta = */ beta};
GEMMAddMMParams params{/* const int ldc = */ ldc,
/* const int fdc = */ fdc,
/* const int64_t batch_stride_c = */ C_batch_stride,
/* const float alpha = */ alpha,
/* const float beta = */ beta};
compute_encoder.set_input_array(c, 2);
compute_encoder.set_bytes(params, 5);
@@ -457,21 +455,20 @@ void steel_matmul_regular_axpby(
int swizzle_log = 0; // tm >= 6 ? 3 : (tm <= 3 ? 0 : 2);
// Prepare steel matmul params
GEMMParams params{
/* const int M = */ M,
/* const int N = */ N,
/* const int K = */ K,
/* const int lda = */ lda,
/* const int ldb = */ ldb,
/* const int ldd = */ ldd,
/* const int tiles_n = */ tn,
/* const int tiles_m = */ tm,
/* const int64_t batch_stride_a = */ A_batch_stride,
/* const int64_t batch_stride_b = */ B_batch_stride,
/* const int64_t batch_stride_d = */ matrix_stride_out,
/* const int swizzle_log = */ swizzle_log,
/* const int gemm_k_iterations_aligned = */ (K / bk),
/* const int batch_ndim = */ int(batch_shape.size())};
GEMMParams params{/* const int M = */ M,
/* const int N = */ N,
/* const int K = */ K,
/* const int lda = */ lda,
/* const int ldb = */ ldb,
/* const int ldd = */ ldd,
/* const int tiles_n = */ tn,
/* const int tiles_m = */ tm,
/* const int64_t batch_stride_a = */ A_batch_stride,
/* const int64_t batch_stride_b = */ B_batch_stride,
/* const int64_t batch_stride_d = */ matrix_stride_out,
/* const int swizzle_log = */ swizzle_log,
/* const int gemm_k_iterations_aligned = */ (K / bk),
/* const int batch_ndim = */ int(batch_shape.size())};
// Prepare launch grid params
int tile = 1 << swizzle_log;
@@ -497,12 +494,11 @@ void steel_matmul_regular_axpby(
int ldc = c.strides()[c.ndim() - 2];
int fdc = c.strides()[c.ndim() - 1];
GEMMAddMMParams params{
/* const int ldc = */ ldc,
/* const int fdc = */ fdc,
/* const int64_t batch_stride_c = */ C_batch_stride,
/* const float alpha = */ alpha,
/* const float beta = */ beta};
GEMMAddMMParams params{/* const int ldc = */ ldc,
/* const int fdc = */ fdc,
/* const int64_t batch_stride_c = */ C_batch_stride,
/* const float alpha = */ alpha,
/* const float beta = */ beta};
compute_encoder.set_input_array(c, 2);
compute_encoder.set_bytes(params, 5);
@@ -1554,21 +1550,20 @@ void BlockMaskedMM::eval_gpu(const std::vector<array>& inputs, array& out) {
int swizzle_log = 0; // tm >= 6 ? 3 : (tm <= 3 ? 0 : 2);
// Prepare steel matmul params
GEMMParams params{
/* const int M = */ M,
/* const int N = */ N,
/* const int K = */ K,
/* const int lda = */ lda,
/* const int ldb = */ ldb,
/* const int ldd = */ N,
/* const int tiles_n = */ tn,
/* const int tiles_m = */ tm,
/* const int64_t batch_stride_a = */ A_batch_str,
/* const int64_t batch_stride_b = */ B_batch_str,
/* const int64_t batch_stride_d = */ matrix_stride_out,
/* const int swizzle_log = */ swizzle_log,
/* const int gemm_k_iterations_aligned = */ (K / bk),
/* const int batch_ndim = */ int(batch_shape.size())};
GEMMParams params{/* const int M = */ M,
/* const int N = */ N,
/* const int K = */ K,
/* const int lda = */ lda,
/* const int ldb = */ ldb,
/* const int ldd = */ N,
/* const int tiles_n = */ tn,
/* const int tiles_m = */ tm,
/* const int64_t batch_stride_a = */ A_batch_str,
/* const int64_t batch_stride_b = */ B_batch_str,
/* const int64_t batch_stride_d = */ matrix_stride_out,
/* const int swizzle_log = */ swizzle_log,
/* const int gemm_k_iterations_aligned = */ (K / bk),
/* const int batch_ndim = */ int(batch_shape.size())};
// Prepare launch grid params
int tile = 1 << swizzle_log;
@@ -2113,23 +2108,22 @@ void gather_mm(
compute_encoder.set_compute_pipeline_state(kernel);
// Prepare the matmul params
steel::GEMMParams params{
/* const int M = */ M,
/* const int N = */ N,
/* const int K = */ K,
/* const int lda = */ static_cast<int>(lda),
/* const int ldb = */ static_cast<int>(ldb),
/* const int ldd = */ N,
/* const int tiles_n = */ (N + bn - 1) / bn,
/* const int tiles_m = */ (M + bm - 1) / bm,
/* const int64_t batch_stride_a = */
(batch_ndim > 0) ? lhs_indices.strides()[0] : 0,
/* const int64_t batch_stride_b = */
(batch_ndim > 0) ? rhs_indices.strides()[0] : 0,
/* const int64_t batch_stride_d = */ M * N,
/* const int swizzle_log = */ 0,
/* const int gemm_k_iterations_aligned = */ (K / bk),
/* const int batch_ndim = */ batch_ndim};
steel::GEMMParams params{/* const int M = */ M,
/* const int N = */ N,
/* const int K = */ K,
/* const int lda = */ static_cast<int>(lda),
/* const int ldb = */ static_cast<int>(ldb),
/* const int ldd = */ N,
/* const int tiles_n = */ (N + bn - 1) / bn,
/* const int tiles_m = */ (M + bm - 1) / bm,
/* const int64_t batch_stride_a = */
(batch_ndim > 0) ? lhs_indices.strides()[0] : 0,
/* const int64_t batch_stride_b = */
(batch_ndim > 0) ? rhs_indices.strides()[0] : 0,
/* const int64_t batch_stride_d = */ M * N,
/* const int swizzle_log = */ 0,
/* const int gemm_k_iterations_aligned = */ (K / bk),
/* const int batch_ndim = */ batch_ndim};
// Prepare the grid
MTL::Size group_dims = MTL::Size(32, wn, wm);
@@ -2317,21 +2311,20 @@ void segmented_mm(
compute_encoder.set_compute_pipeline_state(kernel);
// Prepare the matmul params
steel::GEMMParams params{
/* const int M = */ M,
/* const int N = */ N,
/* const int K = */ K,
/* const int lda = */ static_cast<int>(lda),
/* const int ldb = */ static_cast<int>(ldb),
/* const int ldd = */ N,
/* const int tiles_n = */ (N + bn - 1) / bn,
/* const int tiles_m = */ (M + bm - 1) / bm,
/* const int64_t batch_stride_a = */ 0,
/* const int64_t batch_stride_b = */ 0,
/* const int64_t batch_stride_d = */ M * N,
/* const int swizzle_log = */ 0,
/* const int gemm_k_iterations_aligned = */ 0,
/* const int batch_ndim = */ 0};
steel::GEMMParams params{/* const int M = */ M,
/* const int N = */ N,
/* const int K = */ K,
/* const int lda = */ static_cast<int>(lda),
/* const int ldb = */ static_cast<int>(ldb),
/* const int ldd = */ N,
/* const int tiles_n = */ (N + bn - 1) / bn,
/* const int tiles_m = */ (M + bm - 1) / bm,
/* const int64_t batch_stride_a = */ 0,
/* const int64_t batch_stride_b = */ 0,
/* const int64_t batch_stride_d = */ M * N,
/* const int swizzle_log = */ 0,
/* const int gemm_k_iterations_aligned = */ 0,
/* const int batch_ndim = */ 0};
// Prepare the grid
MTL::Size group_dims = MTL::Size(32, wn, wm);
mlx/distributed/jaccl/jaccl.cpp
+6 -5
View File
@@ -468,11 +468,12 @@ class SideChannel {
}
}
} else {
sockets_.push_back(detail::TCPSocket::connect(
IBV_TAG, address, 4, 1000, [](int attempt, int wait) {
std::cerr << IBV_TAG << " Connection attempt " << attempt
<< " waiting " << wait << " ms" << std::endl;
}));
sockets_.push_back(
detail::TCPSocket::connect(
IBV_TAG, address, 4, 1000, [](int attempt, int wait) {
std::cerr << IBV_TAG << " Connection attempt " << attempt
<< " waiting " << wait << " ms" << std::endl;
}));
sockets_[0].send(IBV_TAG, reinterpret_cast<char*>(&rank_), sizeof(int));
}
}
mlx/distributed/ring/ring.cpp
+42 -39
View File
@@ -353,23 +353,24 @@ std::vector<int> make_connections(
int success;
for (auto& address : addresses) {
sockets.push_back(detail::TCPSocket::connect(
RING_TAG,
address,
CONN_ATTEMPTS,
CONN_WAIT,
[verbose](int attempt, int wait) {
log_info(
verbose,
"Attempt",
attempt,
"waiting",
wait,
"ms (error:",
errno,
")");
})
.detach());
sockets.push_back(
detail::TCPSocket::connect(
RING_TAG,
address,
CONN_ATTEMPTS,
CONN_WAIT,
[verbose](int attempt, int wait) {
log_info(
verbose,
"Attempt",
attempt,
"waiting",
wait,
"ms (error:",
errno,
")");
})
.detach());
}
return sockets;
@@ -510,17 +511,18 @@ class RingGroup : public GroupImpl {
std::vector<std::future<void>> all_gathers;
for (int i = 0; i < n_gathers; i++) {
auto offset = i * bytes_per_gather;
all_gathers.emplace_back(pool_.enqueue(std::bind(
&RingGroup::all_gather_impl,
this,
input_ptr + offset,
output_ptr + offset,
nbytes,
offset + bytes_per_gather > nbytes ? nbytes - offset
: bytes_per_gather,
sockets_right_[i / 2],
sockets_left_[i / 2],
(i % 2) ? -1 : 1)));
all_gathers.emplace_back(pool_.enqueue(
std::bind(
&RingGroup::all_gather_impl,
this,
input_ptr + offset,
output_ptr + offset,
nbytes,
offset + bytes_per_gather > nbytes ? nbytes - offset
: bytes_per_gather,
sockets_right_[i / 2],
sockets_left_[i / 2],
(i % 2) ? -1 : 1)));
}
for (auto& f : all_gathers) {
f.wait();
@@ -634,17 +636,18 @@ class RingGroup : public GroupImpl {
std::vector<std::future<void>> all_sums;
for (int i = 0; i < n_reduces; i++) {
all_sums.emplace_back(pool_.enqueue(std::bind(
&RingGroup::all_reduce_impl<T, ReduceOp>,
this,
reinterpret_cast<T*>(
buffers_.data() + i * ALL_SUM_SIZE * ALL_SUM_BUFFERS),
reinterpret_cast<T*>(out_ptr) + i * step,
std::min(size, (i + 1) * step) - i * step,
sockets_right_[i / 2],
sockets_left_[i / 2],
(i % 2) ? -1 : 1,
reduce_op)));
all_sums.emplace_back(pool_.enqueue(
std::bind(
&RingGroup::all_reduce_impl<T, ReduceOp>,
this,
reinterpret_cast<T*>(
buffers_.data() + i * ALL_SUM_SIZE * ALL_SUM_BUFFERS),
reinterpret_cast<T*>(out_ptr) + i * step,
std::min(size, (i + 1) * step) - i * step,
sockets_right_[i / 2],
sockets_left_[i / 2],
(i % 2) ? -1 : 1,
reduce_op)));
}
for (auto& f : all_sums) {
f.wait();
python/src/ops.cpp
+12 -8
View File
@@ -53,8 +53,9 @@ void init_ops(nb::module_& m) {
"shape"_a,
nb::kw_only(),
"stream"_a = nb::none(),
nb::sig("def reshape(a: array, /, shape: Sequence[int], *, stream: "
"Union[None, Stream, Device] = None) -> array"),
nb::sig(
"def reshape(a: array, /, shape: Sequence[int], *, stream: "
"Union[None, Stream, Device] = None) -> array"),
R"pbdoc(
Reshape an array while preserving the size.
@@ -80,8 +81,9 @@ void init_ops(nb::module_& m) {
"end_axis"_a = -1,
nb::kw_only(),
"stream"_a = nb::none(),
nb::sig("def flatten(a: array, /, start_axis: int = 0, end_axis: int = "
"-1, *, stream: Union[None, Stream, Device] = None) -> array"),
nb::sig(
"def flatten(a: array, /, start_axis: int = 0, end_axis: int = "
"-1, *, stream: Union[None, Stream, Device] = None) -> array"),
R"pbdoc(
Flatten an array.
@@ -182,8 +184,9 @@ void init_ops(nb::module_& m) {
"axis"_a,
nb::kw_only(),
"stream"_a = nb::none(),
nb::sig("def expand_dims(a: array, /, axis: Union[int, Sequence[int]], "
"*, stream: Union[None, Stream, Device] = None) -> array"),
nb::sig(
"def expand_dims(a: array, /, axis: Union[int, Sequence[int]], "
"*, stream: Union[None, Stream, Device] = None) -> array"),
R"pbdoc(
Add a size one dimension at the given axis.
@@ -1675,8 +1678,9 @@ void init_ops(nb::module_& m) {
},
nb::arg(),
"dtype"_a = nb::none(),
nb::sig("def asarray(a: Union[scalar, array, Sequence], dtype: "
"Optional[Dtype] = None) -> array"),
nb::sig(
"def asarray(a: Union[scalar, array, Sequence], dtype: "
"Optional[Dtype] = None) -> array"),
R"pbdoc(
Convert the input to an array.
tests/linalg_tests.cpp
+6 -4
View File
@@ -394,12 +394,14 @@ TEST_CASE("test matrix cholesky") {
linalg::cholesky(array({0.0, 1.0}), /* upper = */ false, Device::cpu));
// Unsupported types throw
CHECK_THROWS(linalg::cholesky(
array({0, 1}, {1, 2}), /* upper = */ false, Device::cpu));
CHECK_THROWS(
linalg::cholesky(
array({0, 1}, {1, 2}), /* upper = */ false, Device::cpu));
// Non-square throws.
CHECK_THROWS(linalg::cholesky(
array({1, 2, 3, 4, 5, 6}, {2, 3}), /* upper = */ false, Device::cpu));
CHECK_THROWS(
linalg::cholesky(
array({1, 2, 3, 4, 5, 6}, {2, 3}), /* upper = */ false, Device::cpu));
const auto prng_key = random::key(220398);
const auto sqrtA = random::normal({5, 5}, prng_key);