Struct TensorData

Source

pub struct TensorData { /* private fields */ }

Expand description

Safe owner for an Objective-C MPSGraphTensorData.

Implementations§

Source §

impl TensorData

Source

pub fn from_bytes( device: &MetalDevice, bytes: &[u8], shape: &[usize], data_type: u32, ) -> Option<Self>

Build tensor data by copying CPU bytes onto the given Metal device.

Examples found in repository ?

examples/06_control_flow_call.rs (line 127)

16fn main() {
17    let device = MetalDevice::system_default().expect("no Metal device available");
18    let queue = device
19        .new_command_queue()
20        .expect("failed to create command queue");
21
22    let callee_graph = Graph::new().expect("callee graph");
23    let callee_input = callee_graph
24        .placeholder(Some(&[2]), data_type::FLOAT32, Some("callee_input"))
25        .expect("callee placeholder");
26    let callee_output = callee_graph
27        .addition(&callee_input, &callee_input, Some("callee_double"))
28        .expect("callee output");
29    let callee_executable = callee_graph
30        .compile(
31            &device,
32            &[FeedDescription::new(&callee_input, &[2], data_type::FLOAT32)],
33            &[&callee_output],
34        )
35        .expect("callee executable");
36
37    let graph = Graph::new().expect("graph");
38    let input = graph
39        .placeholder(Some(&[2]), data_type::FLOAT32, Some("input"))
40        .expect("input placeholder");
41    let predicate = graph
42        .placeholder(Some(&[]), data_type::BOOL, Some("predicate"))
43        .expect("predicate placeholder");
44    let bias = graph.constant_f32_slice(&[1.0, 1.0], &[2]).expect("bias constant");
45
46    let output_type = ShapedType::new(Some(&[2]), data_type::FLOAT32).expect("output type");
47    let call_results = graph
48        .call("double", &[&input], &[&output_type], Some("call"))
49        .expect("call op");
50    let if_results = graph
51        .if_then_else(
52            &predicate,
53            || vec![graph.addition(&input, &bias, None).expect("then add")],
54            || vec![graph.subtraction(&input, &bias, None).expect("else sub")],
55            Some("branch"),
56        )
57        .expect("if/then/else");
58
59    let call_operation = call_results[0].operation().expect("call operation");
60    let dependency = graph
61        .control_dependency(&[&call_operation], || {
62            vec![graph
63                .unary_arithmetic(UnaryArithmeticOp::Identity, &call_results[0], None)
64                .expect("identity")]
65        }, Some("dependency"))
66        .expect("control dependency");
67
68    let number_of_iterations = graph
69        .constant_scalar(4.0, data_type::INT32)
70        .expect("iteration count");
71    let zero = graph
72        .constant_scalar(0.0, data_type::INT32)
73        .expect("zero constant");
74    let one = graph.constant_scalar(1.0, data_type::INT32).expect("one constant");
75    let limit = graph
76        .constant_scalar(3.0, data_type::INT32)
77        .expect("limit constant");
78
79    let for_results = graph
80        .for_loop_iterations(&number_of_iterations, &[&zero], |_index, args| {
81            vec![graph.addition(&args[0], &one, None).expect("for-loop add")]
82        }, Some("for_loop"))
83        .expect("for loop");
84    let while_results = graph
85        .while_loop(
86            &[&zero],
87            |inputs| {
88                let condition = graph
89                    .binary_arithmetic(BinaryArithmeticOp::LessThan, &inputs[0], &limit, None)
90                    .expect("while predicate");
91                let passthrough = graph
92                    .unary_arithmetic(UnaryArithmeticOp::Identity, &inputs[0], None)
93                    .expect("while passthrough");
94                WhileBeforeResult {
95                    predicate: condition,
96                    results: vec![passthrough],
97                }
98            },
99            |inputs| vec![graph.addition(&inputs[0], &one, None).expect("while add")],
100            Some("while_loop"),
101        )
102        .expect("while loop");
103
104    let compile_descriptor = CompilationDescriptor::new().expect("compile descriptor");
105    compile_descriptor
106        .set_callable("double", Some(&callee_executable))
107        .expect("set callable");
108    let executable = graph
109        .compile_with_descriptor(
110            Some(&device),
111            &[
112                FeedDescription::new(&input, &[2], data_type::FLOAT32),
113                FeedDescription::new(&predicate, &[], data_type::BOOL),
114            ],
115            &[
116                &call_results[0],
117                &if_results[0],
118                &dependency[0],
119                &for_results[0],
120                &while_results[0],
121            ],
122            Some(&compile_descriptor),
123        )
124        .expect("compile executable");
125
126    let input_data = TensorData::from_f32_slice(&device, &[3.0, 4.0], &[2]).expect("input data");
127    let predicate_data = TensorData::from_bytes(&device, &[1_u8], &[], data_type::BOOL)
128        .expect("predicate data");
129    let results = executable
130        .run(&queue, &[&input_data, &predicate_data])
131        .expect("run executable");
132
133    println!("call output: {:?}", results[0].read_f32().expect("call output"));
134    println!("if output: {:?}", results[1].read_f32().expect("if output"));
135    println!("dependency output: {:?}", results[2].read_f32().expect("dependency output"));
136    println!("for output: {:?}", read_i32(&results[3]));
137    println!("while output: {:?}", read_i32(&results[4]));
138}

