1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
use objc2::{
__framework_prelude::*,
extern_class, extern_conformance, extern_methods, msg_send,
rc::{Allocated, Retained},
runtime::NSObject,
};
use objc2_foundation::NSObjectProtocol;
extern_class!(
/// A container for tensor rank and extents per dimension.
///
/// [Apple's documentation](https://developer.apple.com/documentation/metal/mtltensorextents?language=objc)
#[unsafe(super(NSObject))]
#[derive(Debug, PartialEq, Eq, Hash)]
pub struct MTLTensorExtents;
);
extern_conformance!(
unsafe impl NSObjectProtocol for MTLTensorExtents {}
);
impl MTLTensorExtents {
extern_methods!(
/// Obtains the rank of the tensor.
///
/// The rank represents the number of dimensions.
#[unsafe(method(rank))]
#[unsafe(method_family = none)]
pub fn rank(&self) -> usize;
/// Returns the extent at an index.
///
/// Parameter `dimensionIndex`: the index of the dimension. The first dimension is the innermost dimension.
/// Returns: the extent at `dimensionIndex`. Returns -1 if `dimensionIndex` is greater than or equal to `rank`.
#[unsafe(method(extentAtDimensionIndex:))]
#[unsafe(method_family = none)]
pub fn extent_at_dimension_index(
&self,
dimension_index: usize,
) -> isize;
);
}
/// Methods declared on superclass `NSObject`.
impl MTLTensorExtents {
extern_methods!(
#[unsafe(method(init))]
#[unsafe(method_family = init)]
pub fn init(this: Allocated<Self>) -> Retained<Self>;
#[unsafe(method(new))]
#[unsafe(method_family = new)]
pub fn new() -> Retained<Self>;
);
}
impl MTLTensorExtents {
/// Creates a new tensor extents with the rank and extent values you provide.
///
/// Zero rank extents represent scalars. `values` can only be `nil`if `rank` is 0.
/// - Parameters:
/// - rank: the number of dimensions.
/// - values: an array of length `rank` that specifies the size of each dimension. The first dimension is the innermost dimension.
/// - Returns: Tensor extents with the rank and extent values you provide. Returns `nil` if `rank` exceeds 0 and `values` is nil or if `rank` exceeds ``MTL_TENSOR_MAX_RANK``.
pub fn new_with_rank_values(
rank: usize,
values: Option<&[isize]>,
) -> Option<Retained<Self>> {
let values_ptr: *const isize = match values {
Some(v) => {
debug_assert_eq!(v.len(), rank);
v.as_ptr()
},
None => {
debug_assert_eq!(rank, 0);
core::ptr::null()
},
};
unsafe {
let this: Allocated<Self> = msg_send![Self::class(), alloc];
msg_send![this, initWithRank: rank, values: values_ptr]
}
}
}