tensor.py

def to_literal(
    self,
    ctx: FlyteContext,
    python_val: T,
    python_type: Type[T],
    expected: LiteralType,
) -> Literal:
    meta = BlobMetadata(
        type=_core_types.BlobType(
            format=self.TENSORFLOW_FORMAT,
            dimensionality=_core_types.BlobType.BlobDimensionality.SINGLE,
        )
    )

    local_path = ctx.file_access.get_random_local_path()
    pathlib.Path(local_path).parent.mkdir(parents=True, exist_ok=True)

    # Save the `tf.tensor` as a file on disk
    filename = "tensor_data"
    tf.io.write_file(filename, tf.io.serialize_tensor(python_val))

    tensor_dtype = python_val.dtype.name

    remote_path = ctx.file_access.get_random_remote_path(local_path)
    ctx.file_access.put_data(local_path, remote_path, is_multipart=False)
    return Literal(
        collection=LiteralCollection(
            literals=[
                Literal(scalar=Scalar(blob=Blob(metadata=meta, uri=remote_path))),
                Literal(scalar=Scalar(primitive=Primitive(string=tensor_dtype))),
            ]
        )
    )


def to_python_value(
    self, ctx: FlyteContext, lv: Literal, python_val: T, expected_python_type: Type[T]
) -> T:
    try:
        uri = lv.collection.literals.scalar.blob.uri
    except AttributeError:
        TypeTransformerFailedError(
            f"Cannot convert from {lv} to {expected_python_type}"
        )

    local_path = ctx.file_access.get_random_local_path()
    ctx.file_access.get_data(uri, local_path, is_multipart=False)
    tensor_dtype = tf.dtypes.as_dtype(lv.collection.scalar.primitive.string)

    serialized_tensor = tf.io.read_file(local_path)
    actual_tensor = tf.io.parse_tensor(read_serial, out_type=tensor_dtype)

    return actual_tensor