Source

pub fn from_f32_slice( device: &MetalDevice, values: &[f32], shape: &[usize], ) -> Option<Self>

Build tensor data from a contiguous f32 slice.

Examples found in repository ?

examples/05_concat_split.rs (line 16)

4fn main() {
5    let device = MetalDevice::system_default().expect("no Metal device available");
6    let graph = Graph::new().expect("graph");
7    let input = graph
8        .placeholder(Some(&[2, 2]), data_type::FLOAT32, Some("input"))
9        .expect("placeholder");
10    let concat = graph
11        .concat_pair(&input, &input, 1, Some("concat"))
12        .expect("concat");
13    let split = graph.split_num(&concat, 2, 1, Some("split"));
14    let stacked = graph.stack(&[&split[0], &split[1]], 0, Some("stack")).expect("stack");
15
16    let input_data = TensorData::from_f32_slice(&device, &[1.0, 2.0, 3.0, 4.0], &[2, 2])
17        .expect("tensor data");
18    let results = graph
19        .run(&[Feed::new(&input, &input_data)], &[&stacked])
20        .expect("run");
21
22    println!("stacked tensor bytes: {}", results[0].byte_len().expect("byte len"));
23}

More examples

Hide additional examples

examples/03_arithmetic_topk.rs (line 20)

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}

examples/01_add_relu.rs (line 21)

4fn main() {
5    let device = MetalDevice::system_default().expect("no Metal device available");
6    let graph = Graph::new().expect("failed to create MPSGraph");
7
8    let input = graph
9        .placeholder(Some(&[2, 2]), data_type::FLOAT32, Some("input"))
10        .expect("failed to create placeholder");
11    let bias = graph
12        .constant_scalar(1.0, data_type::FLOAT32)
13        .expect("failed to create scalar constant");
14    let added = graph
15        .addition(&input, &bias, Some("add"))
16        .expect("failed to create addition op");
17    let output = graph
18        .relu(&added, Some("relu"))
19        .expect("failed to create relu op");
20
21    let input_data = TensorData::from_f32_slice(&device, &[1.0, -2.0, 3.0, -4.0], &[2, 2])
22        .expect("failed to create tensor data");
23    let results = graph
24        .run(&[Feed::new(&input, &input_data)], &[&output])
25        .expect("failed to execute graph");
26    let values = results[0].read_f32().expect("failed to read tensor output");
27
28    assert_eq!(values, vec![2.0, 0.0, 4.0, 0.0]);
29    println!("add+relu smoke passed: {values:?}");
30}

examples/02_compile_matmul.rs (line 32)

