xla 0.1.6

Bindings for the XLA C++ library.
Documentation 
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

use anyhow::Result;
extern crate xla;
use xla::ArrayElement;

fn main() -> Result<()> {
    let client = xla::PjRtClient::cpu()?;
    loop {
        let builder = xla::XlaBuilder::new("test");
        let x = builder.parameter(0, f32::TY, &[2], "x")?;
        let sum = x.reduce_sum(&[], false)?.build()?.compile(&client)?;
        let input = xla::Literal::vec1(&[4.2f32, 1.337f32]);
        let result = sum.execute::<xla::Literal>(&[input])?;
        println!("1");
        let result = result[0][0].to_literal_sync()?;
        drop(sum);
        assert_eq!(result.to_vec::<f32>()?, [4.2, 1.337]);

        let builder = xla::XlaBuilder::new("test");
        let x = builder.parameter(0, f32::TY, &[-2], "x")?;
        let sum = x.reduce_sum(&[0], false)?.build()?.compile(&client)?;
        let input = xla::Literal::vec1(&[4.2f32, 1.337f32]);
        let result = sum.execute::<xla::Literal>(&[input])?;
        println!("2");
        let result = result[0][0].to_literal_sync()?;
        drop(sum);
        assert_eq!(result.to_vec::<f32>()?, [5.5369997]);
        // Dimensions got reduced.
        assert_eq!(result.array_shape()?.dims(), []);

        let builder = xla::XlaBuilder::new("test");
        let x = builder.parameter(0, f32::TY, &[-2], "x")?;
        let sum = x.reduce_sum(&[0], true)?.build()?.compile(&client)?;
        let input = xla::Literal::vec1(&[4.2f32, 1.337f32]);
        let result = sum.execute::<xla::Literal>(&[input])?;
        println!("3");
        let result = result[0][0].to_literal_sync()?;
        drop(sum);
        assert_eq!(result.to_vec::<f32>()?, [5.5369997]);
        // keep_dims = true in this case.
        assert_eq!(result.array_shape()?.dims(), [1]);
        println!("Done!");
    }
}