Skip to content

Commit b169eaf

Browse files
authored
Add fill_op (#7)
Signed-off-by: kerthcet <kerthcet@gmail.com>
1 parent 1950710 commit b169eaf

1 file changed

Lines changed: 179 additions & 23 deletions

File tree

‎crates/mlxcore/src/array.rs‎

Lines changed: 179 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,20 @@ pub struct Array {
1616
handle: sys::mlx_array,
1717
}
1818

19+
/// Returns a pointer to `slice`'s data, or null when it is empty.
20+
///
21+
/// For an empty slice `as_ptr()` is non-null but dangling. The mlx-c functions
22+
/// we call take a `(ptr, len)` pair and build a container from it, so they never
23+
/// dereference the pointer when `len == 0` — but passing an explicit null keeps
24+
/// a dangling pointer from crossing the FFI boundary at all.
25+
fn as_ffi_ptr<T>(slice: &[T]) -> *const T {
26+
if slice.is_empty() {
27+
std::ptr::null()
28+
} else {
29+
slice.as_ptr()
30+
}
31+
}
32+
1933
impl Array {
2034
/// Creates an `Array` from a raw handle, taking ownership of it.
2135
///
@@ -46,17 +60,77 @@ impl Array {
4660
data.len()
4761
);
4862
error::install();
49-
// SAFETY: pointers/len are valid for the duration of the call; mlx
50-
// copies the data into its own buffer.
51-
let handle = unsafe {
52-
sys::mlx_array_new_data(
53-
data.as_ptr() as *const _,
54-
shape.as_ptr(),
55-
shape.len() as i32,
63+
// Both empty cases are reachable here: an empty `shape` for a 0-d
64+
// scalar, empty `data` for a 0-element array.
65+
let data_ptr = as_ffi_ptr(data) as *const _;
66+
let shape_ptr = as_ffi_ptr(shape);
67+
// SAFETY: pointers/len describe valid slices (or null/0) for the
68+
// duration of the call; mlx copies the data into its own buffer.
69+
let handle =
70+
unsafe { sys::mlx_array_new_data(data_ptr, shape_ptr, shape.len() as i32, T::DTYPE) };
71+
unsafe { Self::from_raw(handle) }
72+
}
73+
74+
/// Builds a 0-dimensional array holding a single value.
75+
///
76+
/// The dtype follows the value's type, so `Array::from_scalar(2.0f32)` is a
77+
/// `float32` scalar. Handy as the operand of a broadcasting op.
78+
pub fn from_scalar<T: ArrayElement>(value: T) -> Self {
79+
Self::from_slice(&[value], &[])
80+
}
81+
82+
/// An array of `shape` filled with zeros, with the dtype chosen by `T`.
83+
///
84+
/// Call as `Array::zeros::<f32>(&[2, 3], &stream)`.
85+
pub fn zeros<T: ArrayElement>(shape: &[i32], stream: &Stream) -> Result<Array> {
86+
Self::fill_op(shape, T::DTYPE, stream, sys::mlx_zeros)
87+
}
88+
89+
/// An array of `shape` filled with ones, with the dtype chosen by `T`.
90+
///
91+
/// Call as `Array::ones::<f32>(&[2, 3], &stream)`.
92+
pub fn ones<T: ArrayElement>(shape: &[i32], stream: &Stream) -> Result<Array> {
93+
Self::fill_op(shape, T::DTYPE, stream, sys::mlx_ones)
94+
}
95+
96+
/// An array of `shape` filled with `value`, with the dtype chosen by `T`.
97+
pub fn full<T: ArrayElement>(shape: &[i32], value: T, stream: &Stream) -> Result<Array> {
98+
error::install();
99+
// Held in a local so it outlives the call below.
100+
let vals = Self::from_scalar(value);
101+
let shape_ptr = as_ffi_ptr(shape);
102+
let mut out = unsafe { sys::mlx_array_new() };
103+
// SAFETY: `shape_ptr`/`shape.len()` describe a valid slice (or null/0);
104+
// all handles are valid; the result is written into `out`.
105+
let status = unsafe {
106+
sys::mlx_full(
107+
&mut out,
108+
shape_ptr,
109+
shape.len(),
110+
vals.as_raw(),
56111
T::DTYPE,
112+
stream.as_raw(),
57113
)
58114
};
59-
unsafe { Self::from_raw(handle) }
115+
Self::from_op(out, status)
116+
}
117+
118+
/// Values from `start` (inclusive) to `stop` (exclusive), stepping by `step`.
119+
///
120+
/// The bounds are `f64` because MLX's are; they are cast to the dtype chosen
121+
/// by `T`, as in `Array::arange::<i32>(0.0, 5.0, 1.0, &stream)`.
122+
pub fn arange<T: ArrayElement>(
123+
start: f64,
124+
stop: f64,
125+
step: f64,
126+
stream: &Stream,
127+
) -> Result<Array> {
128+
error::install();
129+
let mut out = unsafe { sys::mlx_array_new() };
130+
// SAFETY: stream is valid; the result is written into `out`.
131+
let status =
132+
unsafe { sys::mlx_arange(&mut out, start, stop, step, T::DTYPE, stream.as_raw()) };
133+
Self::from_op(out, status)
60134
}
61135

62136
/// Total number of elements.
@@ -378,6 +452,29 @@ impl Array {
378452
Self::from_op(out, status)
379453
}
380454

455+
/// Shared plumbing for `res = op(shape, shape_num, dtype, stream)`
456+
/// constructors.
457+
fn fill_op(
458+
shape: &[i32],
459+
dtype: sys::mlx_dtype,
460+
stream: &Stream,
461+
op: unsafe extern "C" fn(
462+
*mut sys::mlx_array,
463+
*const i32,
464+
usize,
465+
sys::mlx_dtype,
466+
sys::mlx_stream,
467+
) -> i32,
468+
) -> Result<Array> {
469+
error::install();
470+
let shape_ptr = as_ffi_ptr(shape);
471+
let mut out = unsafe { sys::mlx_array_new() };
472+
// SAFETY: `shape_ptr`/`shape.len()` describe a valid slice (or null/0);
473+
// stream is valid; `op` writes the result into `out`.
474+
let status = unsafe { op(&mut out, shape_ptr, shape.len(), dtype, stream.as_raw()) };
475+
Self::from_op(out, status)
476+
}
477+
381478
/// Shared plumbing for `res = op(a, shape, shape_num, stream)` shape ops.
382479
fn shape_op(
383480
&self,
@@ -392,13 +489,7 @@ impl Array {
392489
) -> i32,
393490
) -> Result<Array> {
394491
error::install();
395-
// For an empty slice `as_ptr()` is non-null but dangling; pass an
396-
// explicit null pointer so C never receives a bogus pointer.
397-
let shape_ptr = if shape.is_empty() {
398-
std::ptr::null()
399-
} else {
400-
shape.as_ptr()
401-
};
492+
let shape_ptr = as_ffi_ptr(shape);
402493
let mut out = unsafe { sys::mlx_array_new() };
403494
// SAFETY: `shape_ptr`/`shape.len()` describe a valid slice (or null/0)
404495
// for the call; all handles are valid; `op` writes into `out`.
@@ -477,14 +568,7 @@ impl Array {
477568
) -> i32,
478569
) -> Result<Array> {
479570
error::install();
480-
// For an empty slice `as_ptr()` is non-null but dangling; pass an
481-
// explicit null pointer so the FFI call never hands C a bogus pointer
482-
// even if it were to dereference it with `axes_num == 0`.
483-
let axes_ptr = if axes.is_empty() {
484-
std::ptr::null()
485-
} else {
486-
axes.as_ptr()
487-
};
571+
let axes_ptr = as_ffi_ptr(axes);
488572
let mut out = unsafe { sys::mlx_array_new() };
489573
// SAFETY: `axes_ptr`/`axes.len()` describe a valid slice (or null/0) for
490574
// the duration of the call; all handles are valid; `op` writes the
@@ -609,6 +693,78 @@ mod tests {
609693
assert!(a.shape().is_empty());
610694
}
611695

696+
#[test]
697+
fn from_scalar_is_zero_dim() {
698+
let a = Array::from_scalar(42.0f32);
699+
assert_eq!(a.ndim(), 0);
700+
assert_eq!(a.item::<f32>(), 42.0);
701+
}
702+
703+
#[test]
704+
fn empty_array_has_no_elements() {
705+
// Exercises the null-pointer path for empty `data`.
706+
let a = Array::from_slice::<f32>(&[], &[0]);
707+
assert_eq!(a.size(), 0);
708+
assert_eq!(a.shape(), vec![0]);
709+
assert!(a.to_vec::<f32>().is_empty());
710+
}
711+
712+
#[test]
713+
fn zeros_and_ones_fill_shape() {
714+
let s = Stream::cpu();
715+
716+
let z = Array::zeros::<f32>(&[2, 3], &s).unwrap();
717+
assert_eq!(z.shape(), vec![2, 3]);
718+
assert_eq!(z.to_vec::<f32>(), vec![0.0; 6]);
719+
720+
let o = Array::ones::<f32>(&[2, 2], &s).unwrap();
721+
assert_eq!(o.to_vec::<f32>(), vec![1.0; 4]);
722+
723+
// The type parameter picks the dtype, so integer arrays work too.
724+
let i = Array::ones::<i32>(&[3], &s).unwrap();
725+
assert_eq!(i.to_vec::<i32>(), vec![1, 1, 1]);
726+
}
727+
728+
#[test]
729+
fn full_broadcasts_a_scalar() {
730+
let s = Stream::cpu();
731+
let a = Array::full(&[2, 2], 7.5f32, &s).unwrap();
732+
assert_eq!(a.shape(), vec![2, 2]);
733+
assert_eq!(a.to_vec::<f32>(), vec![7.5; 4]);
734+
735+
// An empty shape yields a 0-d array, not an error.
736+
let scalar = Array::full(&[], 3i32, &s).unwrap();
737+
assert_eq!(scalar.ndim(), 0);
738+
assert_eq!(scalar.item::<i32>(), 3);
739+
}
740+
741+
#[test]
742+
fn arange_is_half_open() {
743+
let s = Stream::cpu();
744+
// `stop` is exclusive, so this is 0..5, not 0..=5.
745+
let a = Array::arange::<i32>(0.0, 5.0, 1.0, &s).unwrap();
746+
assert_eq!(a.to_vec::<i32>(), vec![0, 1, 2, 3, 4]);
747+
748+
let f = Array::arange::<f32>(1.0, 2.0, 0.5, &s).unwrap();
749+
assert_eq!(f.to_vec::<f32>(), vec![1.0, 1.5]);
750+
}
751+
752+
#[test]
753+
fn constructors_compose_with_ops() {
754+
let s = Stream::cpu();
755+
// ones + ones == twos, over a shape built by a constructor.
756+
let o = Array::ones::<f32>(&[4], &s).unwrap();
757+
assert_eq!(o.add(&o, &s).unwrap().to_vec::<f32>(), vec![2.0; 4]);
758+
759+
// A 0-d scalar broadcasts against a larger array.
760+
let a = Array::arange::<f32>(0.0, 4.0, 1.0, &s).unwrap();
761+
let two = Array::from_scalar(2.0f32);
762+
assert_eq!(
763+
a.multiply(&two, &s).unwrap().to_vec::<f32>(),
764+
vec![0.0, 2.0, 4.0, 6.0]
765+
);
766+
}
767+
612768
#[test]
613769
#[should_panic(expected = "does not match shape product")]
614770
fn mismatched_len_and_shape_panics() {

0 commit comments

Comments
 (0)