4fn main() {
5    let device = MetalDevice::system_default().expect("no Metal device available");
6    let queue = device
7        .new_command_queue()
8        .expect("failed to create command queue");
9    let graph = Graph::new().expect("failed to create MPSGraph");
10
11    let left = graph
12        .placeholder(Some(&[2, 3]), data_type::FLOAT32, Some("left"))
13        .expect("failed to create left placeholder");
14    let right = graph
15        .placeholder(Some(&[3, 2]), data_type::FLOAT32, Some("right"))
16        .expect("failed to create right placeholder");
17    let output = graph
18        .matrix_multiplication(&left, &right, Some("matmul"))
19        .expect("failed to create matrix multiplication op");
20
21    let executable = graph
22        .compile(
23            &device,
24            &[
25                FeedDescription::new(&left, &[2, 3], data_type::FLOAT32),
26                FeedDescription::new(&right, &[3, 2], data_type::FLOAT32),
27            ],
28            &[&output],
29        )
30        .expect("failed to compile executable");
31
32    let left_data = TensorData::from_f32_slice(&device, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3])
33        .expect("failed to create left tensor data");
34    let right_data =
35        TensorData::from_f32_slice(&device, &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0], &[3, 2])
36            .expect("failed to create right tensor data");
37
38    let results = executable
39        .run(&queue, &[&left_data, &right_data])
40        .expect("failed to run executable");
41    let values = results[0].read_f32().expect("failed to read tensor output");
42    let expected = [58.0_f32, 64.0, 139.0, 154.0];
43    for (actual, expected_value) in values.iter().zip(expected) {
44        assert!(
45            (actual - expected_value).abs() < 1.0e-4,
46            "unexpected matrix multiply result: {values:?}"
47        );
48    }
49
50    println!("compile+matmul smoke passed: {values:?}");
51}

examples/06_control_flow_call.rs (line 126)

16fn main() {
17    let device = MetalDevice::system_default().expect("no Metal device available");
18    let queue = device
19        .new_command_queue()
20        .expect("failed to create command queue");
21
22    let callee_graph = Graph::new().expect("callee graph");
23    let callee_input = callee_graph
24        .placeholder(Some(&[2]), data_type::FLOAT32, Some("callee_input"))
25        .expect("callee placeholder");
26    let callee_output = callee_graph
27        .addition(&callee_input, &callee_input, Some("callee_double"))
28        .expect("callee output");
29    let callee_executable = callee_graph
30        .compile(
31            &device,
32            &[FeedDescription::new(&callee_input, &[2], data_type::FLOAT32)],
33            &[&callee_output],
34        )
35        .expect("callee executable");
36
37    let graph = Graph::new().expect("graph");
38    let input = graph
39        .placeholder(Some(&[2]), data_type::FLOAT32, Some("input"))
40        .expect("input placeholder");
41    let predicate = graph
42        .placeholder(Some(&[]), data_type::BOOL, Some("predicate"))
43        .expect("predicate placeholder");
44    let bias = graph.constant_f32_slice(&[1.0, 1.0], &[2]).expect("bias constant");
45
46    let output_type = ShapedType::new(Some(&[2]), data_type::FLOAT32).expect("output type");
47    let call_results = graph
48        .call("double", &[&input], &[&output_type], Some("call"))
49        .expect("call op");
50    let if_results = graph
51        .if_then_else(
52            &predicate,
53            || vec![graph.addition(&input, &bias, None).expect("then add")],
54            || vec![graph.subtraction(&input, &bias, None).expect("else sub")],
55            Some("branch"),
56        )
57        .expect("if/then/else");
58
59    let call_operation = call_results[0].operation().expect("call operation");
60    let dependency = graph
61        .control_dependency(&[&call_operation], || {
62            vec![graph
63                .unary_arithmetic(UnaryArithmeticOp::Identity, &call_results[0], None)
64                .expect("identity")]
65        }, Some("dependency"))
66        .expect("control dependency");
67
68    let number_of_iterations = graph
69        .constant_scalar(4.0, data_type::INT32)
70        .expect("iteration count");
71    let zero = graph
72        .constant_scalar(0.0, data_type::INT32)
73        .expect("zero constant");
74    let one = graph.constant_scalar(1.0, data_type::INT32).expect("one constant");
75    let limit = graph
76        .constant_scalar(3.0, data_type::INT32)
77        .expect("limit constant");
78
79    let for_results = graph
80        .for_loop_iterations(&number_of_iterations, &[&zero], |_index, args| {
81            vec![graph.addition(&args[0], &one, None).expect("for-loop add")]
82        }, Some("for_loop"))
83        .expect("for loop");
84    let while_results = graph
85        .while_loop(
86            &[&zero],
87            |inputs| {
88                let condition = graph
89                    .binary_arithmetic(BinaryArithmeticOp::LessThan, &inputs[0], &limit, None)
90                    .expect("while predicate");
91                let passthrough = graph
92                    .unary_arithmetic(UnaryArithmeticOp::Identity, &inputs[0], None)
93                    .expect("while passthrough");
94                WhileBeforeResult {
95                    predicate: condition,
96                    results: vec![passthrough],
97                }
98            },
99            |inputs| vec![graph.addition(&inputs[0], &one, None).expect("while add")],
100            Some("while_loop"),
101        )
102        .expect("while loop");
103
104    let compile_descriptor = CompilationDescriptor::new().expect("compile descriptor");
105    compile_descriptor
106        .set_callable("double", Some(&callee_executable))
107        .expect("set callable");
108    let executable = graph
109        .compile_with_descriptor(
110            Some(&device),
111            &[
112                FeedDescription::new(&input, &[2], data_type::FLOAT32),
113                FeedDescription::new(&predicate, &[], data_type::BOOL),
114            ],
115            &[
116                &call_results[0],
117                &if_results[0],
118                &dependency[0],
119                &for_results[0],
120                &while_results[0],
121            ],
122            Some(&compile_descriptor),
123        )
124        .expect("compile executable");
125
126    let input_data = TensorData::from_f32_slice(&device, &[3.0, 4.0], &[2]).expect("input data");
127    let predicate_data = TensorData::from_bytes(&device, &[1_u8], &[], data_type::BOOL)
128        .expect("predicate data");
129    let results = executable
130        .run(&queue, &[&input_data, &predicate_data])
131        .expect("run executable");
132
133    println!("call output: {:?}", results[0].read_f32().expect("call output"));
134    println!("if output: {:?}", results[1].read_f32().expect("if output"));
135    println!("dependency output: {:?}", results[2].read_f32().expect("dependency output"));
136    println!("for output: {:?}", read_i32(&results[3]));
137    println!("while output: {:?}", read_i32(&results[4]));
138}

