mirror of
https://github.com/ml-explore/mlx.git
synced 2025-06-24 17:31:16 +08:00
40c62c1321
* use int64 stride everywhere * fix ext * fix ext * more shape + cleanup * one more * few more
46 lines
1.0 KiB
C++
46 lines
1.0 KiB
C++
// Copyright © 2023 Apple Inc.
|
|
|
|
#include "doctest/doctest.h"
|
|
|
|
#include "mlx/mlx.h"
|
|
|
|
using namespace mlx::core;
|
|
|
|
TEST_CASE("test type promotion") {
|
|
for (auto t : {bool_, uint32, int32, int64, float32}) {
|
|
auto a = array(0, t);
|
|
CHECK_EQ(result_type({a}), t);
|
|
|
|
std::vector<array> arrs = {array(0, t), array(0, t)};
|
|
CHECK_EQ(result_type(arrs), t);
|
|
}
|
|
|
|
{
|
|
std::vector<array> arrs = {array(false), array(0, int32)};
|
|
CHECK_EQ(result_type(arrs), int32);
|
|
}
|
|
|
|
{
|
|
std::vector<array> arrs = {array(0, int32), array(false), array(0.0f)};
|
|
CHECK_EQ(result_type(arrs), float32);
|
|
}
|
|
}
|
|
|
|
TEST_CASE("test normalize axis") {
|
|
struct TestCase {
|
|
int axis;
|
|
int ndim;
|
|
int expected;
|
|
};
|
|
|
|
std::vector<TestCase> testCases = {
|
|
{0, 3, 0}, {1, 3, 1}, {2, 3, 2}, {-1, 3, 2}, {-2, 3, 1}, {-3, 3, 0}};
|
|
|
|
for (const auto& tc : testCases) {
|
|
CHECK_EQ(normalize_axis_index(tc.axis, tc.ndim), tc.expected);
|
|
}
|
|
|
|
CHECK_THROWS(normalize_axis_index(3, 3));
|
|
CHECK_THROWS(normalize_axis_index(-4, 3));
|
|
}
|