Struct TPUOptimizer 

pub struct TPUOptimizer<O, A>
where
    A: Float + ScalarOperand + Debug,
    O: Optimizer<A, Ix1>,
{ /* private fields */ }

TPU-optimized optimizer wrapper

Implementations

impl<O, A> TPUOptimizer<O, A>

where A: Float + Default + Clone + Send + Sync + ScalarOperand + Debug, O: Optimizer<A, Ix1> + Send + Sync,

pub fn new(base_optimizer: O, config: TPUConfig) -> Result<Self>

Create a new TPU optimizer

pub fn initialize_xla_graph(&mut self) -> Result<()>

Initialize XLA computation graph

pub fn compile_step(&mut self, inputshapes: &[XLAShape]) -> Result<String>

Compile optimizer step for TPU execution

pub fn tpu_step<S, DIM>( &mut self, params: &ArrayBase<S, DIM>, gradients: &ArrayBase<S, DIM>, ) -> Result<Array<A, DIM>>

where S: Data<Elem = A>, DIM: Dimension + Clone,

Execute TPU-optimized step

pub fn get_performance_metrics(&self) -> TPUPerformanceMetrics

Get TPU performance metrics

pub fn optimize_memory_layout(&mut self) -> Result<()>

Optimize TPU memory layout

pub fn get_topology_info(&self) -> TPUTopologyInfo

Get TPU topology information