         return TensorShape([Int64](dimensionsBufferPointer).map(Int.init))
     }
 
-    /// Returns the data type of the tensor with the provided name stored in the checkpoint.
-    @inlinable
-    public func dataType(ofTensorNamed name: String) -> TensorDataType {
-        TensorDataType(TF_CheckpointReaderGetVariableDataType(handle, name))
+    /// Returns the scalar type of the tensor with the provided name stored in the checkpoint.
+    public func scalarTypeOfTensor(named name: String) -> Any.Type {
+        let dataType = TensorDataType(TF_CheckpointReaderGetVariableDataType(handle, name))
+        switch dataType._cDataType {
+        case TF_BOOL: return Bool.self
+        case TF_INT8: return Int8.self
+        case TF_UINT8: return UInt8.self
+        case TF_INT16: return Int16.self
+        case TF_UINT16: return UInt16.self
+        case TF_INT32: return Int32.self
+        case TF_UINT32: return UInt32.self
+        case TF_INT64: return Int64.self
+        case TF_UINT64: return UInt64.self
+        case TF_BFLOAT16: return BFloat16.self
+        case TF_FLOAT: return Float.self
+        case TF_DOUBLE: return Double.self
+        case TF_STRING: return String.self
+        default: fatalError("Unhandled type: \(dataType)")
+        }
     }
 
     /// Loads and returns the value of the tensor with the provided name stored in the checkpoint.
-    @inlinable
-    public func load<Scalar: _TensorFlowDataTypeCompatible>(
-        tensorNamed name: String
+    public func loadTensor<Scalar: _TensorFlowDataTypeCompatible>(
+        named name: String
     ) -> ShapedArray<Scalar> {
         let pointer = TF_CheckpointReaderGetTensor(handle, name, status)
         checkOk(status)

// limitations under the License.

import CTensorFlow
import Foundation

/// A TensorFlow checkpoint file reader.
public class TensorFlowCheckpointReader {
    @usableFromInline internal let status: OpaquePointer
    @usableFromInline internal let handle: OpaquePointer

    /// URL of the checkpoint file.
    public let checkpointPath: URL

    /// Number of tensors stored in the checkpoint.
    public var tensorCount: Int { Int(TF_CheckpointReaderSize(handle)) }

    /// Names of the tensors stored in the checkpoint.
    public var tensorNames: [String] {
        (0..<tensorCount).map {
            String(cString: TF_CheckpointReaderGetVariable(handle, Int32($0)))

    /// Creates a new TensorFlow checkpoint reader.
    ///
    /// - Arguments:
    ///   - checkpointPath: Path to the checkpoint file.
    public init(checkpointPath: String) {
        self.status = TF_NewStatus()
        self.handle = TF_NewCheckpointReader(checkpointPath, status)
        checkOk(status)
        self.checkpointPath = checkpointPath
    }

    deinit {
        TF_DeleteStatus(status)
        TF_DeleteCheckpointReader(handle)
    }

    /// Returns `true` if the checkpoint contains a tensor with the provided name.
    public func containsTensor(named name: String) -> Bool {
        TF_CheckpointReaderHasTensor(handle, name) > 0
    }

    /// Returns the shape of the tensor with the provided name stored in the checkpoint.
    public func shapeOfTensor(named name: String) -> TensorShape {
        let rank = TF_CheckpointReaderGetVariableNumDims(handle, name)
        let dimensions = UnsafeMutablePointer<Int64>.allocate(capacity: Int(rank))
        defer { dimensions.deallocate() }
