Skip to content

Commit ad8d821

Browse files
authored
gguf : add tensor shape accessor (ggml-org#24405)
* gguf : add tensor shape accessors * gguf : return tensor shape as const int64_t * * gguf : remove n_dims accessor, keep only gguf_get_tensor_ne
1 parent 91c631b commit ad8d821

3 files changed

Lines changed: 19 additions & 6 deletions

File tree

ggml/include/gguf.h

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -125,12 +125,13 @@ extern "C" {
125125
// get ith C string from array with given key_id
126126
GGML_API const char * gguf_get_arr_str (const struct gguf_context * ctx, int64_t key_id, size_t i);
127127

128-
GGML_API int64_t gguf_get_n_tensors (const struct gguf_context * ctx);
129-
GGML_API int64_t gguf_find_tensor (const struct gguf_context * ctx, const char * name); // returns -1 if the tensor is not found
130-
GGML_API size_t gguf_get_tensor_offset(const struct gguf_context * ctx, int64_t tensor_id);
131-
GGML_API const char * gguf_get_tensor_name (const struct gguf_context * ctx, int64_t tensor_id);
132-
GGML_API enum ggml_type gguf_get_tensor_type (const struct gguf_context * ctx, int64_t tensor_id);
133-
GGML_API size_t gguf_get_tensor_size (const struct gguf_context * ctx, int64_t tensor_id);
128+
GGML_API int64_t gguf_get_n_tensors (const struct gguf_context * ctx);
129+
GGML_API int64_t gguf_find_tensor (const struct gguf_context * ctx, const char * name); // returns -1 if the tensor is not found
130+
GGML_API size_t gguf_get_tensor_offset(const struct gguf_context * ctx, int64_t tensor_id);
131+
GGML_API const char * gguf_get_tensor_name (const struct gguf_context * ctx, int64_t tensor_id);
132+
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
133+
GGML_API enum ggml_type gguf_get_tensor_type (const struct gguf_context * ctx, int64_t tensor_id);
134+
GGML_API size_t gguf_get_tensor_size (const struct gguf_context * ctx, int64_t tensor_id);
134135

135136
// removes key if it exists, returns id that the key had prior to removal (-1 if it didn't exist)
136137
GGML_API int64_t gguf_remove_key(struct gguf_context * ctx, const char * key);

ggml/src/gguf.cpp

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1186,6 +1186,11 @@ const char * gguf_get_tensor_name(const struct gguf_context * ctx, int64_t tenso
11861186
return ctx->info[tensor_id].t.name;
11871187
}
11881188

1189+
const int64_t * gguf_get_tensor_ne(const struct gguf_context * ctx, int64_t tensor_id) {
1190+
GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx));
1191+
return ctx->info[tensor_id].t.ne;
1192+
}
1193+
11891194
enum ggml_type gguf_get_tensor_type(const struct gguf_context * ctx, int64_t tensor_id) {
11901195
GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx));
11911196
return ctx->info[tensor_id].t.type;

tests/test-gguf.cpp

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -662,6 +662,13 @@ static bool handcrafted_check_tensors(const gguf_context * gguf_ctx, const unsig
662662
if (gguf_get_tensor_type(gguf_ctx, id) != type) {
663663
ok = false;
664664
}
665+
666+
const int64_t * ne = gguf_get_tensor_ne(gguf_ctx, id);
667+
for (int j = 0; j < GGML_MAX_DIMS; ++j) {
668+
if (ne[j] != shape[j]) {
669+
ok = false;
670+
}
671+
}
665672
} else {
666673
ok = false;
667674
continue;

0 commit comments

Comments
 (0)