Source

pub fn from_buffer( buffer: &MetalBuffer, shape: &[usize], data_type: u32, ) -> Option<Self>

Alias an existing MTLBuffer as tensor data.

Source

pub fn byte_len(&self) -> Result<usize>

Examples found in repository ?

examples/05_concat_split.rs (line 22)

4fn main() {
5    let device = MetalDevice::system_default().expect("no Metal device available");
6    let graph = Graph::new().expect("graph");
7    let input = graph
8        .placeholder(Some(&[2, 2]), data_type::FLOAT32, Some("input"))
9        .expect("placeholder");
10    let concat = graph
11        .concat_pair(&input, &input, 1, Some("concat"))
12        .expect("concat");
13    let split = graph.split_num(&concat, 2, 1, Some("split"));
14    let stacked = graph.stack(&[&split[0], &split[1]], 0, Some("stack")).expect("stack");
15
16    let input_data = TensorData::from_f32_slice(&device, &[1.0, 2.0, 3.0, 4.0], &[2, 2])
17        .expect("tensor data");
18    let results = graph
19        .run(&[Feed::new(&input, &input_data)], &[&stacked])
20        .expect("run");
21
22    println!("stacked tensor bytes: {}", results[0].byte_len().expect("byte len"));
23}

Source

pub fn read_bytes(&self) -> Result<Vec<u8>>

Examples found in repository ?

examples/06_control_flow_call.rs (line 10)

9fn read_i32(data: &TensorData) -> Vec<i32> {
10    let bytes = data.read_bytes().expect("read bytes");
11    bytes.chunks_exact(core::mem::size_of::<i32>())
12        .map(|chunk| i32::from_ne_bytes(chunk.try_into().expect("i32 chunk")))
13        .collect()
14}

Source

pub fn read_f32(&self) -> Result<Vec<f32>>

Examples found in repository ?

examples/03_arithmetic_topk.rs (line 26)

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}

More examples

Hide additional examples

examples/01_add_relu.rs (line 26)

4fn main() {
5    let device = MetalDevice::system_default().expect("no Metal device available");
6    let graph = Graph::new().expect("failed to create MPSGraph");
7
8    let input = graph
9        .placeholder(Some(&[2, 2]), data_type::FLOAT32, Some("input"))
10        .expect("failed to create placeholder");
11    let bias = graph
12        .constant_scalar(1.0, data_type::FLOAT32)
13        .expect("failed to create scalar constant");
14    let added = graph
15        .addition(&input, &bias, Some("add"))
16        .expect("failed to create addition op");
17    let output = graph
18        .relu(&added, Some("relu"))
19        .expect("failed to create relu op");
20
21    let input_data = TensorData::from_f32_slice(&device, &[1.0, -2.0, 3.0, -4.0], &[2, 2])
22        .expect("failed to create tensor data");
23    let results = graph
24        .run(&[Feed::new(&input, &input_data)], &[&output])
25        .expect("failed to execute graph");
26    let values = results[0].read_f32().expect("failed to read tensor output");
27
28    assert_eq!(values, vec![2.0, 0.0, 4.0, 0.0]);
29    println!("add+relu smoke passed: {values:?}");
30}

