1use runmat_accelerate_api::HostTensorView;
4use runmat_builtins::{
5 BuiltinCompletionPolicy, BuiltinDescriptor, BuiltinErrorDescriptor, BuiltinOutputMode,
6 BuiltinParamArity, BuiltinParamDescriptor, BuiltinParamType, BuiltinSignatureDescriptor,
7 ResolveContext, Tensor, Type, Value,
8};
9use runmat_macros::runtime_builtin;
10
11use super::common::{
12 build_strides, dims_from_tokens, materialize_value, parse_dims, total_elements,
13};
14use crate::builtins::array::type_resolvers::size_vector_len;
15use crate::builtins::common::arg_tokens::tokens_from_context;
16use crate::builtins::common::spec::{
17 BroadcastSemantics, BuiltinFusionSpec, BuiltinGpuSpec, ConstantStrategy, GpuOpKind,
18 ProviderHook, ReductionNaN, ResidencyPolicy, ScalarType, ShapeRequirements,
19};
20use crate::builtins::common::tensor;
21use crate::{build_runtime_error, make_cell, RuntimeError};
22
23#[runmat_macros::register_gpu_spec(builtin_path = "crate::builtins::array::indexing::ind2sub")]
24pub const GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
25 name: "ind2sub",
26 op_kind: GpuOpKind::Custom("indexing"),
27 supported_precisions: &[ScalarType::F32, ScalarType::F64],
28 broadcast: BroadcastSemantics::Matlab,
29 provider_hooks: &[ProviderHook::Custom("ind2sub")],
30 constant_strategy: ConstantStrategy::InlineLiteral,
31 residency: ResidencyPolicy::NewHandle,
32 nan_mode: ReductionNaN::Include,
33 two_pass_threshold: None,
34 workgroup_size: None,
35 accepts_nan_mode: false,
36 notes: "WGPU provider executes `ind2sub` entirely on-device; other providers fall back to the host implementation and re-upload results to preserve residency.",
37};
38
39#[runmat_macros::register_fusion_spec(builtin_path = "crate::builtins::array::indexing::ind2sub")]
40pub const FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
41 name: "ind2sub",
42 shape: ShapeRequirements::Any,
43 constant_strategy: ConstantStrategy::InlineLiteral,
44 elementwise: None,
45 reduction: None,
46 emits_nan: false,
47 notes: "Index conversion is eager and does not participate in fusion today.",
48};
49
50fn ind2sub_type(args: &[Type], ctx: &ResolveContext) -> Type {
51 let Some(dims) = args.first() else {
52 return Type::Unknown;
53 };
54 let length = dims_from_tokens(&tokens_from_context(ctx))
55 .map(|values| values.len())
56 .or_else(|| size_vector_len(dims));
57 Type::Cell {
58 element_type: Some(Box::new(Type::tensor())),
59 length,
60 }
61}
62
63const BUILTIN_NAME: &str = "ind2sub";
64
65const IND2SUB_OUTPUT_CELL: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
66 name: "subs",
67 ty: BuiltinParamType::Any,
68 arity: BuiltinParamArity::Required,
69 default: None,
70 description: "Cell array containing one subscript output per dimension.",
71}];
72
73const IND2SUB_INPUTS: [BuiltinParamDescriptor; 2] = [
74 BuiltinParamDescriptor {
75 name: "sz",
76 ty: BuiltinParamType::SizeArg,
77 arity: BuiltinParamArity::Required,
78 default: None,
79 description: "Size vector describing source array dimensions.",
80 },
81 BuiltinParamDescriptor {
82 name: "ind",
83 ty: BuiltinParamType::Any,
84 arity: BuiltinParamArity::Required,
85 default: None,
86 description: "Linear indices to convert into per-dimension subscripts.",
87 },
88];
89
90const IND2SUB_SIGNATURES: [BuiltinSignatureDescriptor; 1] = [BuiltinSignatureDescriptor {
91 label: "subs = ind2sub(sz, ind)",
92 inputs: &IND2SUB_INPUTS,
93 outputs: &IND2SUB_OUTPUT_CELL,
94}];
95
96const IND2SUB_ERROR_INVALID_INPUT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
97 code: "RM.IND2SUB.INVALID_INPUT",
98 identifier: Some("RunMat:ind2sub:InvalidInput"),
99 when: "Size vector or linear index inputs are malformed or unsupported.",
100 message: "ind2sub: invalid input arguments",
101};
102
103const IND2SUB_ERROR_INDEX_BOUNDS: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
104 code: "RM.IND2SUB.INDEX_BOUNDS",
105 identifier: Some("RunMat:ind2sub:IndexBounds"),
106 when: "At least one provided linear index exceeds array element bounds.",
107 message: "ind2sub: index exceeds array bounds",
108};
109
110const IND2SUB_ERROR_PROVIDER: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
111 code: "RM.IND2SUB.PROVIDER",
112 identifier: Some("RunMat:ind2sub:ProviderError"),
113 when: "Provider-side ind2sub execution fails or returns malformed outputs.",
114 message: "ind2sub: provider execution failed",
115};
116
117const IND2SUB_ERROR_INTERNAL: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
118 code: "RM.IND2SUB.INTERNAL",
119 identifier: Some("RunMat:ind2sub:InternalError"),
120 when: "Internal tensor/materialization logic fails while building outputs.",
121 message: "ind2sub: internal error",
122};
123
124const IND2SUB_ERRORS: [BuiltinErrorDescriptor; 4] = [
125 IND2SUB_ERROR_INVALID_INPUT,
126 IND2SUB_ERROR_INDEX_BOUNDS,
127 IND2SUB_ERROR_PROVIDER,
128 IND2SUB_ERROR_INTERNAL,
129];
130
131pub const IND2SUB_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
132 signatures: &IND2SUB_SIGNATURES,
133 output_mode: BuiltinOutputMode::Fixed,
134 completion_policy: BuiltinCompletionPolicy::Public,
135 errors: &IND2SUB_ERRORS,
136};
137
138fn ind2sub_error_with_message(
139 message: impl Into<String>,
140 error: &'static BuiltinErrorDescriptor,
141) -> RuntimeError {
142 let mut builder = build_runtime_error(message).with_builtin(BUILTIN_NAME);
143 if let Some(identifier) = error.identifier {
144 builder = builder.with_identifier(identifier);
145 }
146 builder.build()
147}
148
149fn ind2sub_input_error(message: impl Into<String>) -> RuntimeError {
150 ind2sub_error_with_message(message, &IND2SUB_ERROR_INVALID_INPUT)
151}
152
153fn ind2sub_internal_error(message: impl Into<String>) -> RuntimeError {
154 ind2sub_error_with_message(message, &IND2SUB_ERROR_INTERNAL)
155}
156
157fn ind2sub_provider_error(message: impl Into<String>) -> RuntimeError {
158 ind2sub_error_with_message(message, &IND2SUB_ERROR_PROVIDER)
159}
160
161#[runtime_builtin(
162 name = "ind2sub",
163 category = "array/indexing",
164 summary = "Convert linear indices to subscripts.",
165 keywords = "ind2sub,linear index,subscripts,column major,gpu indexing",
166 accel = "custom",
167 type_resolver(ind2sub_type),
168 descriptor(crate::builtins::array::indexing::ind2sub::IND2SUB_DESCRIPTOR),
169 builtin_path = "crate::builtins::array::indexing::ind2sub"
170)]
171async fn ind2sub_builtin(dims_val: Value, indices_val: Value) -> crate::BuiltinResult<Value> {
172 let (dims_value, dims_was_gpu) = materialize_value(dims_val, "ind2sub").await?;
173 let dims = parse_dims(&dims_value, "ind2sub").await?;
174 if dims.is_empty() {
175 return Err(ind2sub_error("Size vector must have at least one element."));
176 }
177
178 let total = total_elements(&dims, "ind2sub")?;
179 let strides = build_strides(&dims, "ind2sub")?;
180
181 if let Some(result) = try_gpu_ind2sub(&dims, &strides, total, &indices_val)? {
182 return Ok(result);
183 }
184
185 let (indices_value, indices_was_gpu) = materialize_value(indices_val, "ind2sub").await?;
186 let indices_tensor = tensor::value_into_tensor_for("ind2sub", indices_value)
187 .map_err(|message| ind2sub_error(message))?;
188
189 let subscripts = compute_subscripts(&dims, total, &strides, &indices_tensor)?;
190
191 let want_gpu = (dims_was_gpu || indices_was_gpu) && runmat_accelerate_api::provider().is_some();
192
193 let mut outputs: Vec<Value> = Vec::with_capacity(dims.len());
194 for tensor in subscripts {
195 if want_gpu {
196 #[cfg(all(test, feature = "wgpu"))]
197 {
198 if runmat_accelerate_api::provider().is_none() {
199 let _ = runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(
200 runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default(),
201 );
202 }
203 }
204 if let Some(provider) = runmat_accelerate_api::provider() {
205 let view = HostTensorView {
206 data: &tensor.data,
207 shape: &tensor.shape,
208 };
209 if let Ok(handle) = provider.upload(&view) {
210 outputs.push(Value::GpuTensor(handle));
211 continue;
212 }
213 }
214 }
215 outputs.push(tensor::tensor_into_value(tensor));
216 }
217
218 make_cell(outputs, 1, dims.len()).map_err(|message| ind2sub_error(message))
219}
220
221fn try_gpu_ind2sub(
222 dims: &[usize],
223 strides: &[usize],
224 total: usize,
225 indices: &Value,
226) -> crate::BuiltinResult<Option<Value>> {
227 #[cfg(target_arch = "wasm32")]
228 {
229 let _ = (dims, strides, total, indices);
230 Ok(None)
231 }
232 #[cfg(not(target_arch = "wasm32"))]
233 {
234 let provider = match runmat_accelerate_api::provider() {
235 Some(p) => p,
236 None => return Ok(None),
237 };
238 if !provider.supports_ind2sub() {
239 return Ok(None);
240 }
241 let handle = match indices {
242 Value::GpuTensor(handle) => handle,
243 _ => return Ok(None),
244 };
245 if dims.len() != strides.len() {
246 return Err(ind2sub_error("Size vector must have at least one element."));
247 }
248 if dims.iter().any(|&d| d > u32::MAX as usize)
249 || strides.iter().any(|&s| s > u32::MAX as usize)
250 || total > u32::MAX as usize
251 {
252 return Ok(None);
253 }
254 let len = if handle.shape.is_empty() {
255 1usize
256 } else {
257 handle.shape.iter().copied().product()
258 };
259 if total == 0 && len > 0 {
260 return Err(ind2sub_error(
261 "Index exceeds number of array elements. Index must not exceed 0.",
262 ));
263 }
264 if len > u32::MAX as usize {
265 return Ok(None);
266 }
267 let output_shape = if handle.shape.is_empty() {
268 vec![len, 1]
269 } else {
270 handle.shape.clone()
271 };
272 match provider.ind2sub(dims, strides, handle, total, len, &output_shape) {
273 Ok(handles) => {
274 if handles.len() != dims.len() {
275 return Err(ind2sub_provider_error(
276 "ind2sub: provider returned an unexpected number of outputs.",
277 ));
278 }
279 let values: Vec<Value> = handles.into_iter().map(Value::GpuTensor).collect();
280 make_cell(values, 1, dims.len())
281 .map(Some)
282 .map_err(|message| ind2sub_error(message))
283 }
284 Err(err) => Err(ind2sub_provider_error(err.to_string())),
285 }
286 }
287}
288
289fn compute_subscripts(
290 dims: &[usize],
291 total: usize,
292 strides: &[usize],
293 indices: &Tensor,
294) -> crate::BuiltinResult<Vec<Tensor>> {
295 if strides.len() != dims.len() {
296 return Err(ind2sub_error("Size vector must have at least one element."));
297 }
298
299 let len = indices.data.len();
300 let mut outputs: Vec<Vec<f64>> = dims.iter().map(|_| Vec::with_capacity(len)).collect();
301
302 for &value in &indices.data {
303 let idx = coerce_linear_index(value, total)?;
304 let zero_based = idx - 1;
305 for (dim_index, (&dim, &stride)) in dims.iter().zip(strides.iter()).enumerate() {
306 let coord = ((zero_based / stride) % dim) + 1;
307 outputs[dim_index].push(coord as f64);
308 }
309 }
310
311 let output_shape = if indices.shape.is_empty() {
312 vec![len, 1]
313 } else {
314 indices.shape.clone()
315 };
316
317 let mut tensors = Vec::with_capacity(dims.len());
318 for data in outputs {
319 let tensor = Tensor::new(data, output_shape.clone())
320 .map_err(|e| ind2sub_internal_error(format!("ind2sub: {e}")))?;
321 tensors.push(tensor);
322 }
323 Ok(tensors)
324}
325
326fn coerce_linear_index(value: f64, max_index: usize) -> crate::BuiltinResult<usize> {
327 if !value.is_finite() {
328 return Err(ind2sub_error("Linear indices must be positive integers."));
329 }
330 let rounded = value.round();
331 if (rounded - value).abs() > f64::EPSILON {
332 return Err(ind2sub_error("Linear indices must be positive integers."));
333 }
334 if rounded < 1.0 {
335 return Err(ind2sub_error("Linear indices must be positive integers."));
336 }
337 if rounded > usize::MAX as f64 {
338 return Err(ind2sub_error(
339 "Index exceeds maximum supported size for this platform.",
340 ));
341 }
342 let coerced = rounded as usize;
343 if coerced > max_index {
344 return Err(ind2sub_error_with_message(
345 format!(
346 "Index exceeds number of array elements. Index must not exceed {}.",
347 max_index
348 ),
349 &IND2SUB_ERROR_INDEX_BOUNDS,
350 ));
351 }
352 Ok(coerced)
353}
354
355fn ind2sub_error(message: impl Into<String>) -> RuntimeError {
356 ind2sub_input_error(message)
357}
358
359#[cfg(test)]
360pub(crate) mod tests {
361 use crate::builtins::common::test_support;
362 use futures::executor::block_on;
363 use runmat_accelerate_api::HostTensorView;
364 use runmat_builtins::{ResolveContext, Tensor, Type, Value};
365
366 fn ind2sub_builtin(dims_val: Value, indices_val: Value) -> crate::BuiltinResult<Value> {
367 block_on(super::ind2sub_builtin(dims_val, indices_val))
368 }
369
370 fn cell_to_vec(cell: &runmat_builtins::CellArray) -> Vec<Value> {
371 cell.data.clone()
372 }
373
374 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
375 #[test]
376 fn recovers_tensor_indices() {
377 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
378 let result = ind2sub_builtin(Value::Tensor(dims), Value::Num(8.0)).unwrap();
379 match result {
380 Value::Cell(cell) => {
381 let values = cell_to_vec(&cell);
382 assert_eq!(values.len(), 2);
383 assert_eq!(values[0], Value::Num(2.0));
384 assert_eq!(values[1], Value::Num(3.0));
385 }
386 other => panic!("expected cell output, got {other:?}"),
387 }
388 }
389
390 #[test]
391 fn ind2sub_type_infers_cell_length() {
392 let dims = Type::Tensor {
393 shape: Some(vec![Some(1), Some(3)]),
394 };
395 assert_eq!(
396 super::ind2sub_type(&[dims, Type::Num], &ResolveContext::new(Vec::new())),
397 Type::Cell {
398 element_type: Some(Box::new(Type::tensor())),
399 length: Some(3)
400 }
401 );
402 }
403
404 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
405 #[test]
406 fn handles_vector_indices() {
407 let dims = Tensor::new(vec![3.0, 5.0], vec![1, 2]).unwrap();
408 let idx = Tensor::new(vec![7.0, 8.0, 9.0], vec![1, 3]).unwrap();
409 let result =
410 ind2sub_builtin(Value::Tensor(dims), Value::Tensor(idx)).expect("ind2sub result");
411 match result {
412 Value::Cell(cell) => {
413 let values = cell_to_vec(&cell);
414 assert_eq!(values.len(), 2);
415 match &values[0] {
416 Value::Tensor(t) => {
417 assert_eq!(t.shape, vec![1, 3]);
418 assert_eq!(t.data, vec![1.0, 2.0, 3.0]);
419 }
420 other => panic!("expected tensor rows, got {other:?}"),
421 }
422 match &values[1] {
423 Value::Tensor(t) => {
424 assert_eq!(t.shape, vec![1, 3]);
425 assert_eq!(t.data, vec![3.0, 3.0, 3.0]);
426 }
427 other => panic!("expected tensor cols, got {other:?}"),
428 }
429 }
430 other => panic!("expected cell output, got {other:?}"),
431 }
432 }
433
434 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
435 #[test]
436 fn rejects_non_integer_linear_index_identifier() {
437 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
438 let err = ind2sub_builtin(Value::Tensor(dims), Value::Num(1.25))
439 .expect_err("expected non-integer index error");
440 assert_eq!(
441 err.identifier(),
442 super::IND2SUB_ERROR_INVALID_INPUT.identifier
443 );
444 }
445
446 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
447 #[test]
448 fn rejects_out_of_bounds_linear_index_identifier() {
449 let dims = Tensor::new(vec![2.0, 2.0], vec![1, 2]).unwrap();
450 let err = ind2sub_builtin(Value::Tensor(dims), Value::Num(9.0))
451 .expect_err("expected out-of-bounds index error");
452 assert_eq!(
453 err.identifier(),
454 super::IND2SUB_ERROR_INDEX_BOUNDS.identifier
455 );
456 }
457
458 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
459 #[test]
460 fn recovers_three_dimensional_indices() {
461 let dims = Tensor::new(vec![2.0, 3.0, 4.0], vec![1, 3]).unwrap();
462 let idx = Tensor::new(vec![3.0, 11.0], vec![1, 2]).unwrap();
463 let result =
464 ind2sub_builtin(Value::Tensor(dims), Value::Tensor(idx)).expect("ind2sub result");
465 if let Value::Cell(cell) = result {
466 let values = cell_to_vec(&cell);
467 assert_eq!(values.len(), 3);
468 assert_eq!(
469 values[0],
470 Value::Tensor(Tensor::new(vec![1.0, 1.0], vec![1, 2]).unwrap())
471 );
472 assert_eq!(
473 values[1],
474 Value::Tensor(Tensor::new(vec![2.0, 3.0], vec![1, 2]).unwrap())
475 );
476 assert_eq!(
477 values[2],
478 Value::Tensor(Tensor::new(vec![1.0, 2.0], vec![1, 2]).unwrap())
479 );
480 } else {
481 panic!("expected cell output");
482 }
483 }
484
485 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
486 #[test]
487 fn errors_on_out_of_range_index() {
488 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
489 let err =
490 ind2sub_builtin(Value::Tensor(dims), Value::Num(13.0)).expect_err("expected failure");
491 assert!(
492 err.message()
493 .contains("Index exceeds number of array elements"),
494 "unexpected error: {err}"
495 );
496 }
497
498 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
499 #[test]
500 fn errors_on_zero_index() {
501 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
502 let err =
503 ind2sub_builtin(Value::Tensor(dims), Value::Num(0.0)).expect_err("expected failure");
504 assert!(
505 err.contains("Linear indices must be positive integers"),
506 "unexpected error: {err}"
507 );
508 }
509
510 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
511 #[test]
512 fn errors_on_fractional_index() {
513 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
514 let err =
515 ind2sub_builtin(Value::Tensor(dims), Value::Num(2.5)).expect_err("expected failure");
516 assert!(
517 err.contains("Linear indices must be positive integers"),
518 "unexpected error: {err}"
519 );
520 }
521
522 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
523 #[test]
524 fn errors_on_invalid_size_elements() {
525 let dims = Tensor::new(vec![3.5, 4.0], vec![1, 2]).unwrap();
526 let err = ind2sub_builtin(Value::Tensor(dims), Value::Num(5.0)).expect_err("expected fail");
527 assert!(
528 err.to_string()
529 .contains("Size arguments must be positive integers"),
530 "unexpected error: {err}"
531 );
532 }
533
534 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
535 #[test]
536 fn ind2sub_gpu_roundtrip() {
537 test_support::with_test_provider(|provider| {
538 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
539 let idx_tensor = Tensor::new(vec![10.0, 11.0], vec![2, 1]).unwrap();
540 let view = HostTensorView {
541 data: &idx_tensor.data,
542 shape: &idx_tensor.shape,
543 };
544 let handle = provider.upload(&view).expect("upload indices");
545 let result = ind2sub_builtin(Value::Tensor(dims), Value::GpuTensor(handle)).unwrap();
546 match result {
547 Value::Cell(cell) => {
548 let values = cell_to_vec(&cell);
549 assert_eq!(values.len(), 2);
550 match &values[0] {
551 Value::GpuTensor(_) => {}
552 other => panic!("expected gpu tensor output, got {other:?}"),
553 }
554 match &values[1] {
555 Value::GpuTensor(_) => {}
556 other => panic!("expected gpu tensor output, got {other:?}"),
557 }
558 let rows = test_support::gather(values[0].clone()).expect("gather rows");
559 assert_eq!(rows.shape, vec![2, 1]);
560 assert_eq!(rows.data, vec![1.0, 2.0]);
561 let cols = test_support::gather(values[1].clone()).expect("gather cols");
562 assert_eq!(cols.shape, vec![2, 1]);
563 assert_eq!(cols.data, vec![4.0, 4.0]);
564 }
565 other => panic!("expected cell output, got {other:?}"),
566 }
567 });
568 }
569
570 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
571 #[test]
572 #[cfg(feature = "wgpu")]
573 fn ind2sub_wgpu_matches_cpu() {
574 let provider_init = std::panic::catch_unwind(|| {
575 runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(
576 runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default(),
577 )
578 });
579 if let Ok(Ok(_)) = provider_init {
580 } else {
582 return;
583 }
584
585 let dims_tensor = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
586 let idx_tensor = Tensor::new(vec![7.0, 8.0, 9.0], vec![1, 3]).unwrap();
587
588 let cpu = ind2sub_builtin(
589 Value::Tensor(dims_tensor.clone()),
590 Value::Tensor(idx_tensor.clone()),
591 )
592 .expect("cpu ind2sub");
593
594 let provider = runmat_accelerate_api::provider().unwrap();
595 let view = HostTensorView {
596 data: &idx_tensor.data,
597 shape: &idx_tensor.shape,
598 };
599 let handle = provider.upload(&view).expect("upload indices");
600
601 let gpu = ind2sub_builtin(Value::Tensor(dims_tensor), Value::GpuTensor(handle))
602 .expect("gpu ind2sub");
603
604 let cpu_values = match cpu {
605 Value::Cell(cell) => cell_to_vec(&cell),
606 other => panic!("expected cell output, got {other:?}"),
607 };
608 let gpu_values = match gpu {
609 Value::Cell(cell) => cell_to_vec(&cell),
610 other => panic!("expected cell output, got {other:?}"),
611 };
612
613 assert_eq!(cpu_values.len(), gpu_values.len());
614
615 for (cpu_val, gpu_val) in cpu_values.iter().zip(gpu_values.iter()) {
616 let host_cpu = test_support::gather(cpu_val.clone()).expect("gather cpu");
617 let host_gpu = test_support::gather(gpu_val.clone()).expect("gather gpu");
618 assert_eq!(host_cpu.shape, host_gpu.shape);
619 assert_eq!(host_cpu.data, host_gpu.data);
620 }
621 }
622}