| 445 | |
| 446 | |
| 447 | def get_export_shapes(shapes): |
| 448 | export_shapes = {} |
| 449 | max_num_token = 0 |
| 450 | max_cache_size = 0 |
| 451 | for shape in shapes: |
| 452 | print(f"Shape: {shape}") |
| 453 | num_token = int(shape.split("t")[0]) |
| 454 | cache_size = int(shape.split("t")[1].split("c")[0]) |
| 455 | export_shapes[shape] = [num_token, cache_size] |
| 456 | max_num_token = num_token if num_token > max_num_token else max_num_token |
| 457 | max_cache_size = cache_size if cache_size > max_cache_size else max_cache_size |
| 458 | |
| 459 | return export_shapes, max_num_token, max_cache_size |
| 460 | |
| 461 | |
| 462 | def get_master_pos_emb(config, weight_dir, dtype, **kwargs): |