Struct XlaExecutor 

pub struct XlaExecutor { /* private fields */ }

Experimental peer executor for XLA/PJRT.

Examples

use tenferro_runtime::{GraphCompiler, TracedTensor};
use tenferro_xla::XlaExecutor;

let x = TracedTensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
let mut compiler = GraphCompiler::new();
let program = compiler.compile(&x.neg()).unwrap();
let module = XlaExecutor::default().lower_to_stablehlo(&program).unwrap();
assert!(module.as_str().contains("stablehlo.negate"));

Implementations

impl XlaExecutor

pub fn new(options: XlaExecutorOptions) -> Self

Create an executor with explicit options.

Examples

use tenferro_xla::{XlaExecutor, XlaExecutorOptions};

let executor = XlaExecutor::new(XlaExecutorOptions::default());
assert_eq!(executor.options(), XlaExecutorOptions::default());

pub fn from_env() -> Result<Self>

Create an executor by loading PJRT configuration from environment variables.

Examples

use tenferro_xla::XlaExecutor;

let _ = XlaExecutor::from_env();

pub fn from_env_var(var: &'static str) -> Result<Self>

Create an executor by loading a PJRT plugin path from a specific environment variable.

Examples

use tenferro_xla::XlaExecutor;

let _ = XlaExecutor::from_env_var("__TENFERRO_XLA_DOCS_UNSET");

pub fn options(&self) -> XlaExecutorOptions

Return the executor options.

Examples

use tenferro_xla::XlaExecutor;

assert_eq!(XlaExecutor::default().options(), Default::default());

pub fn has_loaded_pjrt_plugin(&self) -> bool

Return whether this executor owns a loaded PJRT plugin.

Examples

use tenferro_xla::XlaExecutor;

assert!(!XlaExecutor::default().has_loaded_pjrt_plugin());

pub fn lower_to_stablehlo( &self, program: &GraphProgram, ) -> Result<StableHloModule>

Lower a graph program to StableHLO without executing it.

Examples

use tenferro_runtime::{GraphCompiler, TracedTensor};
use tenferro_xla::XlaExecutor;

let x = TracedTensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
let mut compiler = GraphCompiler::new();
let program = compiler.compile(&x.neg()).unwrap();
let module = XlaExecutor::default().lower_to_stablehlo(&program).unwrap();
assert!(module.as_str().contains("stablehlo.negate"));

pub fn run_many_with_inputs( &self, program: &GraphProgram, inputs: &[&Tensor], ) -> Result<Vec<Tensor>>

Execute a graph program through a loaded PJRT plugin and return all outputs.

Inputs must match [GraphProgram::input_specs] exactly. This experimental execution path supports the same exact-static-shape, F32/F64 subset as StableHLO lowering.

Examples

use tenferro_runtime::{GraphCompiler, TracedTensor};
use tenferro_tensor::Tensor;
use tenferro_xla::{Error, XlaExecutor};

let x = TracedTensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
let mut compiler = GraphCompiler::new();
let program = compiler.compile(&x.neg()).unwrap();
let input = Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
let err = XlaExecutor::default()
    .run_many_with_inputs(&program, &[&input])
    .unwrap_err();
assert!(matches!(err, Error::PjrtFeatureDisabled | Error::PjrtPluginNotLoaded));

pub fn run_with_inputs( &self, program: &GraphProgram, inputs: &[&Tensor], ) -> Result<Tensor>

Execute a single-output graph program through a loaded PJRT plugin.

Examples

use tenferro_runtime::{GraphCompiler, TracedTensor};
use tenferro_tensor::Tensor;
use tenferro_xla::{Error, XlaExecutor};

let x = TracedTensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
let mut compiler = GraphCompiler::new();
let program = compiler.compile(&x.neg()).unwrap();
let input = Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
let err = XlaExecutor::default().run_with_inputs(&program, &[&input]).unwrap_err();
assert!(matches!(err, Error::PjrtFeatureDisabled | Error::PjrtPluginNotLoaded));

Trait Implementations

impl Debug for XlaExecutor

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Default for XlaExecutor

fn default() -> Self

Returns the “default value” for a type.