examples/02_compile_matmul.rs (line 41)

4fn main() {
5    let device = MetalDevice::system_default().expect("no Metal device available");
6    let queue = device
7        .new_command_queue()
8        .expect("failed to create command queue");
9    let graph = Graph::new().expect("failed to create MPSGraph");
10
11    let left = graph
12        .placeholder(Some(&[2, 3]), data_type::FLOAT32, Some("left"))
13        .expect("failed to create left placeholder");
14    let right = graph
15        .placeholder(Some(&[3, 2]), data_type::FLOAT32, Some("right"))
16        .expect("failed to create right placeholder");
17    let output = graph
18        .matrix_multiplication(&left, &right, Some("matmul"))
19        .expect("failed to create matrix multiplication op");
20
21    let executable = graph
22        .compile(
23            &device,
24            &[
25                FeedDescription::new(&left, &[2, 3], data_type::FLOAT32),
26                FeedDescription::new(&right, &[3, 2], data_type::FLOAT32),
27            ],
28            &[&output],
29        )
30        .expect("failed to compile executable");
31
32    let left_data = TensorData::from_f32_slice(&device, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3])
33        .expect("failed to create left tensor data");
34    let right_data =
35        TensorData::from_f32_slice(&device, &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0], &[3, 2])
36            .expect("failed to create right tensor data");
37
38    let results = executable
39        .run(&queue, &[&left_data, &right_data])
40        .expect("failed to run executable");
41    let values = results[0].read_f32().expect("failed to read tensor output");
42    let expected = [58.0_f32, 64.0, 139.0, 154.0];
43    for (actual, expected_value) in values.iter().zip(expected) {
44        assert!(
45            (actual - expected_value).abs() < 1.0e-4,
46            "unexpected matrix multiply result: {values:?}"
47        );
48    }
49
50    println!("compile+matmul smoke passed: {values:?}");
51}

examples/06_control_flow_call.rs (line 133)

