@@ -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+
1933impl 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