diff --git a/ggml/include/gguf.h b/ggml/include/gguf.h index 67851ba6f1..b3a1e1230a 100644 --- a/ggml/include/gguf.h +++ b/ggml/include/gguf.h @@ -125,12 +125,13 @@ extern "C" { // get ith C string from array with given key_id GGML_API const char * gguf_get_arr_str (const struct gguf_context * ctx, int64_t key_id, size_t i); - GGML_API int64_t gguf_get_n_tensors (const struct gguf_context * ctx); - GGML_API int64_t gguf_find_tensor (const struct gguf_context * ctx, const char * name); // returns -1 if the tensor is not found - GGML_API size_t gguf_get_tensor_offset(const struct gguf_context * ctx, int64_t tensor_id); - GGML_API const char * gguf_get_tensor_name (const struct gguf_context * ctx, int64_t tensor_id); - GGML_API enum ggml_type gguf_get_tensor_type (const struct gguf_context * ctx, int64_t tensor_id); - GGML_API size_t gguf_get_tensor_size (const struct gguf_context * ctx, int64_t tensor_id); + GGML_API int64_t gguf_get_n_tensors (const struct gguf_context * ctx); + GGML_API int64_t gguf_find_tensor (const struct gguf_context * ctx, const char * name); // returns -1 if the tensor is not found + GGML_API size_t gguf_get_tensor_offset(const struct gguf_context * ctx, int64_t tensor_id); + GGML_API const char * gguf_get_tensor_name (const struct gguf_context * ctx, int64_t tensor_id); + GGML_API const int64_t * gguf_get_tensor_ne (const struct gguf_context * ctx, int64_t tensor_id); // returns ne, an array of GGML_MAX_DIMS elements; ne[dim] is 1 for dim >= n_dims + GGML_API enum ggml_type gguf_get_tensor_type (const struct gguf_context * ctx, int64_t tensor_id); + GGML_API size_t gguf_get_tensor_size (const struct gguf_context * ctx, int64_t tensor_id); // removes key if it exists, returns id that the key had prior to removal (-1 if it didn't exist) GGML_API int64_t gguf_remove_key(struct gguf_context * ctx, const char * key); diff --git a/ggml/src/gguf.cpp b/ggml/src/gguf.cpp index c3ffa1a134..7920b8100b 100644 --- a/ggml/src/gguf.cpp +++ b/ggml/src/gguf.cpp @@ -1186,6 +1186,11 @@ const char * gguf_get_tensor_name(const struct gguf_context * ctx, int64_t tenso return ctx->info[tensor_id].t.name; } +const int64_t * gguf_get_tensor_ne(const struct gguf_context * ctx, int64_t tensor_id) { + GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx)); + return ctx->info[tensor_id].t.ne; +} + enum ggml_type gguf_get_tensor_type(const struct gguf_context * ctx, int64_t tensor_id) { GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx)); return ctx->info[tensor_id].t.type; diff --git a/tests/test-gguf.cpp b/tests/test-gguf.cpp index ddb1b4d948..2875dec806 100644 --- a/tests/test-gguf.cpp +++ b/tests/test-gguf.cpp @@ -662,6 +662,13 @@ static bool handcrafted_check_tensors(const gguf_context * gguf_ctx, const unsig if (gguf_get_tensor_type(gguf_ctx, id) != type) { ok = false; } + + const int64_t * ne = gguf_get_tensor_ne(gguf_ctx, id); + for (int j = 0; j < GGML_MAX_DIMS; ++j) { + if (ne[j] != shape[j]) { + ok = false; + } + } } else { ok = false; continue;