16fn main() {
17    let device = MetalDevice::system_default().expect("no Metal device available");
18    let queue = device
19        .new_command_queue()
20        .expect("failed to create command queue");
21
22    let callee_graph = Graph::new().expect("callee graph");
23    let callee_input = callee_graph
24        .placeholder(Some(&[2]), data_type::FLOAT32, Some("callee_input"))
25        .expect("callee placeholder");
26    let callee_output = callee_graph
27        .addition(&callee_input, &callee_input, Some("callee_double"))
28        .expect("callee output");
29    let callee_executable = callee_graph
30        .compile(
31            &device,
32            &[FeedDescription::new(&callee_input, &[2], data_type::FLOAT32)],
33            &[&callee_output],
34        )
35        .expect("callee executable");
36
37    let graph = Graph::new().expect("graph");
38    let input = graph
39        .placeholder(Some(&[2]), data_type::FLOAT32, Some("input"))
40        .expect("input placeholder");
41    let predicate = graph
42        .placeholder(Some(&[]), data_type::BOOL, Some("predicate"))
43        .expect("predicate placeholder");
44    let bias = graph.constant_f32_slice(&[1.0, 1.0], &[2]).expect("bias constant");
45
46    let output_type = ShapedType::new(Some(&[2]), data_type::FLOAT32).expect("output type");
47    let call_results = graph
48        .call("double", &[&input], &[&output_type], Some("call"))
49        .expect("call op");
50    let if_results = graph
51        .if_then_else(
52            &predicate,
53            || vec![graph.addition(&input, &bias, None).expect("then add")],
54            || vec![graph.subtraction(&input, &bias, None).expect("else sub")],
55            Some("branch"),
56        )
57        .expect("if/then/else");
58
59    let call_operation = call_results[0].operation().expect("call operation");
60    let dependency = graph
61        .control_dependency(&[&call_operation], || {
62            vec![graph
63                .unary_arithmetic(UnaryArithmeticOp::Identity, &call_results[0], None)
64                .expect("identity")]
65        }, Some("dependency"))
66        .expect("control dependency");
67
68    let number_of_iterations = graph
69        .constant_scalar(4.0, data_type::INT32)
70        .expect("iteration count");
71    let zero = graph
72        .constant_scalar(0.0, data_type::INT32)
73        .expect("zero constant");
74    let one = graph.constant_scalar(1.0, data_type::INT32).expect("one constant");
75    let limit = graph
76        .constant_scalar(3.0, data_type::INT32)
77        .expect("limit constant");
78
79    let for_results = graph
80        .for_loop_iterations(&number_of_iterations, &[&zero], |_index, args| {
81            vec![graph.addition(&args[0], &one, None).expect("for-loop add")]
82        }, Some("for_loop"))
83        .expect("for loop");
84    let while_results = graph
85        .while_loop(
86            &[&zero],
87            |inputs| {
88                let condition = graph
89                    .binary_arithmetic(BinaryArithmeticOp::LessThan, &inputs[0], &limit, None)
90                    .expect("while predicate");
91                let passthrough = graph
92                    .unary_arithmetic(UnaryArithmeticOp::Identity, &inputs[0], None)
93                    .expect("while passthrough");
94                WhileBeforeResult {
95                    predicate: condition,
96                    results: vec![passthrough],
97                }
98            },
99            |inputs| vec![graph.addition(&inputs[0], &one, None).expect("while add")],
100            Some("while_loop"),
101        )
102        .expect("while loop");
103
104    let compile_descriptor = CompilationDescriptor::new().expect("compile descriptor");
105    compile_descriptor
106        .set_callable("double", Some(&callee_executable))
107        .expect("set callable");
108    let executable = graph
109        .compile_with_descriptor(
110            Some(&device),
111            &[
112                FeedDescription::new(&input, &[2], data_type::FLOAT32),
113                FeedDescription::new(&predicate, &[], data_type::BOOL),
114            ],
115            &[
116                &call_results[0],
117                &if_results[0],
118                &dependency[0],
119                &for_results[0],
120                &while_results[0],
121            ],
122            Some(&compile_descriptor),
123        )
124        .expect("compile executable");
125
126    let input_data = TensorData::from_f32_slice(&device, &[3.0, 4.0], &[2]).expect("input data");
127    let predicate_data = TensorData::from_bytes(&device, &[1_u8], &[], data_type::BOOL)
128        .expect("predicate data");
129    let results = executable
130        .run(&queue, &[&input_data, &predicate_data])
131        .expect("run executable");
132
133    println!("call output: {:?}", results[0].read_f32().expect("call output"));
134    println!("if output: {:?}", results[1].read_f32().expect("if output"));
135    println!("dependency output: {:?}", results[2].read_f32().expect("dependency output"));
136    println!("for output: {:?}", read_i32(&results[3]));
137    println!("while output: {:?}", read_i32(&results[4]));
138}

examples/07_gather_random_rnn.rs (line 150)

