mirror of
https://github.com/ml-explore/mlx.git
synced 2025-07-14 04:11:12 +08:00
16 lines
244 B
Plaintext
16 lines
244 B
Plaintext
|
// Copyright © 2025 Apple Inc.
|
||
|
|
|
||
|
|
namespace mlx::core::cu {
|
||
|
|
|
||
|
|
template <typename T>
|
||
|
|
struct Arange {
|
||
|
|
const T start;
|
||
|
|
const T step;
|
||
|
|
|
||
|
|
__device__ T operator()(uint32_t i) const {
|
||
|
|
return start + i * step;
|
||
|
|
}
|
||
|
|
};
|
||
|
|
|
||
|
|
} // namespace mlx::core::cu
|