EagerModelHandle Class Reference

Handle wrapping the cpp-eager runtime (embeds CPython; loads from a source .i). More...

#include <NEML2ModelHandle.h>

Inheritance diagram for EagerModelHandle:

Public Member Functions
	EagerModelHandle (const std::string &input_file, const std::string &model_name, const at::Device &device, const std::vector< std::string > &load)
const std::vector< std::string > &	input_names () const override
const std::vector< std::string > &	output_names () const override
const std::vector< std::vector< int64_t > > &	input_base_shapes () const override
const std::vector< std::vector< int64_t > > &	output_base_shapes () const override
const std::map< std::string, std::vector< int64_t > > &	parameter_base_shapes () const override
std::pair< std::map< std::string, at::Tensor >, neml2::aoti::VariablePairJacobian >	jacobian (const std::map< std::string, at::Tensor > &inputs) const override
	Evaluate the model and its full input Jacobian.
std::pair< std::map< std::string, at::Tensor >, neml2::aoti::VariablePairJacobian >	param_jacobian (const std::map< std::string, at::Tensor > &inputs) const override
	Evaluate the model and its parameter Jacobian (d output / d parameter).
void	set_parameter (const std::string &name, const at::Tensor &value) override
	Replace a (runtime-flexible) model parameter's value.
std::map< std::string, at::Tensor >	value (const std::map< std::string, at::Tensor > &inputs) const override
at::Device	device () const override
	The compute device this rank's workload runs on (for cpp-aoti, the scheduler-assigned device).

Private Attributes
neml2::eager::Model	_m

Detailed Description

Handle wrapping the cpp-eager runtime (embeds CPython; loads from a source .i).

Single-device: the model is pinned to one device (cpp-eager has no scheduler / multi-device dispatch).

Definition at line 83 of file NEML2ModelHandle.h.

Constructor & Destructor Documentation

EagerModelHandle()

EagerModelHandle::EagerModelHandle ( const std::string &  input_file, const std::string &  model_name, const at::Device &  device, const std::vector< std::string > &  load  )

inline

Definition at line 86 of file NEML2ModelHandle.h.

    : _m(input_file, model_name, device, load)
  {
  }

Member Function Documentation

device()

at::Device EagerModelHandle::device ( ) const

inlineoverridevirtual

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

Implements NEML2ModelHandle.

Definition at line 128 of file NEML2ModelHandle.h.

{ return _m.device(); }

input_base_shapes()

const std::vector< std::vector< int64_t > > & EagerModelHandle::input_base_shapes ( ) const

inlineoverridevirtual

Implements NEML2ModelHandle.

Definition at line 96 of file NEML2ModelHandle.h.

  {
    return _m.input_base_shapes();
  }

input_names()

const std::vector< std::string > & EagerModelHandle::input_names ( ) const

inlineoverridevirtual

Implements NEML2ModelHandle.

Definition at line 94 of file NEML2ModelHandle.h.

{ return _m.input_names(); }

jacobian()

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

inlineoverridevirtual

Evaluate the model and its full input Jacobian.

Implements NEML2ModelHandle.

Definition at line 110 of file NEML2ModelHandle.h.

  {
    return _m.jacobian(inputs);
  }

output_base_shapes()

const std::vector< std::vector< int64_t > > & EagerModelHandle::output_base_shapes ( ) const

inlineoverridevirtual

Implements NEML2ModelHandle.

Definition at line 100 of file NEML2ModelHandle.h.

  {
    return _m.output_base_shapes();
  }

output_names()

const std::vector< std::string > & EagerModelHandle::output_names ( ) const

inlineoverridevirtual

Implements NEML2ModelHandle.

Definition at line 95 of file NEML2ModelHandle.h.

{ return _m.output_names(); }

param_jacobian()

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

inlineoverridevirtual

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

Implements NEML2ModelHandle.

Definition at line 115 of file NEML2ModelHandle.h.

  {
    return _m.param_jacobian(inputs);
  }

parameter_base_shapes()

const std::map< std::string, std::vector< int64_t > > & EagerModelHandle::parameter_base_shapes ( ) const

inlineoverridevirtual

Implements NEML2ModelHandle.

Definition at line 104 of file NEML2ModelHandle.h.

  {
    return _m.parameter_base_shapes();
  }

set_parameter()

void EagerModelHandle::set_parameter ( const std::string &  name, const at::Tensor &  value  )

inlineoverridevirtual

Replace a (runtime-flexible) model parameter's value.

Implements NEML2ModelHandle.

Definition at line 119 of file NEML2ModelHandle.h.

  {
    _m.set_parameter(name, value);
  }

value()

std::map< std::string, at::Tensor > EagerModelHandle::value ( const std::map< std::string, at::Tensor > &  inputs ) const

inlineoverridevirtual

Implements NEML2ModelHandle.

Definition at line 124 of file NEML2ModelHandle.h.

  {
    return _m.forward(inputs);
  }

Member Data Documentation

_m

neml2::eager::Model EagerModelHandle::_m

private

Definition at line 131 of file NEML2ModelHandle.h.

Referenced by device(), input_base_shapes(), input_names(), jacobian(), output_base_shapes(), output_names(), param_jacobian(), parameter_base_shapes(), set_parameter(), and value().

---

The documentation for this class was generated from the following file:

• include/neml2/interfaces/NEML2ModelHandle.h