Skip to main content
Utilities for inspecting AOT-compiled .pt2 model packages. Extracts target architecture information from the embedded ELF binaries (.cubin for CUDA, .so for CPU) without unpacking the full archive. Also reads model metadata (input/output shapes, dtypes) from the model_metadata.json embedded in the package.

​
PackageArchInfo

Architecture information extracted from a .pt2 package.

​
cuda_compute_capability

Returns the CUDA compute capability as ‘X.Y’ string.

​
cuda_gpu_description

Returns a human-readable description of the target GPU.

​
cpu_description

Returns a human-readable description of the target CPU.

​
get_package_arch_info

Extracts architecture info from a .pt2 package. Reads only the ELF headers of the first .cubin and .so found in the archive. Does not unpack the full files.
​
Path to the .pt2 model package.
Returns: Architecture information for the package. Raises:
  • BioptimusRuntimeError - If the package archive cannot be read.

​
check_device_compatibility

Checks if a .pt2 package is compatible with the target device.
​
Path to the .pt2 model package.
​
The CUDA device to validate against. Accepts a torch.device, an integer device index, or None. When None (or a bare "cuda" device with no index), the current CUDA device is used. This ensures the check inspects the GPU the worker was actually assigned rather than always device 0.
Returns: None if compatible, or a descriptive error message if not.

​
TensorSpec

Specification for a single model input or output tensor. Attributes:
  • name - Name of the tensor (e.g. "x", "histology_tiles").
  • shape - List of dimension sizes. Dynamic dimensions are represented as the string _DYNAMIC; fixed dimensions are integers.
  • dtype - Expected dtype as a string (e.g. "float32", "int64").
  • description - Human-readable description from the model metadata.

​
fixed_shape

Returns a mapping of axis index to fixed size for non-dynamic dims.

​
ndim

Returns the expected number of dimensions.

​
shape_description

Returns a human-readable shape string like ‘(B, 3, 512, 512)’. Returns: Parenthesised comma-separated dimension description.

​
PackageMetadata

Model metadata extracted from a .pt2 package. Attributes:
  • model - Model class name (e.g. "EfficientUnet").
  • inputs - Ordered list of input tensor specifications.
  • outputs - Ordered list of output tensor specifications.
  • raw - The full raw metadata dictionary.

​
load_package_metadata

Loads model metadata from a .pt2 package. Reads the model_metadata.json file from the archive without extracting the full package.
​
Path to the .pt2 model package.
Returns: Parsed metadata with input/output tensor specifications. Raises:
  • BioptimusRuntimeError - If the package archive or metadata JSON cannot be read.

​
validate_input_shapes

Validates tensor shapes against the package metadata. Checks that:
  • All expected inputs are present.
  • Each tensor has the correct number of dimensions.
  • Fixed dimensions match exactly.
  • Dynamic dimensions are consistent across inputs (same batch size).
​
The package metadata with input specifications.
​
Mapping of input name to tensor.
Returns: None if all shapes are valid, or a descriptive error message.