Files
mlx/mlx/backend/cuda/kernels/arange.cuh
2025-05-06 21:26:46 -07:00

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