Runtime-agnostic handle over a NEML2 v3 model. More...

#include <NEML2ModelHandle.h>

Inheritance diagram for NEML2ModelHandle:

Public Member Functions
virtual	~NEML2ModelHandle ()=default

virtual at::Device	device () const =0
	The compute device this rank's workload runs on (for cpp-aoti, the scheduler-assigned device).

Introspection, in graph-call order (params keyed by qualified name)
virtual const std::vector< std::string > &	input_names () const =0

virtual const std::vector< std::string > &	output_names () const =0

virtual const std::vector< std::vector< int64_t > > &	input_base_shapes () const =0

virtual const std::vector< std::vector< int64_t > > &	output_base_shapes () const =0

virtual const std::map< std::string, std::vector< int64_t > > &	parameter_base_shapes () const =0

Evaluation
Evaluate the model only (outputs, no derivatives). Used on pure residual evaluations – line search, damping, JFNK residuals – to skip the (comparatively expensive) derivative.
virtual std::map< std::string, at::Tensor >	value (const std::map< std::string, at::Tensor > &inputs) const =0

virtual std::pair< std::map< std::string, at::Tensor >, neml2::aoti::VariablePairJacobian >	jacobian (const std::map< std::string, at::Tensor > &inputs) const =0
	Evaluate the model and its full input Jacobian.

virtual std::pair< std::map< std::string, at::Tensor >, neml2::aoti::VariablePairJacobian >	param_jacobian (const std::map< std::string, at::Tensor > &inputs) const =0
	Evaluate the model and its parameter Jacobian (d output / d parameter).

virtual void	set_parameter (const std::string &name, const at::Tensor &value)=0
	Replace a (runtime-flexible) model parameter's value.

Detailed Description

Runtime-agnostic handle over a NEML2 v3 model.

NEML2 v3 ships two C++ runtimes with parallel APIs but no common base class: neml2::eager::Model (embeds CPython, loads a model from a source .i) and the cpp-aoti neml2::aoti::DispatchedModel (loads ahead-of-time-compiled artifacts – a shared metadata.json plus per-<device>/<dtype>/ .pt2 binaries – straight from the artifact folder, Python-free, dispatched to a device by a scheduler). This handle exposes the single surface the MOOSE integration needs – introspection (names + base shapes) and evaluation (jacobian / param_jacobian / set_parameter) – so both the setup-time introspection in NEML2Action and the runtime evaluation in NEML2ModelExecutor use one consistent API regardless of which runtime backs the model.

Construct one with makeNEML2ModelHandle.

Definition at line 45 of file NEML2ModelHandle.h.

Constructor & Destructor Documentation

◆ ~NEML2ModelHandle()

virtual NEML2ModelHandle::~NEML2ModelHandle ( )

virtualdefault

Member Function Documentation

◆ device()

virtual at::Device NEML2ModelHandle::device ( ) const

pure virtual

The compute device this rank's workload runs on (for cpp-aoti, the scheduler-assigned device).

Implemented in EagerModelHandle, and AOTIModelHandle.

◆ input_base_shapes()

virtual const std::vector< std::vector< int64_t > > & NEML2ModelHandle::input_base_shapes ( ) const

pure virtual

Implemented in EagerModelHandle, and AOTIModelHandle.

Referenced by NEML2Action::setupDerivativeMappings(), and NEML2Action::setupInputMappings().

◆ input_names()

virtual const std::vector< std::string > & NEML2ModelHandle::input_names ( ) const

pure virtual

Implemented in EagerModelHandle, and AOTIModelHandle.

Referenced by NEML2Action::checkStatefulConsistency(), NEML2Action::setupDerivativeMappings(), and NEML2Action::setupInputMappings().

◆ jacobian()

virtual std::pair< std::map< std::string, at::Tensor >, neml2::aoti::VariablePairJacobian > NEML2ModelHandle::jacobian ( const std::map< std::string, at::Tensor > & inputs ) const

pure virtual

Evaluate the model and its full input Jacobian.

Implemented in EagerModelHandle, and AOTIModelHandle.

Referenced by NEML2ModelExecutor::solve().

◆ output_base_shapes()

virtual const std::vector< std::vector< int64_t > > & NEML2ModelHandle::output_base_shapes ( ) const

pure virtual

Implemented in EagerModelHandle, and AOTIModelHandle.

Referenced by NEML2Action::setupDerivativeMappings(), NEML2Action::setupOutputMappings(), and NEML2Action::setupParameterDerivativeMappings().

◆ output_names()

virtual const std::vector< std::string > & NEML2ModelHandle::output_names ( ) const

pure virtual

Implemented in EagerModelHandle, and AOTIModelHandle.

Referenced by NEML2Action::checkStatefulConsistency(), NEML2Action::setupDerivativeMappings(), NEML2Action::setupInputMappings(), NEML2Action::setupOutputMappings(), and NEML2Action::setupParameterDerivativeMappings().

◆ param_jacobian()

virtual std::pair< std::map< std::string, at::Tensor >, neml2::aoti::VariablePairJacobian > NEML2ModelHandle::param_jacobian ( const std::map< std::string, at::Tensor > & inputs ) const

pure virtual

Evaluate the model and its parameter Jacobian (d output / d parameter).

Implemented in EagerModelHandle, and AOTIModelHandle.

Referenced by NEML2ModelExecutor::solve().