Function make_tensor

onnx/helper.py:363–483

Make a TensorProto with specified arguments. If raw is False, this function will choose the corresponding proto field to store the values based on data_type. If raw is True, use "raw_data" proto field to store the values, and values should be of type bytes in this case. Args:

(
    name: str,
    data_type: int,
    dims: Sequence[int],
    vals: Sequence[int | float] | bytes | np.ndarray,
    raw: bool = False,
)

def make_tensor(
    name: str,
    data_type: int,
    dims: Sequence[int],
    vals: Sequence[int | float] | bytes | np.ndarray,
    raw: bool = False,
) -> TensorProto:
    """Make a TensorProto with specified arguments.  If raw is False, this
    function will choose the corresponding proto field to store the
    values based on data_type. If raw is True, use "raw_data" proto
    field to store the values, and values should be of type bytes in
    this case.

    Args:
        name: tensor name
        data_type: a value such as onnx.TensorProto.FLOAT
        dims: shape
        vals: values
        raw: if True, vals contains the serialized content of the tensor,
            otherwise, vals should be a list of values of the type defined by ``data_type``.

    Returns:
        TensorProto
    """
    tensor = TensorProto()
    tensor.data_type = data_type
    tensor.name = name
    tensor.dims.extend(dims)

    if data_type == TensorProto.STRING and raw:
        raise TypeError("Can not use raw_data to store string type.")

    np_dtype = tensor_dtype_to_np_dtype(data_type)

    if raw:
        # NumPy doesn't have INT2/INT4/FP4. It is packed in couples to UINT8 buffers.
        if data_type in {TensorProto.UINT4, TensorProto.INT4, TensorProto.FLOAT4E2M1}:
            expected_size_bytes = 0.5
        elif data_type in {TensorProto.UINT2, TensorProto.INT2}:
            expected_size_bytes = 0.25
        else:
            expected_size_bytes = np_dtype.itemsize
        expected_size_bytes *= math.prod(dims)
        expected_size_bytes = math.ceil(expected_size_bytes)
        if isinstance(vals, np.ndarray):
            if data_type in {
                TensorProto.INT4,
                TensorProto.UINT4,
                TensorProto.FLOAT4E2M1,
            }:
                vals = onnx.numpy_helper._pack_4bitx2(vals)
            elif data_type in {TensorProto.UINT2, TensorProto.INT2}:
                vals = onnx.numpy_helper._pack_2bitx4(vals)

            raw_data = onnx.numpy_helper.tobytes_little_endian(vals)
        elif isinstance(vals, bytes):
            raw_data = vals
        else: