onnx_diagnostic.reference.torch_ops

modules

OpRunKernel

class onnx_diagnostic.reference.torch_ops.OpRunKernel(node: NodeProto, version: int | None = None, verbose: int = 0, custom_kernels: Dict[Tuple[str, str], type] | None = None)[source][source]

Main class. Every kernel should inherit from it. It does not copy the proto.

classmethod device_dependent() → bool[source][source]

Returns True if the kernel needs a device to be efficiently initialized.

get_attribute_float(node: NodeProto, name: str, default_value: float | None = None) → float | None[source][source]

Returns an attribute as an int.

Parameters:

  • node – NodeProto

  • name – name

  • default_value – default_value

Returns:

value

get_attribute_int(node: NodeProto, name: str, default_value: int | None = None) → int | None[source][source]

Returns an attribute as an int.

Parameters:

  • node – NodeProto

  • name – name

  • default_value – default_value

Returns:

value

get_attribute_ints(node: NodeProto, name: str, default_value: Tuple[int, ...] | None = None) → Tuple[int, ...] | None[source][source]

Returns an attribute as a tuple of ints.

Parameters:

  • node – NodeProto

  • name – name

  • default_value – default_value

Returns:

value

get_attribute_string(node: NodeProto, name: str, default_value: str | None = None) → str | None[source][source]

Returns an attribute as a tuple of ints.

Parameters:

  • node – NodeProto

  • name – name

  • default_value – default_value

Returns:

value

get_attribute_tensor(node: NodeProto, name: str) → Tensor | None[source][source]

Returns an attribute as a torch tensor.

Parameters:

  • node – NodeProto

  • name – name

  • default_value – default_value

Returns:

value

classmethod has_subgraphs() → bool[source][source]

Returns True if the kernel has subgraphs.

run(*args: OpRunValue | None) → OpRunValue | Tuple[OpRunValue | None, ...][source][source]

Kernel implementation.

same_device(*tensors: Tensor) → Tuple[Tensor, ...][source][source]

Puts all tensors on the same device.

OpRunTensor

class onnx_diagnostic.reference.torch_ops.OpRunTensor(tensor, is_constant: bool = False, may_cpu: bool = False)[source][source]

Wrapper around a tensor.

Parameters:

  • tensor – torch.Tensor

  • is_constant – is it a constant

  • may_cpu – change the device the tensor is if more appropriate

property as_tuple_int: Tuple[int, ...]

value as int

copy() → OpRunTensor[source][source]

Shallow copy.

property device

Returns the device.

property dtype
get_device() → int[source][source]

Returns the device id.

classmethod is_sequence() → bool[source][source]

Tells if it is sequence.

numel() → int[source][source]

Returns the number of elements.

property shape
string_type() → str[source][source]

Returns information about the value as a string.

property tensor_or_sequence: Tensor | List[Tensor]

Returns either a tensor or a sequence.

to(to: Any) → OpRunTensor[source][source]

Changes the device.

OpRunValue

class onnx_diagnostic.reference.torch_ops.OpRunValue[source][source]

Defines a value for the runtime, a tensor or a sequence.

classmethod is_sequence() → bool[source][source]

Tells if it is sequence.

OpRunSequence

class onnx_diagnostic.reference.torch_ops.OpRunSequence(sequence: List[Tensor] | None = None, dtype: dtype = torch.float32)[source][source]

Defines a sequence.

copy() → OpRunSequence[source][source]

Shallow copy.

property dtype
insert_at(tensor: Tensor, position: OpRunTensor | None = None) → OpRunSequence[source][source]

Inserts a value at a given position.

classmethod is_sequence() → bool[source][source]

Tells if it is sequence.

string_type() → str[source][source]

Returns a string which can be printed.

property tensor_or_sequence: Tensor | List[Tensor]

Returns either a tensor or a sequence.

OpRunFunction

class onnx_diagnostic.reference.torch_ops.OpRunFunction(runtime: onnx_diagnostic.reference.TorchOnnxEvaluator, node: NodeProto, version: int | None = None, verbose: int = 0)[source][source]

Defines a kernel based on a local functions.

run(*args: OpRunValue | None) → OpRunValue | Tuple[OpRunValue | None, ...][source][source]

Kernel implementation.

Other functions