Export AnyTensor models to ONNX (unstable guide).
This subpackage is not a stable library API and not in
anytensor.__all__. Import it explicitly::
from anytensor import export
The recommended serialization target is ONNX. ONNX Runtime (ORT) is
well tested and a common engine for serving a frozen graph. AnyTensor does
not run ops on ORT: keep the portable body, bind weights on Torch or
TensorFlow, then serialize so ORT (or another ONNX runner) can deploy it.
Library authors should keep helpers portable (at.shape, segment ops) and
leave export to the application. These names may change.
Pathway. Keep the AnyTensor body; bind weights; serialize:
- Lightning /
nn.Module — nn.Parameter weights +
:func:to_onnx_torch. Named ONNX initializers.
- Keras / TensorFlow — :func:
as_tensorflow_fn creates named constants
inside the trace, then :func:to_onnx_tensorflow. Outer tensors leak as
graph inputs.
- Flax — :func:
numpy_leaves then (1) or (2). Do not jax2tf.
Symbolic lengths. num_segments = at.shape(nodes)[0], not a Python
int.
as_tensorflow_fn(fn, params)
Bind params as named constants created inside the traced function.
Outer tf.constant / tf.Variable objects become extra ONNX inputs.
Constants constructed during tracing embed as initializers (W:0).
fn is called as fn(*args, params=<rebuilt pytree>, **kwargs).
Source code in anytensor/export/_bind.py
| def as_tensorflow_fn(fn: Callable, params) -> Callable:
"""Bind ``params`` as **named constants created inside** the traced function.
Outer ``tf.constant`` / ``tf.Variable`` objects become extra ONNX inputs.
Constants constructed during tracing embed as initializers (``W:0``).
``fn`` is called as ``fn(*args, params=<rebuilt pytree>, **kwargs)``.
"""
import numpy as np
tf = module_if_loaded("tensorflow", raises=True)
treedef, slots, weights = _named_weight_leaves(params)
stored = {name: np.array(leaf, copy=True) for name, leaf in weights.items()}
def wrapped(*args, **kwargs):
leaves = []
for kind, payload in slots:
if kind == "weight":
leaves.append(tf.constant(stored[payload], name=payload))
else:
leaves.append(payload)
return fn(*args, params=tree.unflatten(treedef, leaves), **kwargs)
wrapped.__name__ = getattr(fn, "__name__", "anytensor_tf_fn")
wrapped.__qualname__ = wrapped.__name__
return wrapped
|
as_torch_module(fn, params=None, *, buffers=False)
Wrap fn as an nn.Module for torch.onnx.export / Lightning.
forward keeps fn's argument names so dynamic_shapes can use
them. When params is given, array leaves are registered as named
nn.Parameter values (or buffers) and fn is called as
fn(*args, params=<rebuilt pytree>). Those Parameters become named
ONNX initializers, not graph inputs.
Source code in anytensor/export/_bind.py
| def as_torch_module(fn: Callable, params=None, *, buffers: bool = False) -> Any:
"""Wrap ``fn`` as an ``nn.Module`` for ``torch.onnx.export`` / Lightning.
``forward`` keeps ``fn``'s argument names so ``dynamic_shapes`` can use
them. When ``params`` is given, array leaves are registered as named
``nn.Parameter`` values (or buffers) and ``fn`` is called as
``fn(*args, params=<rebuilt pytree>)``. Those Parameters become **named
ONNX initializers**, not graph inputs.
"""
import numpy as np
torch = module_if_loaded("torch", raises=True)
names = _forward_arg_names(fn)
bound = params is not None
treedef = slots = weights = None
if bound:
treedef, slots, weights = _named_weight_leaves(params)
class _Fn(torch.nn.Module):
def __init__(self):
super().__init__()
self._treedef = treedef
self._slots = slots
self._weight_names = []
if weights:
for name, leaf in weights.items():
tensor = torch.from_numpy(np.array(leaf, copy=True))
if buffers:
self.register_buffer(name, tensor)
else:
self.register_parameter(name, torch.nn.Parameter(tensor))
self._weight_names.append(name)
def _params_tree(self):
leaves = []
for kind, payload in self._slots:
if kind == "weight":
leaves.append(getattr(self, payload))
else:
leaves.append(payload)
return tree.unflatten(self._treedef, leaves)
def forward(self, *args, **kwargs):
if self._treedef is None:
return fn(*args, **kwargs)
return fn(*args, params=self._params_tree(), **kwargs)
_patch_forward_signature(_Fn.forward, names)
_Fn.__name__ = getattr(fn, "__name__", "AnyTensorModule")
_Fn.__qualname__ = _Fn.__name__
return _Fn()
|
Map array leaves to NumPy so a Flax (or JAX) pytree can rebind onto TF / Torch.
Source code in anytensor/export/_bind.py
| def numpy_leaves(structure):
"""Map array leaves to NumPy so a Flax (or JAX) pytree can rebind onto TF / Torch."""
import numpy as np
def _leaf(x):
if hasattr(x, "__array__"):
return np.asarray(x)
return x
return tree.map(_leaf, _plain_pytree(structure))
|
assert_embedded_weights(
model, params, *, require_names=True
)
Require every array leaf of params to be an ONNX initializer, not a feed.
require_names (default True) also demands the initializer name is the
pytree path (Torch W) or TF's W:0 suffix. Returns
{param_name: initializer_name}.
Source code in anytensor/export/_graph.py
| def assert_embedded_weights(
model,
params,
*,
require_names: bool = True,
) -> dict[str, str]:
"""Require every array leaf of ``params`` to be an ONNX initializer, not a feed.
``require_names`` (default True) also demands the initializer name is the
pytree path (Torch ``W``) or TF's ``W:0`` suffix. Returns
``{param_name: initializer_name}``.
"""
proto = _model_proto(model)
inits = _initializer_arrays(proto)
feeds = set(_required_input_names(proto))
_treedef, _slots, weights = _named_weight_leaves(params)
if not weights:
raise AssertionError("params has no array leaves to embed")
used: set[str] = set()
matched: dict[str, str] = {}
leaked = sorted(n for n in weights if n in feeds or f"{n}:0" in feeds)
if leaked:
raise AssertionError(
f"weights leaked as ONNX graph inputs (pass them via as_torch_module "
f"/ as_tensorflow_fn, not as extra arguments): {leaked}"
)
for name, value in weights.items():
key = _match_initializer(inits, name, value, used, require_names=require_names)
if key is None:
raise AssertionError(
f"weight {name!r} is not an embedded ONNX initializer "
f"(names={sorted(inits)}, require_names={require_names})"
)
used.add(key)
matched[name] = key
return matched
|
assert_symbolic_lengths(
model, *, inputs=None, outputs=None
)
Require listed axes to be symbolic (dim_param), not a baked dim_value.
inputs / outputs map value names (or unique suffixes) to axis
indices that must be symbolic. Omit a map to require every rank≥1
input or output to have at least one symbolic axis.
Source code in anytensor/export/_graph.py
| def assert_symbolic_lengths(
model,
*,
inputs: Mapping[str, Sequence[int]] | None = None,
outputs: Mapping[str, Sequence[int]] | None = None,
) -> dict[str, tuple[str | int, ...]]:
"""Require listed axes to be symbolic (``dim_param``), not a baked ``dim_value``.
``inputs`` / ``outputs`` map value names (or unique suffixes) to axis
indices that must be symbolic. Omit a map to require **every** rank≥1
input or output to have at least one symbolic axis.
"""
dims = _symbolic_dims(model)
proto = _model_proto(model)
if inputs is None:
_require_any_symbolic(dims, [v.name for v in proto.graph.input], kind="input")
else:
_require_axes(dims, inputs, kind="input")
if outputs is None:
_require_any_symbolic(dims, [v.name for v in proto.graph.output], kind="output")
else:
_require_axes(dims, outputs, kind="output")
return dims
|
to_onnx_tensorflow(
fn, input_signature, *, params=None, opset=18
)
tf.function + tf2onnx.convert.from_function (None dims stay symbolic).
Pass params= to embed weights via :func:as_tensorflow_fn (named
constants created inside the traced function).
Source code in anytensor/export/_serialize.py
| def to_onnx_tensorflow(fn: Callable, input_signature, *, params=None, opset: int = 18):
"""``tf.function`` + ``tf2onnx.convert.from_function`` (``None`` dims stay symbolic).
Pass ``params=`` to embed weights via :func:`as_tensorflow_fn` (named
constants created inside the traced function).
"""
tf = module_if_loaded("tensorflow", raises=True)
try:
import tf2onnx
except ImportError as exc:
raise RuntimeError(
"to_onnx_tensorflow requires tf2onnx (pip install tf2onnx)"
) from exc
if params is not None:
fn = as_tensorflow_fn(fn, params)
wrapped = tf.function(fn, input_signature=input_signature)
proto, _storage = tf2onnx.convert.from_function(
wrapped, input_signature=input_signature, opset=opset
)
return proto
|
to_onnx_torch(
model,
args,
*,
params=None,
dynamic_shapes=None,
input_names=None,
output_names=None,
**kwargs,
)
torch.onnx.export with the dynamo / torch.export path (dynamic shapes).
Pass params= to bind a weight pytree as nn.Parameter initializers.
Source code in anytensor/export/_serialize.py
| def to_onnx_torch(
model,
args,
*,
params=None,
dynamic_shapes=None,
input_names=None,
output_names=None,
**kwargs,
):
"""``torch.onnx.export`` with the dynamo / ``torch.export`` path (dynamic shapes).
Pass ``params=`` to bind a weight pytree as ``nn.Parameter`` initializers.
"""
torch = module_if_loaded("torch", raises=True)
if params is not None:
if isinstance(model, torch.nn.Module):
raise TypeError(
"params= is for callables; register nn.Parameter on the module instead"
)
model = as_torch_module(model, params)
elif not isinstance(model, torch.nn.Module):
model = as_torch_module(model)
model.eval()
kwargs.setdefault("dynamo", True)
return torch.onnx.export(
model,
args,
dynamic_shapes=dynamic_shapes,
input_names=input_names,
output_names=output_names,
**kwargs,
)
|