03_arithmetic_topk/
03_arithmetic_topk.rs

1use apple_metal::MetalDevice;
2use apple_mpsgraph::{
3    data_type, Feed, Graph, ReductionAxesOp, TensorData, UnaryArithmeticOp,
4};
5
6fn main() {
7    let device = MetalDevice::system_default().expect("no Metal device available");
8    let graph = Graph::new().expect("graph");
9    let input = graph
10        .placeholder(Some(&[2, 3]), data_type::FLOAT32, Some("input"))
11        .expect("placeholder");
12    let squared = graph
13        .unary_arithmetic(UnaryArithmeticOp::Square, &input, Some("square"))
14        .expect("square");
15    let row_sum = graph
16        .reduce_axes(ReductionAxesOp::Sum, &squared, &[1], Some("row_sum"))
17        .expect("reduce");
18    let topk = graph.top_k(&input, 2, Some("topk")).expect("topk");
19
20    let input_data = TensorData::from_f32_slice(&device, &[1.0, 3.0, 2.0, 4.0, 6.0, 5.0], &[2, 3])
21        .expect("tensor data");
22    let results = graph
23        .run(&[Feed::new(&input, &input_data)], &[&row_sum, &topk.0])
24        .expect("run");
25
26    println!("row sums: {:?}", results[0].read_f32().expect("row sums"));
27    println!("top-k values: {:?}", results[1].read_f32().expect("topk values"));
28}
03_arithmetic_topk/03_arithmetic_topk.rs

03_arithmetic_topk/
03_arithmetic_topk.rs
