Struct ExtensionShapeContext

Source

pub struct ExtensionShapeContext<'a> { /* private fields */ }

Expand description

Input metadata and equality requirements for one extension inference call.

The context records requirements declaratively. It does not prove, reject, normalize, or solve them.

§Examples

use tenferro_ops::ExtensionShapeContext;

fn infer(ctx: &mut ExtensionShapeContext<'_>) -> tenferro_tensor::Result<()> {
    let lhs = ctx.input_axis(0, 0)?;
    let rhs = ctx.input_axis(1, 0)?;
    ctx.require_equal(lhs, 2 * rhs)?;
    Ok(())
}

Implementations§

Source §

impl<'a> ExtensionShapeContext<'a>

Source

pub fn input_dtype(&self, input: usize) -> Result<DType, ExtensionShapeError>

Return the dtype of one extension input.

§Examples

use tenferro_ops::ExtensionShapeContext;
use tenferro_tensor::DType;

fn infer(ctx: &ExtensionShapeContext<'_>) -> tenferro_tensor::Result<DType> {
    Ok(ctx.input_dtype(0)?)
}

§Errors

Returns ExtensionShapeError::InputOutOfBounds when input is not present in the extension metadata.

Source

pub fn input_shape( &self, input: usize, ) -> Result<&[SymDim], ExtensionShapeError>

Return the symbolic shape of one extension input.

§Examples

use tenferro_ops::ExtensionShapeContext;

fn infer(ctx: &ExtensionShapeContext<'_>) -> tenferro_tensor::Result<()> {
    let _shape = ctx.input_shape(0)?;
    Ok(())
}

§Errors

Returns ExtensionShapeError::InputOutOfBounds when input is not present in the extension metadata.

Source

pub fn input_axis( &self, input: usize, axis: usize, ) -> Result<SymDim, ExtensionShapeError>

Return one symbolic axis expression from an extension input.

§Examples

use tenferro_ops::ExtensionShapeContext;
use tenferro_ops::SymDim;

fn infer(ctx: &ExtensionShapeContext<'_>) -> tenferro_tensor::Result<SymDim> {
    Ok(ctx.input_axis(0, 0)?)
}

§Errors

Returns ExtensionShapeError::InputOutOfBounds for an unknown input or ExtensionShapeError::AxisOutOfBounds for an axis outside that input’s rank.

Source

pub fn require_equal( &mut self, lhs: SymDim, rhs: SymDim, ) -> Result<(), ExtensionShapeError>

Record equality of two symbolic dimension expressions.

This method records the expressions without trying to solve them.

§Examples

use tenferro_ops::ExtensionShapeContext;

fn infer(ctx: &mut ExtensionShapeContext<'_>) -> tenferro_tensor::Result<()> {
    let lhs = ctx.input_axis(0, 0)?;
    let rhs = ctx.input_axis(1, 0)?;
    ctx.require_equal(lhs, 2 * rhs)?;
    Ok(())
}

§Errors

Returns Ok(()); this method only records a symbolic equality and does not evaluate it or produce an ExtensionShapeError.

Source

pub fn require_axes_equal( &mut self, lhs: (usize, usize), rhs: (usize, usize), ) -> Result<(), ExtensionShapeError>

Record equality of two input axes.

§Examples

use tenferro_ops::ExtensionShapeContext;

fn infer(ctx: &mut ExtensionShapeContext<'_>) -> tenferro_tensor::Result<()> {
    ctx.require_axes_equal((0, 0), (1, 0))?;
    Ok(())
}

§Errors

Returns ExtensionShapeError::InputOutOfBounds or ExtensionShapeError::AxisOutOfBounds when either referenced axis is absent from the extension metadata.

Source

pub fn require_same_shape( &mut self, lhs_input: usize, rhs_input: usize, ) -> Result<(), ExtensionShapeError>

Require two extension inputs to have the same rank and axis extents.

A rank mismatch is returned before any equality is recorded.

§Examples

use tenferro_ops::ExtensionShapeContext;

fn infer(ctx: &mut ExtensionShapeContext<'_>) -> tenferro_tensor::Result<()> {
    ctx.require_same_shape(0, 1)?;
    Ok(())
}