Struct ExtensionRuleSet 

pub struct ExtensionRuleSet { /* private fields */ }

Explicit, owned set of extension AD rules.

This is the rule container used by higher-level AD contexts. Extension AD intentionally has no process-global fallback; callers must pass the rule set that their graph needs.

Implementations

impl ExtensionRuleSet

pub fn new() -> ExtensionRuleSet

Create an empty extension rule set.

Examples

use tenferro_ops::ExtensionRuleSet;

let rules = ExtensionRuleSet::new();
assert!(!rules.is_linearize_registered("example.missing.v1"));

pub fn register_linearize( &mut self, rule: Arc<dyn ExtensionLinearizeRule>, ) -> Result<(), ExtensionRegistryError>

Add one linearize rule to this owned set.

Examples

use std::sync::Arc;
use tidu::ADRuleResult;
use computegraph::types::{LocalValueId, ValueKey};
use tenferro_ops::ad::PrimitiveRuleBuilder;
use tenferro_ops::ext_op::{ExtensionLinearizeRule, ExtensionOp};
use tenferro_ops::{ExtensionRuleSet, ShapeGuardContext};
use tenferro_ops::std_tensor_op::StdTensorOp;

#[derive(Debug)]
struct Rule;

impl ExtensionLinearizeRule for Rule {
    fn family_id(&self) -> &'static str { "example.register_linearize.v1" }
    fn linearize(
        &self,
        _op: &dyn ExtensionOp,
        _builder: &mut dyn PrimitiveRuleBuilder,
        _primal_in: &[ValueKey<StdTensorOp>],
        _primal_out: &[ValueKey<StdTensorOp>],
        tangent_in: &[Option<LocalValueId>],
        _ctx: &mut ShapeGuardContext,
    ) -> ADRuleResult<Vec<Option<LocalValueId>>> {
        Ok(tangent_in.to_vec())
    }
}

let mut rules = ExtensionRuleSet::new();
rules.register_linearize(Arc::new(Rule)).unwrap();
assert!(rules.lookup_linearize("example.register_linearize.v1").is_some());

pub fn register_linear_transpose( &mut self, rule: Arc<dyn ExtensionLinearTransposeRule>, ) -> Result<(), ExtensionRegistryError>

Add one linear-transpose rule to this owned set.

pub fn register_primal_vjp( &mut self, rule: Arc<dyn ExtensionPrimalVjpRule>, ) -> Result<(), ExtensionRegistryError>

Add one primal-VJP rule to this owned set.

pub fn with_linearize( self, rule: Arc<dyn ExtensionLinearizeRule>, ) -> Result<ExtensionRuleSet, ExtensionRegistryError>

Return a new rule set containing a linearize rule.

Examples

use std::sync::Arc;
use tidu::ADRuleResult;
use computegraph::types::{LocalValueId, ValueKey};
use tenferro_ops::ad::PrimitiveRuleBuilder;
use tenferro_ops::ext_op::{ExtensionLinearizeRule, ExtensionOp};
use tenferro_ops::{ExtensionRuleSet, ShapeGuardContext};
use tenferro_ops::std_tensor_op::StdTensorOp;

#[derive(Debug)]
struct Rule;

impl ExtensionLinearizeRule for Rule {
    fn family_id(&self) -> &'static str { "example.with_linearize.v1" }
    fn linearize(
        &self,
        _op: &dyn ExtensionOp,
        _builder: &mut dyn PrimitiveRuleBuilder,
        _primal_in: &[ValueKey<StdTensorOp>],
        _primal_out: &[ValueKey<StdTensorOp>],
        tangent_in: &[Option<LocalValueId>],
        _ctx: &mut ShapeGuardContext,
    ) -> ADRuleResult<Vec<Option<LocalValueId>>> {
        Ok(tangent_in.to_vec())
    }
}

let rules = ExtensionRuleSet::new().with_linearize(Arc::new(Rule)).unwrap();
assert!(rules.is_linearize_registered("example.with_linearize.v1"));

pub fn with_linear_transpose( self, rule: Arc<dyn ExtensionLinearTransposeRule>, ) -> Result<ExtensionRuleSet, ExtensionRegistryError>

Return a new rule set containing a linear-transpose rule.

pub fn with_primal_vjp( self, rule: Arc<dyn ExtensionPrimalVjpRule>, ) -> Result<ExtensionRuleSet, ExtensionRegistryError>

Return a new rule set containing a primal-VJP rule.

pub fn merge( &mut self, other: ExtensionRuleSet, ) -> Result<(), ExtensionRegistryError>

Merge another owned rule set into this one.

The merge is atomic: if any rule in other is invalid or duplicates an existing family, self is left unchanged.

Examples

use tenferro_ops::ExtensionRuleSet;

let mut rules = ExtensionRuleSet::new();
rules.merge(ExtensionRuleSet::new()).unwrap();
assert!(!rules.is_linearize_registered("example.missing.v1"));

pub fn lookup_linearize( &self, family_id: &str, ) -> Option<Arc<dyn ExtensionLinearizeRule>>

Look up a linearize rule in this set.

Examples

use tenferro_ops::ExtensionRuleSet;

let rules = ExtensionRuleSet::new();
assert!(rules.lookup_linearize("example.missing.v1").is_none());

pub fn lookup_linear_transpose( &self, family_id: &str, ) -> Option<Arc<dyn ExtensionLinearTransposeRule>>

Look up a linear-transpose rule in this set.

pub fn lookup_primal_vjp( &self, family_id: &str, ) -> Option<Arc<dyn ExtensionPrimalVjpRule>>

Look up a primal-VJP rule in this set.

pub fn is_linearize_registered(&self, family_id: &str) -> bool

Return whether family_id has a linearize rule.

Examples

use tenferro_ops::ExtensionRuleSet;

let rules = ExtensionRuleSet::new();
assert!(!rules.is_linearize_registered("example.missing.v1"));

pub fn is_linear_transpose_registered(&self, family_id: &str) -> bool

Return whether family_id has a linear-transpose rule.

pub fn is_primal_vjp_registered(&self, family_id: &str) -> bool

Return whether family_id has a primal-VJP rule.

Trait Implementations

impl Clone for ExtensionRuleSet

fn clone(&self) -> ExtensionRuleSet

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for ExtensionRuleSet

Available on crate feature autodiff only.

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

Formats the value using the given formatter.

impl Default for ExtensionRuleSet

fn default() -> ExtensionRuleSet

Returns the “default value” for a type.