AOTIModelHandle Class Reference

Handle wrapping the cpp-aoti runtime: a neml2::aoti::DispatchedModel backed by the self-describing artifact folder that neml2-compile writes (a shared metadata.json plus per-<device>/<dtype>/ .pt2 binaries). More...

#include <NEML2ModelHandle.h>

Inheritance diagram for AOTIModelHandle:

Public Member Functions
	AOTIModelHandle (const std::string &input, const std::string &model_name, const std::vector< std::string > &devices, const std::vector< std::size_t > &batch_sizes, const void *comm)
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).

Static Private Member Functions
static neml2::aoti::DispatchedModel	load (const std::string &input, const std::string &model_name, const std::vector< std::string > &devices, const std::vector< std::size_t > &batch_sizes, const void *comm)
	Build the MPI scheduler (round-robin device list over the node's ranks) and load the dispatched aoti model.

Private Attributes
neml2::aoti::DispatchedModel	_m

Detailed Description

Handle wrapping the cpp-aoti runtime: a neml2::aoti::DispatchedModel backed by the self-describing artifact folder that neml2-compile writes (a shared metadata.json plus per-<device>/<dtype>/ .pt2 binaries).

That folder carries everything the runtime needs – structure, solver config, I/O names, promoted parameters – so the neml2-compile stub .i is no longer required (it was only a name -> artifact_path indirection).

For a friendly UX the input is resolved flexibly (see load): the artifact folder itself, the metadata.json inside it, or the legacy stub .i are all accepted. An MPISimpleScheduler pins this rank's workload to a device, round-robining the device list over the MPI ranks on each node (the single "cpu" entry sends every rank's local batch to CPU). This matches MOOSE's MPI domain decomposition – each rank evaluates its own local batch.

Definition at line 147 of file NEML2ModelHandle.h.

Constructor & Destructor Documentation

AOTIModelHandle()

AOTIModelHandle::AOTIModelHandle ( const std::string &  input, const std::string &  model_name, const std::vector< std::string > &  devices, const std::vector< std::size_t > &  batch_sizes, const void *  comm  )

inline

Definition at line 150 of file NEML2ModelHandle.h.

    : _m(load(input, model_name, devices, batch_sizes, comm))
  {
  }

Member Function Documentation

device()

at::Device AOTIModelHandle::device ( ) const

inlineoverridevirtual

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

Implements NEML2ModelHandle.

Definition at line 193 of file NEML2ModelHandle.h.

{ return _m.device(); }

input_base_shapes()

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

inlineoverridevirtual

Implements NEML2ModelHandle.

Definition at line 161 of file NEML2ModelHandle.h.

  {
    return _m.input_base_shapes();
  }

input_names()

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

inlineoverridevirtual

Implements NEML2ModelHandle.

Definition at line 159 of file NEML2ModelHandle.h.

{ return _m.input_names(); }

jacobian()

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

inlineoverridevirtual

Evaluate the model and its full input Jacobian.

Implements NEML2ModelHandle.

Definition at line 175 of file NEML2ModelHandle.h.

  {
    return _m.jacobian(inputs);
  }

load()

static neml2::aoti::DispatchedModel AOTIModelHandle::load ( const std::string &  input, const std::string &  model_name, const std::vector< std::string > &  devices, const std::vector< std::size_t > &  batch_sizes, const void *  comm  )

inlinestaticprivate

Build the MPI scheduler (round-robin device list over the node's ranks) and load the dispatched aoti model.

input is resolved flexibly for a friendlier UX:

• a directory -> the artifact root, loaded directly (self-describing);
• a metadata.json file -> its parent directory (metadata.json always sits at the root);
• anything else (a stub .i)-> parsed by neml2::aoti::load_model, which resolves the recorded artifact_path (absolute, or relative to the stub). The folder and metadata.json forms are self-describing and ignore model_name; only the stub form selects a named model, so model_name is required there (as before). Prefer the folder: it is the relocatable unit – the stub references its folder by path (increasingly a relative one in neml2), so moving the stub alone breaks that link, whereas the folder moves cleanly.

Definition at line 206 of file NEML2ModelHandle.h.

  {
    neml2::aoti::MPISimpleScheduler::Config config;
    config.devices = devices;
    config.batch_sizes = batch_sizes;
    config.comm = comm;
    auto scheduler = std::make_shared<neml2::aoti::MPISimpleScheduler>(config);
 
    const std::filesystem::path path(input);
    if (std::filesystem::is_directory(path))
      return neml2::aoti::DispatchedModel(path, scheduler);
    if (path.filename() == "metadata.json")
      return neml2::aoti::DispatchedModel(path.parent_path(), scheduler);
    return neml2::aoti::load_model(input, model_name, scheduler);
  }

output_base_shapes()

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

inlineoverridevirtual

Implements NEML2ModelHandle.

Definition at line 165 of file NEML2ModelHandle.h.

  {
    return _m.output_base_shapes();
  }

output_names()

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

inlineoverridevirtual

Implements NEML2ModelHandle.

Definition at line 160 of file NEML2ModelHandle.h.

{ return _m.output_names(); }

param_jacobian()

std::pair< std::map< std::string, at::Tensor >, neml2::aoti::VariablePairJacobian > AOTIModelHandle::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 180 of file NEML2ModelHandle.h.

  {
    return _m.param_jacobian(inputs);
  }

parameter_base_shapes()

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

inlineoverridevirtual

Implements NEML2ModelHandle.

Definition at line 169 of file NEML2ModelHandle.h.

  {
    return _m.parameter_base_shapes();
  }

set_parameter()

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

inlineoverridevirtual

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

Implements NEML2ModelHandle.

Definition at line 184 of file NEML2ModelHandle.h.

  {
    _m.set_parameter(name, value);
  }

value()

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

inlineoverridevirtual

Implements NEML2ModelHandle.

Definition at line 189 of file NEML2ModelHandle.h.

  {
    return _m.forward(inputs);
  }

Member Data Documentation

_m

neml2::aoti::DispatchedModel AOTIModelHandle::_m

private

Definition at line 226 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