Struct PreparedDiagonalPlan 

pub struct PreparedDiagonalPlan { /* private fields */ }

Caller-owned, device-bound indices for repeated-axis diagonal extraction.

Axis ids represent the input side of an einsum equation: equal ids are the repeated labels whose diagonal is retained, and first occurrence order determines the output axes. Preparing once avoids rebuilding and uploading the same u32 index tensor for repeated calls with the same shape.

Implementations

impl PreparedDiagonalPlan

pub fn new( input_shape: &[usize], axis_ids: &[usize], device: &Device, ) -> Result<Self>

Prepares reusable diagonal indices for one exact shape and device.

pub fn repeated( extent: usize, multiplicity: usize, device: &Device, ) -> Result<Self>

Convenience constructor for i i ... -> i extraction.

pub fn interleaved( first_extent: usize, second_extent: usize, device: &Device, ) -> Result<Self>

Convenience constructor for i j i j -> i j extraction.

pub fn execute(&self, input: &Tensor) -> Result<Tensor>

Extracts the prepared diagonal from one contiguous tensor.

pub fn input_shape(&self) -> &[usize]

Exact input shape captured by this plan.

pub fn axis_ids(&self) -> &[usize]

Axis ids captured from the input-side equation.

pub fn output_shape(&self) -> &[usize]

Unique-axis output shape in first-occurrence order.

pub const fn index_dtype(&self) -> DType

Device index representation used by the prepared gather.

pub fn indices(&self) -> &Tensor

Prepared device-local indices, exposed for diagnostics and benchmarking.

Trait Implementations

impl Debug for PreparedDiagonalPlan

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

Formats the value using the given formatter.