15fn main() {
16    let graph = Graph::new().expect("graph");
17    let updates = graph
18        .constant_f32_slice(&[10.0, 20.0, 30.0, 40.0, 50.0, 60.0], &[2, 3])
19        .expect("updates");
20    let gather_indices = graph
21        .constant_bytes(&i32_bytes(&[2, 0]), &[2], data_type::INT32)
22        .expect("gather indices");
23    let gather_nd_indices = graph
24        .constant_bytes(&i32_bytes(&[0, 1, 1, 0]), &[2, 2], data_type::INT32)
25        .expect("gather nd indices");
26    let along_indices = graph
27        .constant_bytes(&i32_bytes(&[2, 1, 0, 0, 1, 2]), &[2, 3], data_type::INT32)
28        .expect("gather along indices");
29    let axis_tensor = graph
30        .constant_scalar(1.0, data_type::INT32)
31        .expect("axis tensor");
32
33    let gather = graph
34        .gather(&updates, &gather_indices, 1, 0, Some("gather"))
35        .expect("gather");
36    let gather_nd = graph
37        .gather_nd(&updates, &gather_nd_indices, 0, Some("gather_nd"))
38        .expect("gather nd");
39    let gather_axis = graph
40        .gather_along_axis(1, &updates, &along_indices, Some("gather_axis"))
41        .expect("gather along axis");
42    let gather_axis_tensor = graph
43        .gather_along_axis_tensor(&axis_tensor, &updates, &along_indices, Some("gather_axis_tensor"))
44        .expect("gather along axis tensor");
45
46    let descriptor = RandomOpDescriptor::new(random_distribution::UNIFORM, data_type::FLOAT32)
47        .expect("random descriptor");
48    descriptor.set_min(0.0).expect("random min");
49    descriptor.set_max(1.0).expect("random max");
50    let random = graph
51        .random_tensor_seed(&[4], &descriptor, 7, Some("random"))
52        .expect("random tensor");
53    let dropout = graph.dropout(&updates, 1.0, Some("dropout")).expect("dropout");
54
55    let single_gate_descriptor = SingleGateRNNDescriptor::new().expect("single gate descriptor");
56    single_gate_descriptor
57        .set_activation(rnn_activation::RELU)
58        .expect("single gate activation");
59    let single_gate_source = graph
60        .constant_f32_slice(&[0.5], &[1, 1, 1])
61        .expect("single gate source");
62    let single_gate_recurrent = graph
63        .constant_f32_slice(&[0.0], &[1, 1])
64        .expect("single gate recurrent");
65    let single_gate = graph
66        .single_gate_rnn(
67            &single_gate_source,
68            &single_gate_recurrent,
69            None,
70            None,
71            None,
72            None,
73            &single_gate_descriptor,
74            Some("single_gate"),
75        )
76        .expect("single gate rnn");
77
78    let lstm_descriptor = LSTMDescriptor::new().expect("lstm descriptor");
79    lstm_descriptor
80        .set_produce_cell(true)
81        .expect("set produce cell");
82    let lstm_source = graph
83        .constant_f32_slice(&[0.0; 4], &[1, 1, 4])
84        .expect("lstm source");
85    let lstm_recurrent = graph
86        .constant_f32_slice(&[0.0; 4], &[4, 1])
87        .expect("lstm recurrent");
88    let lstm = graph
89        .lstm(
90            &lstm_source,
91            &lstm_recurrent,
92            None,
93            None,
94            None,
95            None,
96            None,
97            None,
98            &lstm_descriptor,
99            Some("lstm"),
100        )
101        .expect("lstm");
102
103    let gru_descriptor = GRUDescriptor::new().expect("gru descriptor");
104    gru_descriptor.set_training(true).expect("set gru training");
105    gru_descriptor
106        .set_reset_after(true)
107        .expect("set gru reset_after");
108    let gru_source = graph
109        .constant_f32_slice(&[0.0; 3], &[1, 1, 3])
110        .expect("gru source");
111    let gru_recurrent = graph
112        .constant_f32_slice(&[0.0; 3], &[3, 1])
113        .expect("gru recurrent");
114    let gru_secondary_bias = graph
115        .constant_f32_slice(&[0.0], &[1])
116        .expect("gru secondary bias");
117    let gru = graph
118        .gru(
119            &gru_source,
120            &gru_recurrent,
121            None,
122            None,
123            None,
124            None,
125            Some(&gru_secondary_bias),
126            &gru_descriptor,
127            Some("gru"),
128        )
129        .expect("gru");
130
131    let results = graph
132        .run(
133            &[],
134            &[
135                &gather,
136                &gather_nd,
137                &gather_axis,
138                &gather_axis_tensor,
139                &random,
140                &dropout,
141                &single_gate[0],
142                &lstm[0],
143                &lstm[1],
144                &gru[0],
145                &gru[1],
146            ],
147        )
148        .expect("run graph");
149
150    println!("gather: {:?}", results[0].read_f32().expect("gather"));
151    println!("gather_nd: {:?}", results[1].read_f32().expect("gather_nd"));
152    println!("gather_axis: {:?}", results[2].read_f32().expect("gather_axis"));
153    println!(
154        "gather_axis_tensor: {:?}",
155        results[3].read_f32().expect("gather_axis_tensor")
156    );
157    println!("random: {:?}", results[4].read_f32().expect("random"));
158    println!("dropout: {:?}", results[5].read_f32().expect("dropout"));
159    println!("single_gate: {:?}", results[6].read_f32().expect("single_gate"));
160    println!("lstm state: {:?}", results[7].read_f32().expect("lstm state"));
161    println!("lstm cell: {:?}", results[8].read_f32().expect("lstm cell"));
162    println!("gru state: {:?}", results[9].read_f32().expect("gru state"));
163    println!("gru training: {:?}", results[10].read_f32().expect("gru training"));
164}