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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
mod scalars_or_tensors;
pub use scalars_or_tensors::SpatialAxesBatchAxisBlockDimensionsScalarsOrTensors;
use crate::{Graph, Tensor, ns_number_array_from_slice};
use objc2::{extern_methods, msg_send, rc::Retained};
use objc2_foundation::NSString;
impl Graph {
/// Creates a *space-to-batch* operation.
///
/// Values from the `spatial_axes` dimensions are packed into blocks of
/// `block_dimensions` and moved to the `batch_axis` dimension.
/// Setting `use_pixel_shuffle_order` controls whether the blocks are stored
/// contiguously (pixel-shuffle order) or interleaved.
///
/// # Arguments
///
/// * `tensor` – Input tensor.
/// * `spatial_axes_batch_axis_block_dimensions` – Tuple of spatial axes,
/// destination batch axis, and per-axis block sizes. Accepts scalars or
/// tensors via [`SpatialAxesBatchAxisBlockDimensionsScalarsOrTensors`].
/// * `use_pixel_shuffle_order` – If `true`, blocks are laid out
/// contiguously along the batch axis.
/// * `name` – Optional debug label.
///
/// # Returns
///
/// A [`Tensor`] with spatial data moved into the batch dimension.
pub fn space_to_batch<'a>(
&self,
tensor: &Tensor,
spatial_axes_batch_axis_block_dimensions: SpatialAxesBatchAxisBlockDimensionsScalarsOrTensors<'a>,
use_pixel_shuffle_order: bool,
name: Option<&str>,
) -> Retained<Tensor> {
match spatial_axes_batch_axis_block_dimensions {
SpatialAxesBatchAxisBlockDimensionsScalarsOrTensors::Scalars {
spatial_axes,
batch_axis,
block_dimensions,
} => unsafe {
msg_send![
self,
spaceToBatchTensor: tensor,
spatialAxes: &*ns_number_array_from_slice(spatial_axes),
batchAxis: batch_axis,
blockDimensions: &*ns_number_array_from_slice(block_dimensions),
usePixelShuffleOrder: use_pixel_shuffle_order,
name: name.map(NSString::from_str).as_deref(),
]
},
SpatialAxesBatchAxisBlockDimensionsScalarsOrTensors::Tensors {
spatial_axes_tensor,
batch_axis_tensor,
block_dimensions_tensor,
} => unsafe {
msg_send![
self,
spaceToBatchTensor: tensor,
spatialAxesTensor: spatial_axes_tensor,
batchAxisTensor: batch_axis_tensor,
blockDimensionsTensor: block_dimensions_tensor,
usePixelShuffleOrder: use_pixel_shuffle_order,
name: name.map(NSString::from_str).as_deref(),
]
},
}
}
/// Creates a *batch-to-space* operation (inverse of space-to-batch).
///
/// Values from the `batch_axis` are unpacked into spatial blocks of
/// `block_dimensions` and distributed across the `spatial_axes`.
///
/// # Arguments
///
/// * `tensor` – Input tensor.
/// * `spatial_axes_batch_axis_block_dimensions` – Tuple of spatial axes,
/// source batch axis, and per-axis block sizes.
/// * `use_pixel_shuffle_order` – If `true`, expects contiguous
/// pixel-shuffle layout in the batch dimension.
/// * `name` – Optional debug label.
///
/// # Returns
///
/// A [`Tensor`] with batch data moved back into spatial dimensions.
pub fn batch_to_space<'a>(
&self,
tensor: &Tensor,
spatial_axes_batch_axis_block_dimensions: SpatialAxesBatchAxisBlockDimensionsScalarsOrTensors<'a>,
use_pixel_shuffle_order: bool,
name: Option<&str>,
) -> Retained<Tensor> {
match spatial_axes_batch_axis_block_dimensions {
SpatialAxesBatchAxisBlockDimensionsScalarsOrTensors::Scalars {
spatial_axes,
batch_axis,
block_dimensions,
} => unsafe {
msg_send![
self,
batchToSpaceTensor: tensor,
spatialAxes: &*ns_number_array_from_slice(spatial_axes),
batchAxis: batch_axis,
blockDimensions: &*ns_number_array_from_slice(block_dimensions),
usePixelShuffleOrder: use_pixel_shuffle_order,
name: name.map(NSString::from_str).as_deref(),
]
},
SpatialAxesBatchAxisBlockDimensionsScalarsOrTensors::Tensors {
spatial_axes_tensor,
batch_axis_tensor,
block_dimensions_tensor,
} => unsafe {
msg_send![
self,
batchToSpaceTensor: tensor,
spatialAxesTensor: spatial_axes_tensor,
batchAxisTensor: batch_axis_tensor,
blockDimensionsTensor: block_dimensions_tensor,
usePixelShuffleOrder: use_pixel_shuffle_order,
name: name.map(NSString::from_str).as_deref(),
]
},
}
}
}