mirror of
https://github.com/ml-explore/mlx.git
synced 2025-09-18 01:50:16 +08:00
Close a couple edge case bugs: hadamard and addmm on empty inputs (#2177)
* handle hadamard and addmm on empty inputs * fix
This commit is contained in:
@@ -132,6 +132,10 @@ void AddMM::eval_cpu(const std::vector<array>& inputs, array& out) {
|
||||
throw std::runtime_error(
|
||||
"[AddMM::eval_cpu] Currently only supports float32.");
|
||||
}
|
||||
if (out.size() == 0) {
|
||||
out.set_data(allocator::malloc(out.nbytes()));
|
||||
return;
|
||||
}
|
||||
|
||||
// Fill output with C
|
||||
auto& c = inputs[2];
|
||||
@@ -139,7 +143,9 @@ void AddMM::eval_cpu(const std::vector<array>& inputs, array& out) {
|
||||
? CopyType::Scalar
|
||||
: (c.flags().row_contiguous ? CopyType::Vector : CopyType::General);
|
||||
copy(c, out, ctype, stream());
|
||||
|
||||
if (inputs[0].shape(-1) == 0) {
|
||||
return;
|
||||
}
|
||||
matmul_general(inputs[0], inputs[1], out, stream(), alpha_, beta_);
|
||||
}
|
||||
|
||||
|
@@ -716,6 +716,23 @@ void AddMM::eval_gpu(const std::vector<array>& inputs, array& out) {
|
||||
throw std::runtime_error(
|
||||
"[matmul] Does not yet support non-floating point types.");
|
||||
}
|
||||
|
||||
// Return 0s if either input is empty
|
||||
if (out.size() == 0) {
|
||||
out.set_data(allocator::malloc(out.nbytes()));
|
||||
return;
|
||||
}
|
||||
|
||||
// Copy c into out and return
|
||||
if (inputs[0].shape(-1) == 0) {
|
||||
copy_gpu(
|
||||
inputs[2],
|
||||
out,
|
||||
inputs[2].flags().row_contiguous ? CopyType::Vector : CopyType::General,
|
||||
stream());
|
||||
return;
|
||||
}
|
||||
|
||||
out.set_data(allocator::malloc(out.nbytes()));
|
||||
auto& s = stream();
|
||||
auto& d = metal::device(s.device);
|
||||
|
@@ -472,6 +472,10 @@ array hadamard_transform(
|
||||
const array& a,
|
||||
std::optional<float> scale_ /* = std::nullopt */,
|
||||
StreamOrDevice s /* = {} */) {
|
||||
if (a.size() == 0) {
|
||||
throw std::invalid_argument(
|
||||
"[hadamard_transform] Does not support empty arrays.");
|
||||
}
|
||||
// Default to an orthonormal Hadamard matrix scaled by 1/sqrt(N)
|
||||
int n = a.ndim() > 0 ? a.shape(-1) : 1;
|
||||
float scale = scale_.has_value() ? *scale_ : 1.0f / std::sqrt(n);
|
||||
@@ -4326,6 +4330,10 @@ array addmm(
|
||||
|
||||
c = reshape(c, c_reshape, s);
|
||||
}
|
||||
if (c.shape() != out_shape) {
|
||||
throw std::invalid_argument(
|
||||
"[addmm] input c must broadcast to the output shape");
|
||||
}
|
||||
|
||||
auto out = array(
|
||||
std::move(out_shape),
|
||||
|
Reference in New Issue
Block a user