Coverage for src/bioimageio/spec/model/v0_4.py: 91%
631 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 10:20 +0000
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 10:20 +0000
1from __future__ import annotations
3from collections.abc import Mapping, Sequence
4from typing import (
5 TYPE_CHECKING,
6 Annotated,
7 Any,
8 Callable,
9 ClassVar,
10 Literal,
11 Union,
12 cast,
13)
15import numpy as np
16from annotated_types import Ge, Interval, MaxLen, MinLen, MultipleOf
17from loguru import logger
18from numpy.typing import NDArray
19from pydantic import (
20 AllowInfNan,
21 Discriminator,
22 Field,
23 RootModel,
24 SerializationInfo,
25 SerializerFunctionWrapHandler,
26 StringConstraints,
27 TypeAdapter,
28 ValidationInfo,
29 WrapSerializer,
30 field_validator,
31 model_validator,
32)
33from typing_extensions import Self, assert_never, get_args
35from .._internal.common_nodes import (
36 KwargsNode,
37 Node,
38 NodeWithExplicitlySetFields,
39)
40from .._internal.constants import SHA256_HINT
41from .._internal.field_validation import validate_unique_entries
42from .._internal.field_warning import issue_warning, warn
43from .._internal.io import BioimageioYamlContent, WithSuffix
44from .._internal.io import FileDescr as FileDescr
45from .._internal.io_basics import Sha256 as Sha256
46from .._internal.io_packaging import FileSource_package, include_in_package
47from .._internal.io_utils import load_array
48from .._internal.packaging_context import packaging_context_var
49from .._internal.types import Datetime as Datetime
50from .._internal.types import FileSource, LowerCaseIdentifier
51from .._internal.types import Identifier as Identifier
52from .._internal.types import LicenseId as LicenseId
53from .._internal.types import NotEmpty as NotEmpty
54from .._internal.url import HttpUrl as HttpUrl
55from .._internal.validated_string_with_inner_node import ValidatedStringWithInnerNode
56from .._internal.validator_annotations import AfterValidator, RestrictCharacters
57from .._internal.version_type import Version as Version
58from .._internal.warning_levels import ALERT, INFO
59from ..dataset.v0_2 import VALID_COVER_IMAGE_EXTENSIONS as VALID_COVER_IMAGE_EXTENSIONS
60from ..dataset.v0_2 import DatasetDescr as DatasetDescr
61from ..dataset.v0_2 import LinkedDataset as LinkedDataset
62from ..generic.v0_2 import AttachmentsDescr as AttachmentsDescr
63from ..generic.v0_2 import Author as Author
64from ..generic.v0_2 import BadgeDescr as BadgeDescr
65from ..generic.v0_2 import CiteEntry as CiteEntry
66from ..generic.v0_2 import Doi as Doi
67from ..generic.v0_2 import GenericModelDescrBase
68from ..generic.v0_2 import LinkedResource as LinkedResource
69from ..generic.v0_2 import Maintainer as Maintainer
70from ..generic.v0_2 import OrcidId as OrcidId
71from ..generic.v0_2 import RelativeFilePath as RelativeFilePath
72from ..generic.v0_2 import ResourceId as ResourceId
73from ..generic.v0_2 import Uploader as Uploader
74from ._v0_4_converter import convert_from_older_format
77class ModelId(ResourceId):
78 pass
81AxesStr = Annotated[
82 str, RestrictCharacters("bitczyx"), AfterValidator(validate_unique_entries)
83]
84AxesInCZYX = Annotated[
85 str, RestrictCharacters("czyx"), AfterValidator(validate_unique_entries)
86]
88PostprocessingName = Literal[
89 "binarize",
90 "clip",
91 "scale_linear",
92 "sigmoid",
93 "zero_mean_unit_variance",
94 "scale_range",
95 "scale_mean_variance",
96]
97PreprocessingName = Literal[
98 "binarize",
99 "clip",
100 "scale_linear",
101 "sigmoid",
102 "zero_mean_unit_variance",
103 "scale_range",
104]
107class TensorName(LowerCaseIdentifier):
108 pass
111class CallableFromDepencencyNode(Node):
112 _submodule_adapter: ClassVar[TypeAdapter[Identifier]] = TypeAdapter(Identifier)
114 module_name: str
115 """The Python module that implements **callable_name**."""
117 @field_validator("module_name", mode="after")
118 def _check_submodules(cls, module_name: str) -> str:
119 for submod in module_name.split("."):
120 _ = cls._submodule_adapter.validate_python(submod)
122 return module_name
124 callable_name: Identifier
125 """The callable Python identifier implemented in module **module_name**."""
128class CallableFromDepencency(ValidatedStringWithInnerNode[CallableFromDepencencyNode]):
129 _inner_node_class = CallableFromDepencencyNode
130 root_model: ClassVar[type[RootModel[Any]]] = RootModel[
131 Annotated[
132 str,
133 StringConstraints(strip_whitespace=True, pattern=r"^.+\..+$"),
134 ]
135 ]
137 @classmethod
138 def _get_data(cls, valid_string_data: str):
139 *mods, callname = valid_string_data.split(".")
140 return {"module_name": ".".join(mods), "callable_name": callname}
142 @property
143 def module_name(self):
144 """The Python module that implements **callable_name**."""
145 return self._inner_node.module_name
147 @property
148 def callable_name(self):
149 """The callable Python identifier implemented in module **module_name**."""
150 return self._inner_node.callable_name
153class CallableFromFileNode(Node):
154 source_file: Annotated[
155 RelativeFilePath | HttpUrl,
156 Field(union_mode="left_to_right"),
157 include_in_package,
158 ]
159 """The Python source file that implements **callable_name**."""
160 callable_name: Identifier
161 """The callable Python identifier implemented in **source_file**."""
164class CallableFromFile(ValidatedStringWithInnerNode[CallableFromFileNode]):
165 _inner_node_class = CallableFromFileNode
166 root_model: ClassVar[type[RootModel[Any]]] = RootModel[
167 Annotated[
168 str,
169 StringConstraints(strip_whitespace=True, pattern=r"^.+:.+$"),
170 ]
171 ]
173 @classmethod
174 def _get_data(cls, valid_string_data: str):
175 *file_parts, callname = valid_string_data.split(":")
176 return {"source_file": ":".join(file_parts), "callable_name": callname}
178 @property
179 def source_file(self):
180 """The Python source file that implements **callable_name**."""
181 return self._inner_node.source_file
183 @property
184 def callable_name(self):
185 """The callable Python identifier implemented in **source_file**."""
186 return self._inner_node.callable_name
189CustomCallable = Annotated[
190 Union[CallableFromFile, CallableFromDepencency], Field(union_mode="left_to_right")
191]
194class DependenciesNode(Node):
195 manager: Annotated[NotEmpty[str], Field(examples=["conda", "maven", "pip"])]
196 """Dependency manager"""
198 file: FileSource_package
199 """Dependency file"""
202class Dependencies(ValidatedStringWithInnerNode[DependenciesNode]):
203 _inner_node_class = DependenciesNode
204 root_model: ClassVar[type[RootModel[Any]]] = RootModel[
205 Annotated[
206 str,
207 StringConstraints(strip_whitespace=True, pattern=r"^.+:.+$"),
208 ]
209 ]
211 @classmethod
212 def _get_data(cls, valid_string_data: str):
213 manager, *file_parts = valid_string_data.split(":")
214 return {"manager": manager, "file": ":".join(file_parts)}
216 @property
217 def manager(self):
218 """Dependency manager"""
219 return self._inner_node.manager
221 @property
222 def file(self):
223 """Dependency file"""
224 return self._inner_node.file
227WeightsFormat = Literal[
228 "keras_hdf5",
229 "onnx",
230 "pytorch_state_dict",
231 "tensorflow_js",
232 "tensorflow_saved_model_bundle",
233 "torchscript",
234]
237class WeightsEntryDescrBase(FileDescr):
238 type: ClassVar[WeightsFormat]
239 weights_format_name: ClassVar[str] # human readable
241 source: FileSource_package
242 """The weights file."""
244 attachments: Annotated[
245 AttachmentsDescr | None,
246 warn(None, "Weights entry depends on additional attachments.", ALERT),
247 ] = None
248 """Attachments that are specific to this weights entry."""
250 authors: list[Author] | None = None
251 """Authors
252 Either the person(s) that have trained this model resulting in the original weights file.
253 (If this is the initial weights entry, i.e. it does not have a `parent`)
254 Or the person(s) who have converted the weights to this weights format.
255 (If this is a child weight, i.e. it has a `parent` field)
256 """
258 dependencies: Annotated[
259 Dependencies | None,
260 warn(
261 None,
262 "Custom dependencies ({value}) specified. Avoid this whenever possible "
263 + "to allow execution in a wider range of software environments.",
264 ),
265 Field(
266 examples=[
267 "conda:environment.yaml",
268 "maven:./pom.xml",
269 "pip:./requirements.txt",
270 ]
271 ),
272 ] = None
273 """Dependency manager and dependency file, specified as `<dependency manager>:<relative file path>`."""
275 parent: Annotated[WeightsFormat | None, Field(examples=["pytorch_state_dict"])] = (
276 None
277 )
278 """The source weights these weights were converted from.
279 For example, if a model's weights were converted from the `pytorch_state_dict` format to `torchscript`,
280 The `pytorch_state_dict` weights entry has no `parent` and is the parent of the `torchscript` weights.
281 All weight entries except one (the initial set of weights resulting from training the model),
282 need to have this field."""
284 @model_validator(mode="after")
285 def check_parent_is_not_self(self) -> Self:
286 if self.type == self.parent:
287 raise ValueError("Weights entry can't be it's own parent.")
289 return self
292class KerasHdf5WeightsDescr(WeightsEntryDescrBase):
293 type: ClassVar[WeightsFormat] = "keras_hdf5"
294 weights_format_name: ClassVar[str] = "Keras HDF5"
295 tensorflow_version: Version | None = None
296 """TensorFlow version used to create these weights"""
298 @field_validator("tensorflow_version", mode="after")
299 @classmethod
300 def _tfv(cls, value: Any):
301 if value is None:
302 issue_warning(
303 "missing. Please specify the TensorFlow version"
304 + " these weights were created with.",
305 value=value,
306 severity=ALERT,
307 field="tensorflow_version",
308 )
309 return value
312class OnnxWeightsDescr(WeightsEntryDescrBase):
313 type: ClassVar[WeightsFormat] = "onnx"
314 weights_format_name: ClassVar[str] = "ONNX"
315 opset_version: Annotated[int, Ge(7)] | None = None
316 """ONNX opset version"""
318 @field_validator("opset_version", mode="after")
319 @classmethod
320 def _ov(cls, value: Any):
321 if value is None:
322 issue_warning(
323 "Missing ONNX opset version (aka ONNX opset number). "
324 + "Please specify the ONNX opset version these weights were created"
325 + " with.",
326 value=value,
327 severity=ALERT,
328 field="opset_version",
329 )
330 return value
333class PytorchStateDictWeightsDescr(WeightsEntryDescrBase):
334 type: ClassVar[WeightsFormat] = "pytorch_state_dict"
335 weights_format_name: ClassVar[str] = "Pytorch State Dict"
336 architecture: CustomCallable = Field(
337 examples=["my_function.py:MyNetworkClass", "my_module.submodule.get_my_model"]
338 )
339 """callable returning a torch.nn.Module instance.
340 Local implementation: `<relative path to file>:<identifier of implementation within the file>`.
341 Implementation in a dependency: `<dependency-package>.<[dependency-module]>.<identifier>`."""
343 architecture_sha256: Annotated[
344 Sha256 | None,
345 Field(
346 description=(
347 "The SHA256 of the architecture source file, if the architecture is not"
348 " defined in a module listed in `dependencies`\n"
349 )
350 + SHA256_HINT,
351 ),
352 ] = None
353 """The SHA256 of the architecture source file,
354 if the architecture is not defined in a module listed in `dependencies`"""
356 @model_validator(mode="after")
357 def check_architecture_sha256(self) -> Self:
358 if isinstance(self.architecture, CallableFromFile):
359 if self.architecture_sha256 is None:
360 raise ValueError(
361 "Missing required `architecture_sha256` for `architecture` with"
362 + " source file."
363 )
364 elif self.architecture_sha256 is not None:
365 raise ValueError(
366 "Got `architecture_sha256` for architecture that does not have a source"
367 + " file."
368 )
370 return self
372 kwargs: dict[str, Any] = Field(
373 default_factory=cast(Callable[[], dict[str, Any]], dict)
374 )
375 """key word arguments for the `architecture` callable"""
377 pytorch_version: Version | None = None
378 """Version of the PyTorch library used.
379 If `depencencies` is specified it should include pytorch and the verison has to match.
380 (`dependencies` overrules `pytorch_version`)"""
382 @field_validator("pytorch_version", mode="after")
383 @classmethod
384 def _ptv(cls, value: Any):
385 if value is None:
386 issue_warning(
387 "missing. Please specify the PyTorch version these"
388 + " PyTorch state dict weights were created with.",
389 value=value,
390 severity=ALERT,
391 field="pytorch_version",
392 )
393 return value
396class TorchscriptWeightsDescr(WeightsEntryDescrBase):
397 type: ClassVar[WeightsFormat] = "torchscript"
398 weights_format_name: ClassVar[str] = "TorchScript"
399 pytorch_version: Version | None = None
400 """Version of the PyTorch library used."""
402 @field_validator("pytorch_version", mode="after")
403 @classmethod
404 def _ptv(cls, value: Any):
405 if value is None:
406 issue_warning(
407 "missing. Please specify the PyTorch version these"
408 + " Torchscript weights were created with.",
409 value=value,
410 severity=ALERT,
411 field="pytorch_version",
412 )
413 return value
416class TensorflowJsWeightsDescr(WeightsEntryDescrBase):
417 type: ClassVar[WeightsFormat] = "tensorflow_js"
418 weights_format_name: ClassVar[str] = "Tensorflow.js"
419 tensorflow_version: Version | None = None
420 """Version of the TensorFlow library used."""
422 @field_validator("tensorflow_version", mode="after")
423 @classmethod
424 def _tfv(cls, value: Any):
425 if value is None:
426 issue_warning(
427 "missing. Please specify the TensorFlow version"
428 + " these TensorflowJs weights were created with.",
429 value=value,
430 severity=ALERT,
431 field="tensorflow_version",
432 )
433 return value
435 source: FileSource_package
436 """The multi-file weights.
437 All required files/folders should be a zip archive."""
440class TensorflowSavedModelBundleWeightsDescr(WeightsEntryDescrBase):
441 type: ClassVar[WeightsFormat] = "tensorflow_saved_model_bundle"
442 weights_format_name: ClassVar[str] = "Tensorflow Saved Model"
443 tensorflow_version: Version | None = None
444 """Version of the TensorFlow library used."""
446 @field_validator("tensorflow_version", mode="after")
447 @classmethod
448 def _tfv(cls, value: Any):
449 if value is None:
450 issue_warning(
451 "missing. Please specify the TensorFlow version"
452 + " these Tensorflow saved model bundle weights were created with.",
453 value=value,
454 severity=ALERT,
455 field="tensorflow_version",
456 )
457 return value
460class WeightsDescr(Node):
461 keras_hdf5: KerasHdf5WeightsDescr | None = None
462 onnx: OnnxWeightsDescr | None = None
463 pytorch_state_dict: PytorchStateDictWeightsDescr | None = None
464 tensorflow_js: TensorflowJsWeightsDescr | None = None
465 tensorflow_saved_model_bundle: TensorflowSavedModelBundleWeightsDescr | None = None
466 torchscript: TorchscriptWeightsDescr | None = None
468 @model_validator(mode="after")
469 def check_one_entry(self) -> Self:
470 if all(
471 entry is None
472 for entry in [
473 self.keras_hdf5,
474 self.onnx,
475 self.pytorch_state_dict,
476 self.tensorflow_js,
477 self.tensorflow_saved_model_bundle,
478 self.torchscript,
479 ]
480 ):
481 raise ValueError("Missing weights entry")
483 return self
485 def __getitem__(
486 self,
487 key: WeightsFormat,
488 ):
489 if key == "keras_hdf5":
490 ret = self.keras_hdf5
491 elif key == "onnx":
492 ret = self.onnx
493 elif key == "pytorch_state_dict":
494 ret = self.pytorch_state_dict
495 elif key == "tensorflow_js":
496 ret = self.tensorflow_js
497 elif key == "tensorflow_saved_model_bundle":
498 ret = self.tensorflow_saved_model_bundle
499 elif key == "torchscript":
500 ret = self.torchscript
501 else:
502 raise KeyError(key)
504 if ret is None:
505 raise KeyError(key)
507 return ret
509 @property
510 def available_formats(self):
511 return {
512 **({} if self.keras_hdf5 is None else {"keras_hdf5": self.keras_hdf5}),
513 **({} if self.onnx is None else {"onnx": self.onnx}),
514 **(
515 {}
516 if self.pytorch_state_dict is None
517 else {"pytorch_state_dict": self.pytorch_state_dict}
518 ),
519 **(
520 {}
521 if self.tensorflow_js is None
522 else {"tensorflow_js": self.tensorflow_js}
523 ),
524 **(
525 {}
526 if self.tensorflow_saved_model_bundle is None
527 else {
528 "tensorflow_saved_model_bundle": self.tensorflow_saved_model_bundle
529 }
530 ),
531 **({} if self.torchscript is None else {"torchscript": self.torchscript}),
532 }
534 @property
535 def missing_formats(self):
536 return {
537 wf for wf in get_args(WeightsFormat) if wf not in self.available_formats
538 }
541class ParameterizedInputShape(Node):
542 """A sequence of valid shapes given by `shape_k = min + k * step for k in {0, 1, ...}`."""
544 min: NotEmpty[list[int]]
545 """The minimum input shape"""
547 step: NotEmpty[list[int]]
548 """The minimum shape change"""
550 def __len__(self) -> int:
551 return len(self.min)
553 @model_validator(mode="after")
554 def matching_lengths(self) -> Self:
555 if len(self.min) != len(self.step):
556 raise ValueError("`min` and `step` required to have the same length")
558 return self
561class ImplicitOutputShape(Node):
562 """Output tensor shape depending on an input tensor shape.
563 `shape(output_tensor) = shape(input_tensor) * scale + 2 * offset`"""
565 reference_tensor: TensorName
566 """Name of the reference tensor."""
568 scale: NotEmpty[list[float | None]]
569 """output_pix/input_pix for each dimension.
570 'null' values indicate new dimensions, whose length is defined by 2*`offset`"""
572 offset: NotEmpty[list[int | Annotated[float, MultipleOf(0.5)]]]
573 """Position of origin wrt to input."""
575 def __len__(self) -> int:
576 return len(self.scale)
578 @model_validator(mode="after")
579 def matching_lengths(self) -> Self:
580 if len(self.scale) != len(self.offset):
581 raise ValueError(
582 f"scale {self.scale} has to have same length as offset {self.offset}!"
583 )
584 # if we have an expanded dimension, make sure that it's offet is not zero
585 for sc, off in zip(self.scale, self.offset):
586 if sc is None and not off:
587 raise ValueError("`offset` must not be zero if `scale` is none/zero")
589 return self
592class TensorDescrBase(Node):
593 name: TensorName
594 """Tensor name. No duplicates are allowed."""
596 description: str = ""
598 axes: AxesStr
599 """Axes identifying characters. Same length and order as the axes in `shape`.
600 | axis | description |
601 | --- | --- |
602 | b | batch (groups multiple samples) |
603 | i | instance/index/element |
604 | t | time |
605 | c | channel |
606 | z | spatial dimension z |
607 | y | spatial dimension y |
608 | x | spatial dimension x |
609 """
611 data_range: (
612 tuple[Annotated[float, AllowInfNan(True)], Annotated[float, AllowInfNan(True)]]
613 | None
614 ) = None
615 """Tuple `(minimum, maximum)` specifying the allowed range of the data in this tensor.
616 If not specified, the full data range that can be expressed in `data_type` is allowed."""
619class BinarizeKwargs(KwargsNode):
620 """key word arguments for `BinarizeDescr`"""
622 threshold: float
623 """The fixed threshold"""
626class BinarizeDescr(NodeWithExplicitlySetFields):
627 """BinarizeDescr the tensor with a fixed `BinarizeKwargs.threshold`.
628 Values above the threshold will be set to one, values below the threshold to zero.
629 """
631 implemented_name: ClassVar[Literal["binarize"]] = "binarize"
632 if TYPE_CHECKING:
633 name: Literal["binarize"] = "binarize"
634 else:
635 name: Literal["binarize"]
637 kwargs: BinarizeKwargs
640class ClipKwargs(KwargsNode):
641 """key word arguments for `ClipDescr`"""
643 min: float
644 """minimum value for clipping"""
645 max: float
646 """maximum value for clipping"""
649class ClipDescr(NodeWithExplicitlySetFields):
650 """Clip tensor values to a range.
652 Set tensor values below `ClipKwargs.min` to `ClipKwargs.min`
653 and above `ClipKwargs.max` to `ClipKwargs.max`.
654 """
656 implemented_name: ClassVar[Literal["clip"]] = "clip"
657 if TYPE_CHECKING:
658 name: Literal["clip"] = "clip"
659 else:
660 name: Literal["clip"]
662 kwargs: ClipKwargs
665class ScaleLinearKwargs(KwargsNode):
666 """key word arguments for `ScaleLinearDescr`"""
668 axes: Annotated[AxesInCZYX | None, Field(examples=["xy"])] = None
669 """The subset of axes to scale jointly.
670 For example xy to scale the two image axes for 2d data jointly."""
672 gain: float | list[float] = 1.0
673 """multiplicative factor"""
675 offset: float | list[float] = 0.0
676 """additive term"""
678 @model_validator(mode="after")
679 def either_gain_or_offset(self) -> Self:
680 if (
681 self.gain == 1.0
682 or isinstance(self.gain, list)
683 and all(g == 1.0 for g in self.gain)
684 ) and (
685 self.offset == 0.0
686 or isinstance(self.offset, list)
687 and all(off == 0.0 for off in self.offset)
688 ):
689 raise ValueError(
690 "Redunt linear scaling not allowd. Set `gain` != 1.0 and/or `offset` !="
691 + " 0.0."
692 )
694 return self
697class ScaleLinearDescr(NodeWithExplicitlySetFields):
698 """Fixed linear scaling."""
700 implemented_name: ClassVar[Literal["scale_linear"]] = "scale_linear"
701 if TYPE_CHECKING:
702 name: Literal["scale_linear"] = "scale_linear"
703 else:
704 name: Literal["scale_linear"]
706 kwargs: ScaleLinearKwargs
709class SigmoidDescr(NodeWithExplicitlySetFields):
710 """The logistic sigmoid funciton, a.k.a. expit function."""
712 implemented_name: ClassVar[Literal["sigmoid"]] = "sigmoid"
713 if TYPE_CHECKING:
714 name: Literal["sigmoid"] = "sigmoid"
715 else:
716 name: Literal["sigmoid"]
718 @property
719 def kwargs(self) -> KwargsNode:
720 """empty kwargs"""
721 return KwargsNode()
724class ZeroMeanUnitVarianceKwargs(KwargsNode):
725 """key word arguments for `ZeroMeanUnitVarianceDescr`"""
727 mode: Literal["fixed", "per_dataset", "per_sample"] = "fixed"
728 """Mode for computing mean and variance.
729 | mode | description |
730 | ----------- | ------------------------------------ |
731 | fixed | Fixed values for mean and variance |
732 | per_dataset | Compute for the entire dataset |
733 | per_sample | Compute for each sample individually |
734 """
735 axes: Annotated[AxesInCZYX, Field(examples=["xy"])]
736 """The subset of axes to normalize jointly.
737 For example `xy` to normalize the two image axes for 2d data jointly."""
739 mean: Annotated[
740 float | NotEmpty[list[float]] | None, Field(examples=[(1.1, 2.2, 3.3)])
741 ] = None
742 """The mean value(s) to use for `mode: fixed`.
743 For example `[1.1, 2.2, 3.3]` in the case of a 3 channel image with `axes: xy`."""
744 # todo: check if means match input axes (for mode 'fixed')
746 std: Annotated[
747 float | NotEmpty[list[float]] | None, Field(examples=[(0.1, 0.2, 0.3)])
748 ] = None
749 """The standard deviation values to use for `mode: fixed`. Analogous to mean."""
751 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
752 """epsilon for numeric stability: `out = (tensor - mean) / (std + eps)`."""
754 @model_validator(mode="after")
755 def mean_and_std_match_mode(self) -> Self:
756 if self.mode == "fixed" and (self.mean is None or self.std is None):
757 raise ValueError("`mean` and `std` are required for `mode: fixed`.")
758 elif self.mode != "fixed" and (self.mean is not None or self.std is not None):
759 raise ValueError(f"`mean` and `std` not allowed for `mode: {self.mode}`")
761 return self
764class ZeroMeanUnitVarianceDescr(NodeWithExplicitlySetFields):
765 """Subtract mean and divide by variance."""
767 implemented_name: ClassVar[Literal["zero_mean_unit_variance"]] = (
768 "zero_mean_unit_variance"
769 )
770 if TYPE_CHECKING:
771 name: Literal["zero_mean_unit_variance"] = "zero_mean_unit_variance"
772 else:
773 name: Literal["zero_mean_unit_variance"]
775 kwargs: ZeroMeanUnitVarianceKwargs
778class ScaleRangeKwargs(KwargsNode):
779 """key word arguments for `ScaleRangeDescr`
781 For `min_percentile`=0.0 (the default) and `max_percentile`=100 (the default)
782 this processing step normalizes data to the [0, 1] intervall.
783 For other percentiles the normalized values will partially be outside the [0, 1]
784 intervall. Use `ScaleRange` followed by `ClipDescr` if you want to limit the
785 normalized values to a range.
786 """
788 mode: Literal["per_dataset", "per_sample"]
789 """Mode for computing percentiles.
790 | mode | description |
791 | ----------- | ------------------------------------ |
792 | per_dataset | compute for the entire dataset |
793 | per_sample | compute for each sample individually |
794 """
795 axes: Annotated[AxesInCZYX, Field(examples=["xy"])]
796 """The subset of axes to normalize jointly.
797 For example xy to normalize the two image axes for 2d data jointly."""
799 min_percentile: Annotated[int | float, Interval(ge=0, lt=100)] = 0.0
800 """The lower percentile used to determine the value to align with zero."""
802 max_percentile: Annotated[int | float, Interval(gt=1, le=100)] = 100.0
803 """The upper percentile used to determine the value to align with one.
804 Has to be bigger than `min_percentile`.
805 The range is 1 to 100 instead of 0 to 100 to avoid mistakenly
806 accepting percentiles specified in the range 0.0 to 1.0."""
808 @model_validator(mode="after")
809 def min_smaller_max(self, info: ValidationInfo) -> Self:
810 if self.min_percentile >= self.max_percentile:
811 raise ValueError(
812 f"min_percentile {self.min_percentile} >= max_percentile"
813 + f" {self.max_percentile}"
814 )
816 return self
818 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
819 """Epsilon for numeric stability.
820 `out = (tensor - v_lower) / (v_upper - v_lower + eps)`;
821 with `v_lower,v_upper` values at the respective percentiles."""
823 reference_tensor: TensorName | None = None
824 """Tensor name to compute the percentiles from. Default: The tensor itself.
825 For any tensor in `inputs` only input tensor references are allowed.
826 For a tensor in `outputs` only input tensor refereences are allowed if `mode: per_dataset`"""
829class ScaleRangeDescr(NodeWithExplicitlySetFields):
830 """Scale with percentiles."""
832 implemented_name: ClassVar[Literal["scale_range"]] = "scale_range"
833 if TYPE_CHECKING:
834 name: Literal["scale_range"] = "scale_range"
835 else:
836 name: Literal["scale_range"]
838 kwargs: ScaleRangeKwargs
841class ScaleMeanVarianceKwargs(KwargsNode):
842 """key word arguments for `ScaleMeanVarianceDescr`"""
844 mode: Literal["per_dataset", "per_sample"]
845 """Mode for computing mean and variance.
846 | mode | description |
847 | ----------- | ------------------------------------ |
848 | per_dataset | Compute for the entire dataset |
849 | per_sample | Compute for each sample individually |
850 """
852 reference_tensor: TensorName
853 """Name of tensor to match."""
855 axes: Annotated[AxesInCZYX | None, Field(examples=["xy"])] = None
856 """The subset of axes to scale jointly.
857 For example xy to normalize the two image axes for 2d data jointly.
858 Default: scale all non-batch axes jointly."""
860 eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
861 """Epsilon for numeric stability:
862 "`out = (tensor - mean) / (std + eps) * (ref_std + eps) + ref_mean."""
865class ScaleMeanVarianceDescr(NodeWithExplicitlySetFields):
866 """Scale the tensor s.t. its mean and variance match a reference tensor."""
868 implemented_name: ClassVar[Literal["scale_mean_variance"]] = "scale_mean_variance"
869 if TYPE_CHECKING:
870 name: Literal["scale_mean_variance"] = "scale_mean_variance"
871 else:
872 name: Literal["scale_mean_variance"]
874 kwargs: ScaleMeanVarianceKwargs
877PreprocessingDescr = Annotated[
878 Union[
879 BinarizeDescr,
880 ClipDescr,
881 ScaleLinearDescr,
882 SigmoidDescr,
883 ZeroMeanUnitVarianceDescr,
884 ScaleRangeDescr,
885 ],
886 Discriminator("name"),
887]
888PostprocessingDescr = Annotated[
889 Union[
890 BinarizeDescr,
891 ClipDescr,
892 ScaleLinearDescr,
893 SigmoidDescr,
894 ZeroMeanUnitVarianceDescr,
895 ScaleRangeDescr,
896 ScaleMeanVarianceDescr,
897 ],
898 Discriminator("name"),
899]
902class InputTensorDescr(TensorDescrBase):
903 data_type: Literal["float32", "uint8", "uint16"]
904 """For now an input tensor is expected to be given as `float32`.
905 The data flow in bioimage.io models is explained
906 [in this diagram.](https://docs.google.com/drawings/d/1FTw8-Rn6a6nXdkZ_SkMumtcjvur9mtIhRqLwnKqZNHM/edit)."""
908 shape: Annotated[
909 Sequence[int] | ParameterizedInputShape,
910 Field(
911 examples=[(1, 512, 512, 1), {"min": (1, 64, 64, 1), "step": (0, 32, 32, 0)}]
912 ),
913 ]
914 """Specification of input tensor shape."""
916 preprocessing: list[PreprocessingDescr] = Field(
917 default_factory=cast( # TODO: (py>3.8) use list[PreprocessingDesr]
918 Callable[[], list[PreprocessingDescr]], list
919 )
920 )
921 """Description of how this input should be preprocessed."""
923 @model_validator(mode="after")
924 def zero_batch_step_and_one_batch_size(self) -> Self:
925 bidx = self.axes.find("b")
926 if bidx == -1:
927 return self
929 if isinstance(self.shape, ParameterizedInputShape):
930 step = self.shape.step
931 shape = self.shape.min
932 if step[bidx] != 0:
933 raise ValueError(
934 "Input shape step has to be zero in the batch dimension (the batch"
935 + " dimension can always be increased, but `step` should specify how"
936 + " to increase the minimal shape to find the largest single batch"
937 + " shape)"
938 )
939 else:
940 shape = self.shape
942 if shape[bidx] != 1:
943 raise ValueError("Input shape has to be 1 in the batch dimension b.")
945 return self
947 @model_validator(mode="after")
948 def validate_preprocessing_kwargs(self) -> Self:
949 for p in self.preprocessing:
950 kwargs_axes = p.kwargs.get("axes")
951 if isinstance(kwargs_axes, str) and any(
952 a not in self.axes for a in kwargs_axes
953 ):
954 raise ValueError("`kwargs.axes` needs to be subset of `axes`")
956 return self
959class OutputTensorDescr(TensorDescrBase):
960 data_type: Literal[
961 "float32",
962 "float64",
963 "uint8",
964 "int8",
965 "uint16",
966 "int16",
967 "uint32",
968 "int32",
969 "uint64",
970 "int64",
971 "bool",
972 ]
973 """Data type.
974 The data flow in bioimage.io models is explained
975 [in this diagram.](https://docs.google.com/drawings/d/1FTw8-Rn6a6nXdkZ_SkMumtcjvur9mtIhRqLwnKqZNHM/edit)."""
977 @property
978 def dtype(self):
979 """alias for `data_type`"""
980 return self.data_type
982 shape: Sequence[int] | ImplicitOutputShape
983 """Output tensor shape."""
985 halo: Sequence[int] | None = None
986 """The `halo` that should be cropped from the output tensor to avoid boundary effects.
987 The `halo` is to be cropped from both sides, i.e. `shape_after_crop = shape - 2 * halo`.
988 To document a `halo` that is already cropped by the model `shape.offset` has to be used instead."""
990 postprocessing: list[PostprocessingDescr] = Field(
991 default_factory=cast(Callable[[], list[PostprocessingDescr]], list)
992 )
993 """Description of how this output should be postprocessed."""
995 @model_validator(mode="after")
996 def matching_halo_length(self) -> Self:
997 if self.halo and len(self.halo) != len(self.shape):
998 raise ValueError(
999 f"halo {self.halo} has to have same length as shape {self.shape}!"
1000 )
1002 return self
1004 @model_validator(mode="after")
1005 def validate_postprocessing_kwargs(self) -> Self:
1006 for p in self.postprocessing:
1007 kwargs_axes = p.kwargs.get("axes", "")
1008 if not isinstance(kwargs_axes, str):
1009 raise ValueError(f"Expected {kwargs_axes} to be a string")
1011 if any(a not in self.axes for a in kwargs_axes):
1012 raise ValueError("`kwargs.axes` needs to be subset of axes")
1014 return self
1017KnownRunMode = Literal["deepimagej"]
1020class RunMode(Node):
1021 name: Annotated[
1022 KnownRunMode | str, warn(KnownRunMode, "Unknown run mode '{value}'.")
1023 ]
1024 """Run mode name"""
1026 kwargs: dict[str, Any] = Field(
1027 default_factory=cast(Callable[[], dict[str, Any]], dict)
1028 )
1029 """Run mode specific key word arguments"""
1032class LinkedModel(Node):
1033 """Reference to a bioimage.io model."""
1035 id: Annotated[ModelId, Field(examples=["affable-shark", "ambitious-sloth"])]
1036 """A valid model `id` from the bioimage.io collection."""
1038 version_number: int | None = None
1039 """version number (n-th published version, not the semantic version) of linked model"""
1042def package_weights(
1043 value: Node, # Union[v0_4.WeightsDescr, v0_5.WeightsDescr]
1044 handler: SerializerFunctionWrapHandler,
1045 info: SerializationInfo,
1046):
1047 ctxt = packaging_context_var.get()
1048 if ctxt is not None and ctxt.weights_priority_order is not None:
1049 for wf in ctxt.weights_priority_order:
1050 w = getattr(value, wf, None)
1051 if w is not None:
1052 break
1053 else:
1054 raise ValueError(
1055 "None of the weight formats in `weights_priority_order`"
1056 + f" ({ctxt.weights_priority_order}) is present in the given model."
1057 )
1059 assert isinstance(w, Node), type(w)
1060 # construct WeightsDescr with new single weight format entry
1061 new_w = w.model_construct(**{k: v for k, v in w if k != "parent"})
1062 value = value.model_construct(None, **{wf: new_w})
1064 return handler(value)
1067class ModelDescr(GenericModelDescrBase):
1068 """Specification of the fields used in a bioimage.io-compliant RDF that describes AI models with pretrained weights.
1070 These fields are typically stored in a YAML file which we call a model resource description file (model RDF).
1071 """
1073 implemented_format_version: ClassVar[Literal["0.4.10"]] = "0.4.10"
1074 if TYPE_CHECKING:
1075 format_version: Literal["0.4.10"] = "0.4.10"
1076 else:
1077 format_version: Literal["0.4.10"]
1078 """Version of the bioimage.io model description specification used.
1079 When creating a new model always use the latest micro/patch version described here.
1080 The `format_version` is important for any consumer software to understand how to parse the fields.
1081 """
1083 implemented_type: ClassVar[Literal["model"]] = "model"
1084 if TYPE_CHECKING:
1085 type: Literal["model"] = "model"
1086 else:
1087 type: Literal["model"]
1088 """Specialized resource type 'model'"""
1090 id: ModelId | None = None
1091 """bioimage.io-wide unique resource identifier
1092 assigned by bioimage.io; version **un**specific."""
1094 authors: NotEmpty[ # pyright: ignore[reportGeneralTypeIssues] # make mandatory
1095 list[Author]
1096 ]
1097 """The authors are the creators of the model RDF and the primary points of contact."""
1099 documentation: Annotated[
1100 FileSource_package,
1101 Field(
1102 examples=[
1103 "https://raw.githubusercontent.com/bioimage-io/spec-bioimage-io/main/example_descriptions/models/unet2d_nuclei_broad/README.md",
1104 "README.md",
1105 ],
1106 ),
1107 ]
1108 """URL or relative path to a markdown file with additional documentation.
1109 The recommended documentation file name is `README.md`. An `.md` suffix is mandatory.
1110 The documentation should include a '[#[#]]# Validation' (sub)section
1111 with details on how to quantitatively validate the model on unseen data."""
1113 inputs: NotEmpty[list[InputTensorDescr]]
1114 """Describes the input tensors expected by this model."""
1116 license: Annotated[
1117 LicenseId | str,
1118 warn(LicenseId, "Unknown license id '{value}'."),
1119 Field(examples=["CC0-1.0", "MIT", "BSD-2-Clause"]),
1120 ]
1121 """A [SPDX license identifier](https://spdx.org/licenses/).
1122 We do notsupport custom license beyond the SPDX license list, if you need that please
1123 [open a GitHub issue](https://github.com/bioimage-io/spec-bioimage-io/issues/new/choose
1124 ) to discuss your intentions with the community."""
1126 name: Annotated[
1127 str,
1128 MinLen(1),
1129 warn(Annotated[str, MinLen(5)], "Name shorter than 5 characters.", INFO),
1130 warn(Annotated[str, MaxLen(64)], "Name longer than 64 characters.", INFO),
1131 ]
1132 """A human-readable name of this model.
1133 It should be no longer than 64 characters and only contain letter, number, underscore, minus or space characters."""
1135 outputs: NotEmpty[list[OutputTensorDescr]]
1136 """Describes the output tensors."""
1138 @field_validator("inputs", "outputs")
1139 @classmethod
1140 def unique_tensor_descr_names(
1141 cls, value: Sequence[InputTensorDescr | OutputTensorDescr]
1142 ) -> Sequence[InputTensorDescr | OutputTensorDescr]:
1143 unique_names = {str(v.name) for v in value}
1144 if len(unique_names) != len(value):
1145 raise ValueError("Duplicate tensor descriptor names")
1147 return value
1149 @model_validator(mode="after")
1150 def unique_io_names(self) -> Self:
1151 unique_names = {str(ss.name) for s in (self.inputs, self.outputs) for ss in s}
1152 if len(unique_names) != (len(self.inputs) + len(self.outputs)):
1153 raise ValueError("Duplicate tensor descriptor names across inputs/outputs")
1155 return self
1157 @model_validator(mode="after")
1158 def minimum_shape2valid_output(self) -> Self:
1159 tensors_by_name: dict[TensorName, InputTensorDescr | OutputTensorDescr] = {
1160 t.name: t for t in self.inputs + self.outputs
1161 }
1163 for out in self.outputs:
1164 if isinstance(out.shape, ImplicitOutputShape):
1165 ndim_ref = len(tensors_by_name[out.shape.reference_tensor].shape)
1166 ndim_out_ref = len(
1167 [scale for scale in out.shape.scale if scale is not None]
1168 )
1169 if ndim_ref != ndim_out_ref:
1170 expanded_dim_note = (
1171 " Note that expanded dimensions (`scale`: null) are not"
1172 + f" counted for {out.name}'sdimensionality here."
1173 if None in out.shape.scale
1174 else ""
1175 )
1176 raise ValueError(
1177 f"Referenced tensor '{out.shape.reference_tensor}' with"
1178 + f" {ndim_ref} dimensions does not match output tensor"
1179 + f" '{out.name}' with"
1180 + f" {ndim_out_ref} dimensions.{expanded_dim_note}"
1181 )
1183 min_out_shape = self._get_min_shape(out, tensors_by_name)
1184 if out.halo:
1185 halo = out.halo
1186 halo_msg = f" for halo {out.halo}"
1187 else:
1188 halo = [0] * len(min_out_shape)
1189 halo_msg = ""
1191 if any(s - 2 * h < 1 for s, h in zip(min_out_shape, halo)):
1192 raise ValueError(
1193 f"Minimal shape {min_out_shape} of output {out.name} is too"
1194 + f" small{halo_msg}."
1195 )
1197 return self
1199 @classmethod
1200 def _get_min_shape(
1201 cls,
1202 t: InputTensorDescr | OutputTensorDescr,
1203 tensors_by_name: dict[TensorName, InputTensorDescr | OutputTensorDescr],
1204 ) -> Sequence[int]:
1205 """output with subtracted halo has to result in meaningful output even for the minimal input
1206 see https://github.com/bioimage-io/spec-bioimage-io/issues/392
1207 """
1208 if isinstance(t.shape, Sequence):
1209 return t.shape
1210 elif isinstance(t.shape, ParameterizedInputShape):
1211 return t.shape.min
1212 elif isinstance(t.shape, ImplicitOutputShape):
1213 pass
1214 else:
1215 assert_never(t.shape)
1217 ref_shape = cls._get_min_shape(
1218 tensors_by_name[t.shape.reference_tensor], tensors_by_name
1219 )
1221 if None not in t.shape.scale:
1222 scale: Sequence[float, ...] = t.shape.scale # type: ignore
1223 else:
1224 expanded_dims = [idx for idx, sc in enumerate(t.shape.scale) if sc is None]
1225 new_ref_shape: list[int] = []
1226 for idx in range(len(t.shape.scale)):
1227 ref_idx = idx - sum(int(exp < idx) for exp in expanded_dims)
1228 new_ref_shape.append(1 if idx in expanded_dims else ref_shape[ref_idx])
1230 ref_shape = new_ref_shape
1231 assert len(ref_shape) == len(t.shape.scale)
1232 scale = [0.0 if sc is None else sc for sc in t.shape.scale]
1234 offset = t.shape.offset
1235 assert len(offset) == len(scale)
1236 return [int(rs * s + 2 * off) for rs, s, off in zip(ref_shape, scale, offset)]
1238 @model_validator(mode="after")
1239 def validate_tensor_references_in_inputs(self) -> Self:
1240 for t in self.inputs:
1241 for proc in t.preprocessing:
1242 if "reference_tensor" not in proc.kwargs:
1243 continue
1245 ref_tensor = proc.kwargs["reference_tensor"]
1246 if ref_tensor is not None and str(ref_tensor) not in {
1247 str(t.name) for t in self.inputs
1248 }:
1249 raise ValueError(f"'{ref_tensor}' not found in inputs")
1251 if ref_tensor == t.name:
1252 raise ValueError(
1253 f"invalid self reference for preprocessing of tensor {t.name}"
1254 )
1256 return self
1258 @model_validator(mode="after")
1259 def validate_tensor_references_in_outputs(self) -> Self:
1260 for t in self.outputs:
1261 for proc in t.postprocessing:
1262 if "reference_tensor" not in proc.kwargs:
1263 continue
1264 ref_tensor = proc.kwargs["reference_tensor"]
1265 if ref_tensor is not None and str(ref_tensor) not in {
1266 str(t.name) for t in self.inputs
1267 }:
1268 raise ValueError(f"{ref_tensor} not found in inputs")
1270 return self
1272 packaged_by: list[Author] = Field(
1273 default_factory=cast(Callable[[], list[Author]], list)
1274 )
1275 """The persons that have packaged and uploaded this model.
1276 Only required if those persons differ from the `authors`."""
1278 parent: LinkedModel | None = None
1279 """The model from which this model is derived, e.g. by fine-tuning the weights."""
1281 @field_validator("parent", mode="before")
1282 @classmethod
1283 def ignore_url_parent(cls, parent: Any):
1284 if isinstance(parent, dict):
1285 return None
1287 else:
1288 return parent
1290 run_mode: RunMode | None = None
1291 """Custom run mode for this model: for more complex prediction procedures like test time
1292 data augmentation that currently cannot be expressed in the specification.
1293 No standard run modes are defined yet."""
1295 sample_inputs: list[FileSource_package] = Field(
1296 default_factory=cast(Callable[[], list[FileSource]], list)
1297 )
1298 """URLs/relative paths to sample inputs to illustrate possible inputs for the model,
1299 for example stored as PNG or TIFF images.
1300 The sample files primarily serve to inform a human user about an example use case"""
1302 sample_outputs: list[FileSource_package] = Field(
1303 default_factory=cast(Callable[[], list[FileSource]], list)
1304 )
1305 """URLs/relative paths to sample outputs corresponding to the `sample_inputs`."""
1307 test_inputs: NotEmpty[
1308 list[Annotated[FileSource_package, WithSuffix(".npy", case_sensitive=True)]]
1309 ]
1310 """Test input tensors compatible with the `inputs` description for a **single test case**.
1311 This means if your model has more than one input, you should provide one URL/relative path for each input.
1312 Each test input should be a file with an ndarray in
1313 [numpy.lib file format](https://numpy.org/doc/stable/reference/generated/numpy.lib.format.html#module-numpy.lib.format).
1314 The extension must be '.npy'."""
1316 test_outputs: NotEmpty[
1317 list[Annotated[FileSource_package, WithSuffix(".npy", case_sensitive=True)]]
1318 ]
1319 """Analog to `test_inputs`."""
1321 timestamp: Datetime
1322 """Timestamp in [ISO 8601](#https://en.wikipedia.org/wiki/ISO_8601) format
1323 with a few restrictions listed [here](https://docs.python.org/3/library/datetime.html#datetime.datetime.fromisoformat)."""
1325 training_data: LinkedDataset | DatasetDescr | None = None
1326 """The dataset used to train this model"""
1328 weights: Annotated[WeightsDescr, WrapSerializer(package_weights)]
1329 """The weights for this model.
1330 Weights can be given for different formats, but should otherwise be equivalent.
1331 The available weight formats determine which consumers can use this model."""
1333 @model_validator(mode="before")
1334 @classmethod
1335 def _convert_from_older_format(
1336 cls, data: BioimageioYamlContent, /
1337 ) -> BioimageioYamlContent:
1338 convert_from_older_format(data)
1339 return data
1341 @staticmethod
1342 def get_batch_size(tensor_sizes: Mapping[TensorName, Mapping[AxesStr, int]]) -> int:
1343 batch_size = 1
1344 tensor_with_batchsize: TensorName | None = None
1345 for tid in tensor_sizes:
1346 for aid, s in tensor_sizes[tid].items():
1347 if aid != "b" or s == 1 or s == batch_size:
1348 continue
1350 if batch_size != 1:
1351 assert tensor_with_batchsize is not None
1352 raise ValueError(
1353 f"batch size mismatch for tensors '{tensor_with_batchsize}' ({batch_size}) and '{tid}' ({s})"
1354 )
1356 batch_size = s
1357 tensor_with_batchsize = tid
1359 return batch_size
1361 def get_output_tensor_sizes(
1362 self, input_sizes: Mapping[TensorName, Mapping[AxesStr, int]]
1363 ) -> dict[TensorName, dict[AxesStr, int]]:
1364 """backport of model.v0_5.ModelDescr.get_output_tensor_sizes()"""
1365 batch_size = self.get_batch_size(input_sizes)
1366 input_sizes_refs = {
1367 ipt.name: {
1368 ax: (ax, input_sizes[ipt.name][ax])
1369 for ax in ipt.axes
1370 if ax != "b" and "b" in input_sizes[ipt.name]
1371 }
1372 for ipt in self.inputs
1373 if ipt.name in input_sizes
1374 }
1375 output_sizes: dict[TensorName, dict[AxesStr, int]] = {
1376 out.name: {} for out in self.outputs
1377 }
1378 for out in self.outputs:
1379 for i, ax in enumerate(out.axes):
1380 if ax == "b":
1381 s = batch_size
1382 elif isinstance(out.shape, Sequence):
1383 s = out.shape[i]
1384 elif isinstance(out.shape, ImplicitOutputShape):
1385 if (
1386 out.shape.reference_tensor not in input_sizes_refs
1387 or ax not in input_sizes_refs[out.shape.reference_tensor]
1388 ):
1389 logger.warning(
1390 "Output axis {}.{} references input axis {}.{}, which is not present in the provided input sizes."
1391 + " Assuming the output is optional and not produced for this input.",
1392 out.name,
1393 ax,
1394 out.shape.reference_tensor,
1395 ax,
1396 )
1397 continue
1399 _ref_ax, ref_size = input_sizes_refs[out.shape.reference_tensor][ax]
1400 s = (
1401 ref_size * scale
1402 if (scale := out.shape.scale[i]) is not None
1403 else 0
1404 ) + 2 * out.shape.offset[i]
1405 assert s == int(s)
1406 s = int(s)
1407 else:
1408 assert_never(out.shape)
1410 output_sizes[out.name][ax] = s
1412 return output_sizes
1414 def get_input_test_arrays(self) -> list[NDArray[Any]]:
1415 data = [load_array(ipt) for ipt in self.test_inputs]
1416 assert all(isinstance(d, np.ndarray) for d in data)
1417 return data
1419 def get_output_test_arrays(self) -> list[NDArray[Any]]:
1420 data = [load_array(out) for out in self.test_outputs]
1421 assert all(isinstance(d, np.ndarray) for d in data)
1422 return data