Ngôn Ngữ Lập Trình¶

Ngôn ngữ lập trình là giao diện giữa ý định của con người và sự thực thi của máy móc. File này bao gồm các mô hình ngôn ngữ (language paradigms), hệ thống kiểu (type systems), chiến lược quản lý bộ nhớ, quy trình biên dịch (compilation pipeline), thông dịch và biên dịch JIT, các tính năng ngôn ngữ chính, ngôn ngữ chuyên dụng (domain-specific languages) và các đánh đổi trong thiết kế.

Mọi phần mềm, mọi mô hình ML, mọi hệ điều hành đều được viết bằng một ngôn ngữ lập trình. Nhưng có hàng trăm ngôn ngữ, mỗi loại với thế mạnh riêng. Tại sao? Bởi vì thiết kế ngôn ngữ liên quan đến những sự đánh đổi căn bản: hiệu năng so với an toàn, tính biểu đạt so với sự đơn giản, kiểm soát so với sự trừu tượng hóa. Hiểu những sự đánh đổi này giúp bạn chọn đúng công cụ cho công việc và hiểu các ràng buộc bạn đang làm việc trong đó.

Mô Hình Ngôn Ngữ¶

Một mô hình (paradigm) là một phong cách lập trình: một tập hợp các nguyên tắc hướng dẫn cách bạn tổ chức mã và suy nghĩ về vấn đề.
Lập trình mệnh lệnh (imperative) mô tả tính toán như một chuỗi các lệnh thay đổi trạng thái. "Gán x bằng 5. Cộng 3 vào x. Nếu x > 7, in ra." C, Python và Java mệnh lệnh ở cốt lõi. Mô hình tinh thần là một cỗ máy với bộ nhớ mà bạn sửa đổi từng bước.
Lập trình hướng đối tượng (OOP) tổ chức mã xung quanh các đối tượng: các gói dữ liệu (thuộc tính) và hành vi (phương thức). Các đối tượng tương tác bằng cách gửi thông điệp lẫn nhau. Các ý tưởng chính là đóng gói (encapsulation) (che giấu trạng thái nội bộ đằng sau một giao diện công khai), kế thừa (inheritance) (tạo lớp mới bằng cách mở rộng các lớp hiện có), và đa hình (polymorphism) (xử lý các kiểu khác nhau một cách thống nhất qua một giao diện dùng chung). Java, C++ và Python hỗ trợ OOP.
Lập trình hàm (functional programming - FP) coi tính toán như là việc đánh giá các hàm toán học. Các nguyên lý cốt lõi: bất biến (immutability) (dữ liệu không thay đổi sau khi được tạo ra), hàm thuần túy (pure functions) (đầu ra chỉ phụ thuộc vào đầu vào, không có hiệu ứng phụ), và hàm là công dân hạng nhất (first-class functions) (hàm là giá trị có thể được truyền như đối số, trả về từ hàm khác và lưu trong biến). Haskell là thuần hàm. Python, JavaScript và Scala hỗ trợ phong cách hàm.
Hàm thuần túy dễ lập luận, dễ kiểm tra và dễ song song hóa (không có trạng thái thay đổi dùng chung nghĩa là không có điều kiện đua). Đây là lý do các ý tưởng hàm ngày càng được sử dụng trong các hệ thống phân tán và pipeline dữ liệu. JAX (được dùng xuyên suốt cuốn sách này) mang tính hàm: jax.grad hoạt động vì các hàm của JAX là thuần túy.
Lập trình logic mô tả cái gì nên đúng, chứ không phải làm thế nào để tính toán nó. Bạn nêu các sự kiện và quy tắc, và runtime tìm ra các giải pháp. Prolog là ví dụ kinh điển: với "Socrates là người" và "mọi người đều phải chết," engine suy ra "Socrates phải chết." Lập trình logic được sử dụng trong các cơ sở tri thức AI và kiểm tra kiểu.
Hầu hết các ngôn ngữ hiện đại là đa mô hình (multi-paradigm): Python hỗ trợ phong cách mệnh lệnh, OOP và hàm. Rust hỗ trợ mệnh lệnh và hàm. Mô hình là một công cụ, không phải một tôn giáo.

Hệ Thống Kiểu¶

Một kiểu (type) phân loại các giá trị và xác định phép toán nào là hợp lệ. Số nguyên 3 và chuỗi "3" là các kiểu khác nhau: bạn có thể cộng các số nguyên nhưng không thể cộng các chuỗi (thực ra bạn có thể nối chuỗi, nhưng đó là một phép toán khác).
Kiểu tĩnh (Static typing): các kiểu được kiểm tra tại thời điểm biên dịch, trước khi chương trình chạy. Lỗi kiểu được phát hiện sớm. C, Java, Rust và Go là kiểu tĩnh. Bạn phải khai báo kiểu (hoặc trình biên dịch suy luận chúng):

let x: i32 = 5;     // Rust: x là số nguyên 32-bit
let y: f64 = 3.14;  // y là float 64-bit
// let z = x + y;    // lỗi biên dịch: không thể cộng i32 và f64

Kiểu động (Dynamic typing): các kiểu được kiểm tra tại thời điểm chạy, khi các phép toán thực sự được thực thi. Linh hoạt hơn nhưng lỗi kiểu chỉ xuất hiện khi mã chạy. Python, JavaScript và Ruby là kiểu động:

x = 5       # x là int (tạm thời)
x = "hello" # giờ x là string -- không có lỗi

Kiểu mạnh (Strong typing): ngôn ngữ ngăn chặn sự ép kiểu ngầm định. Python là kiểu mạnh: "3" + 5 gây ra TypeError. Kiểu yếu (Weak typing): ngôn ngữ âm thầm chuyển đổi kiểu. JavaScript là kiểu yếu: "3" + 5 cho ra "35" (số được ép thành chuỗi). C là kiểu yếu: bạn có thể cast một con trỏ thành số nguyên.
Suy luận kiểu (Type inference) cho phép trình biên dịch suy ra kiểu mà không cần chú thích rõ ràng:

let x = 5;        // compiler suy ra: i32
let y = x + 3.0;  // lỗi biên dịch: kiểu hỗn hợp, dù có suy luận

Generics (đa hình tham số) cho phép bạn viết mã làm việc với bất kỳ kiểu nào:

fn largest<T: PartialOrd>(list: &[T]) -> &T {
    let mut max = &list[0];
    for item in &list[1..] {
        if item > max { max = item; }
    }
    max
}
// Hoạt động với số nguyên, float, chuỗi -- bất kỳ kiểu nào hỗ trợ so sánh

Đối với ML: kiểu động của Python giúp thử nghiệm nhanh nhưng che giấu lỗi. Các hệ thống ML trong sản xuất ngày càng sử dụng type hints (def train(model: nn.Module, lr: float) -> float) và các công cụ phân tích tĩnh (mypy) để phát hiện lỗi trước khi triển khai. PyTorch và JAX dùng Python vì tính linh hoạt; TensorRT và ONNX Runtime dùng C++ vì hiệu năng.

Quản Lý Bộ Nhớ¶

Mọi chương trình đều cấp phát và giải phóng bộ nhớ. Cách quản lý việc này là một trong những quyết định thiết kế ngôn ngữ hệ quả nhất.

Bố cục bộ nhớ: ngăn xếp mở rộng xuống dưới từ địa chỉ cao, heap mở rộng lên trên, với mã và dữ liệu ở dưới cùng

Ngăn xếp (Stack) lưu trữ các biến địa phương và các khung gọi hàm. Việc cấp phát là tầm thường (di chuyển con trỏ ngăn xếp) và giải phóng là tự động (pop khung khi hàm trả về). Truy cập ngăn xếp nhanh vì nó luôn nằm trong cache. Nhưng ngăn xếp có kích thước cố định (thường 1-8 MB) và chỉ hỗ trợ cấp phát LIFO (vào sau ra trước).
Heap lưu trữ dữ liệu được cấp phát động (các đối tượng, mảng, chuỗi có kích thước không được biết tại thời điểm biên dịch). Cấp phát heap chậm hơn (cần tìm một khối trống) và yêu cầu giải phóng tường minh hoặc tự động. Heap có thể mở rộng để lấp đầy bộ nhớ có sẵn.
Quản lý bộ nhớ thủ công (C, C++): lập trình viên tường minh cấp phát (malloc) và giải phóng (free) bộ nhớ heap. Kiểm soát và hiệu năng tối đa, nhưng cực kỳ dễ lỗi:
- Use-after-free: truy cập bộ nhớ đã được giải phóng. Gây ra sập hoặc lỗ hổng bảo mật.
- Double free: giải phóng cùng một bộ nhớ hai lần. Làm hỏng cấu trúc dữ liệu nội bộ của bộ cấp phát.
- Rò rỉ bộ nhớ (Memory leak): cấp phát bộ nhớ nhưng không bao giờ giải phóng. Chương trình từ từ tiêu thụ hết RAM có sẵn.
Garbage collection (GC - thu gom rác): runtime tự động phát hiện và giải phóng bộ nhớ không còn được tham chiếu. Lập trình viên không bao giờ gọi free.
- Tracing GC (Java, Go, trình thu gom chu trình của Python): định kỳ duyệt tất cả các đối tượng có thể truy cập từ các "gốc" (biến ngăn xếp, biến toàn cục) và giải phóng các đối tượng không thể truy cập. Đơn giản nhưng gây ra tạm dừng GC (GC pauses): chương trình dừng lại trong khi bộ thu gom chạy. Các bộ thu gom hiện đại (GC đồng thời của Go, ZGC của Java) giảm thiểu thời gian tạm dừng xuống dưới mili giây.
- Đếm tham chiếu (Reference counting) (cơ chế chính của Python, Swift, Objective-C): mỗi đối tượng theo dõi có bao nhiêu tham chiếu trỏ đến nó. Khi số đếm giảm xuống 0, đối tượng được giải phóng ngay lập tức. Không có tạm dừng, nhưng không thể xử lý các chu trình (cycles) (A tham chiếu B, B tham chiếu A, cả hai có số đếm > 0 nhưng không cái nào có thể truy cập được). Python dùng một bộ phát hiện chu trình riêng để xử lý điều này.
Quyền sở hữu (Ownership) (Rust): trình biên dịch thực thi các quy tắc an toàn bộ nhớ tại thời điểm biên dịch, với chi phí runtime bằng không.
- Mỗi giá trị có chính xác một chủ sở hữu (owner). Khi chủ sở hữu ra khỏi phạm vi, giá trị bj hủy (giải phóng).
- Các giá trị có thể được mượn (borrowed) (tham chiếu) nhưng trình biên dịch thực thi: hoặc một tham chiếu có thể thay đổi HOẶC bất kỳ số lượng tham chiếu bất biến nào, không bao giờ đồng thời cả hai.
- Điều này ngăn chặn use-after-free, double free, điều kiện đua dữ liệu (data races) và con trỏ treo (dangling pointers), tất cả tại thời điểm biên dịch. Không cần GC, không có chi phí runtime.
Bộ kiểm tra mượn (borrow checker) là tính năng đặc trưng của Rust và cũng là đường cong học tập dốc nhất của nó. Nó đảm bảo an toàn bộ nhớ và an toàn luồng mà không cần thu gom rác, đó là lý do Rust ngày càng được dùng cho các hệ thống quan trọng về hiệu năng (kernel hệ điều hành, game engine, runtime suy luận ML như Candle và Burn).

Quy Trình Biên Dịch¶

Một trình biên dịch (compiler) dịch mã nguồn thành mã máy (hoặc một ngôn ngữ đích khác) trước khi chương trình chạy. Quy trình có nhiều giai đoạn:

Quy trình biên dịch: nguồn → từ vùng (lexer) → phân tích cú pháp (parser) → phân tích ngữ nghĩa → tối ưu hóa → sinh mã → mã máy

Từ vùng (Lexing) (tokenisation): chuyển đổi văn bản nguồn thành một luồng các token. x = 3 + y trở thành [IDENT("x"), EQUALS, INT(3), PLUS, IDENT("y")]. Bộ từ vùng loại bỏ khoảng trắng và chú thích.
Phân tích cú pháp (Parsing): xây dựng một Cây cú pháp trừu tượng (AST) từ luồng token. AST biểu diễn cấu trúc phân cấp của chương trình. 3 + y * 2 được phân tích thành Add(3, Mul(y, 2)) (phép nhân có độ ưu tiên cao hơn). Bộ phân tích kiểm tra cú pháp: các dấu ngoặc không khớp và thiếu dấu chấm phẩy được phát hiện ở đây.
Phân tích ngữ nghĩa (Semantic analysis): kiểm tra kiểu, giải quyết tên biến, xác minh rằng các hàm được gọi với đối số đúng. Đây là nơi kiểm tra kiểu tĩnh diễn ra. Đầu ra là một AST đã được gán kiểu và chú thích.
Tối ưu hóa (Optimisation): biến đổi chương trình để chạy nhanh hơn mà không thay đổi hành vi của nó. Các tối ưu hóa phổ biến:
- Gấp hằng số (Constant folding): tính 3 + 5 tại thời điểm biên dịch, thay thế nó bằng 8.
- Loại bỏ mã chết (Dead code elimination): loại bỏ mã không bao giờ có thể thực thi.
- Mở vòng lặp (Loop unrolling): thay thế một vòng lặp bằng mã nội tuyến lặp lại để giảm chi phí rẽ nhánh.
- Nội tuyến hóa (Inlining): thay thế một lời gọi hàm bằng phần thân của hàm, loại bỏ chi phí gọi hàm.
Sinh mã (Code generation): dịch biểu diễn đã tối ưu hóa thành mã máy đích (x86, ARM) hoặc một biểu diễn trung gian.
LLVM là cơ sở hạ tầng trình biên dịch thống trị. Nó cung cấp một biểu diễn trung gian chung (LLVM IR) mà nhiều ngôn ngữ biên dịch xuống. Bộ tối ưu hóa của LLVM làm việc trên IR này, và các backend của nó sinh mã máy cho nhiều nền tảng. Clang (C/C++), Rust, Swift, Julia và nhiều ngôn ngữ khác dùng LLVM. Điều này có nghĩa là các cải tiến đối với bộ tối ưu hóa của LLVM đều có lợi cho tất cả các ngôn ngữ này cùng lúc.

Thông Dịch và Biên Dịch JIT¶

Một trình thông dịch (interpreter) thực thi chương trình từng dòng một (hoặc từng câu lệnh) mà không tạo ra mã máy. Điều này giúp khởi động nhanh và phát triển tương tác, nhưng thực thi chậm hơn (mỗi dòng được phân tích lại mỗi lần nó chạy).
Hầu hết các ngôn ngữ thông dịch thực sự biên dịch thành bytecode: một biểu diễn trung gian đơn giản hơn mã nguồn nhưng không dành riêng cho máy. Bytecode chạy trên một máy ảo (virtual machine - VM).
- CPython (hiện thực Python tiêu chuẩn) biên dịch mã nguồn Python thành bytecode (file .pyc), được thực thi bởi VM của CPython. VM thông dịch bytecode từng chỉ thị một. Đây là lý do Python chậm hơn C khoảng 100x cho mã tính toán nặng.
- JVM (Java Virtual Machine): Java biên dịch thành bytecode JVM (file .class). Ban đầu JVM thông dịch bytecode, sau đó JIT-biên dịch các đường dẫn mã thường xuyên được thực thi ("hot spots") thành mã máy gốc. Đây là lý do Java khởi động chậm hơn C (chi phí thông dịch) nhưng tiếp cận tốc độ của C cho các chương trình chạy lâu (các hot path đã được JIT tối ưu hóa).
Biên dịch JIT (Just-In-Time) biên dịch mã thành mã máy tại thời điểm chạy, sử dụng thông tin chỉ có sẵn trong quá trình thực thi. Một JIT có thể tối ưu hóa dựa trên dữ liệu runtime thực tế: nếu một hàm luôn được gọi với đối số nguyên, JIT sinh ra mã máy chuyên biệt chỉ dành cho số nguyên, bỏ qua các kiểm tra kiểu.
PyPy là một hiện thực Python thay thế với một trình biên dịch JIT. Nó chạy hầu hết mã Python nhanh hơn 5-10x so với CPython bằng cách JIT-biên dịch các hot loops thành mã máy. Tuy nhiên, nó có khả năng tương thích hạn chế với các module mở rộng C (NumPy, PyTorch), điều này hạn chế việc sử dụng nó trong ML.
Phổ từ thông dịch đến biên dịch không phải là nhị phân:
- Thông dịch thuần túy: script Bash shell.
- Thông dịch bytecode: CPython.
- Bytecode + JIT: JVM, .NET CLR, LuaJIT, PyPy.
- Biên dịch trước khi chạy (AOT): C, C++, Rust, Go.
- AOT + sinh mã runtime: jax.jit của JAX biên dịch các hàm Python thành mã XLA tối ưu hóa tại lần gọi đầu tiên, sau đó cache phiên bản đã biên dịch.

Các Tính Năng Ngôn Ngữ Chính¶

Closure (bao đóng): một hàm nắm bắt các biến từ phạm vi bao quanh nó. Hàm "đóng" môi trường nơi nó được định nghĩa:

def make_adder(n):
    def add(x):
        return x + n  # n được nắm bắt từ phạm vi bao quanh
    return add

add5 = make_adder(5)
print(add5(3))  # 8

Closures là cơ chế đằng sau các callback, decorator và ứng dụng một phần (partial application). Chúng là nền tảng của lập trình hàm.
So khớp mẫu (Pattern matching): một cơ chế điều khiển luồng mạnh mẽ phân rã dữ liệu và rẽ nhánh dựa trên hình dạng của nó:

match value {
    Some(x) if x > 0 => println!("Positive: {}", x),
    Some(0)           => println!("Zero"),
    Some(x)           => println!("Negative: {}", x),
    None              => println!("Nothing"),
}

Pattern matching biểu đạt hơn các chuỗi if-else: nó kiểm tra cấu trúc của dữ liệu (nó là Some hay None? nó có chứa giá trị khớp với một điều kiện nào đó không?), không chỉ sự bằng nhau. Python thêm pattern matching cấu trúc trong 3.10 (match/case).
Kiểu dữ liệu đại số (Algebraic data types - ADTs): các kiểu có thể là một trong nhiều biến thể, mỗi biến thể mang dữ liệu khác nhau. Kiểu Result hoặc là Ok(value) hoặc Err(error). Kiểu Tree hoặc là Leaf(value) hoặc Node(left, right). ADTs kết hợp với pattern matching cho phép xử lý triệt để tất cả các trường hợp, loại bỏ toàn bộ các lớp lỗi (null pointer exception, mã lỗi không được xử lý).
Traits và interfaces: định nghĩa một tập các phương thức mà một kiểu phải hiện thực, mà không chỉ rõ cách thức. Điều này cho phép đa hình: một hàm nhận "bất cứ thứ gì hiện thực trait Display" hoạt động với số nguyên, chuỗi và các kiểu tùy chỉnh. Rust dùng traits, Java dùng interfaces, Go dùng interfaces ngầm định, Python dùng duck typing ("nếu nó đi như con vịt...").

Ngôn Ngữ Chuyên Dụng (Domain-Specific Languages)¶

Một ngôn ngữ chuyên dụng (DSL) được thiết kế cho một miền bài toán cụ thể, đánh đổi tính tổng quát để lấy tính biểu đạt trong miền đó.
SQL: ngôn ngữ của cơ sở dữ liệu quan hệ. SELECT name FROM users WHERE age > 30 dễ đọc và tối ưu hóa hơn nhiều so với vòng lặp mệnh lệnh tương đương. Engine cơ sở dữ liệu tối ưu hóa kế hoạch thực thi truy vấn, tự động chọn chiến lược join và sử dụng chỉ mục.
Biểu thức chính quy (Regular expressions - RegEx): một mini-ngôn ngữ cho việc so khớp mẫu trong văn bản. \d{3}-\d{4} khớp với các số điện thoại như "555-1234." Các engine regex biên dịch các mẫu thành các automaton hữu hạn để so khớp hiệu quả.
Ngôn ngữ shader (GLSL, HLSL, Metal Shading Language): các chương trình chạy trên lõi GPU để tính màu pixel, vị trí đỉnh (vertex) hoặc các phép toán tính toán. Shader song song hàng loạt: mỗi lần gọi xử lý một pixel hoặc một phần tử độc lập. Đây là cùng mô hình thực thi được CUDA sử dụng cho tính toán ML.
Trong ML, các framework như PyTorch và JAX về cơ bản là các DSL cho tính toán tensor, nhúng trong Python. Chúng cung cấp các trừu tượng hóa (abstractions) chuyên biệt (tensor, tự động vi phân, đặt thiết bị (device placement)) trong khi tận dụng hệ sinh thái của Python.

Các Sự Đánh Đổi Trong Thiết Kế Ngôn Ngữ¶

Không có ngôn ngữ nào là tốt nhất cho mọi thứ. Thiết kế là về việc chọn những sự đánh đổi nào:
Hiệu năng so với an toàn: C cho tốc độ thô và kiểm soát phần cứng nhưng để bạn làm hỏng bộ nhớ. Rust cho tốc độ tương đương với an toàn bộ nhớ tại thời điểm biên dịch. Java cho an toàn bộ nhớ với chi phí thu gom rác. Python cho an toàn và biểu đạt tối đa với tốc độ chậm hơn 100x.
Tính biểu đạt so với đơn giản: Hệ thống kiểu của Haskell có thể biểu đạt các ràng buộc rất chính xác nhưng có đường cong học tập dốc. Go cố tình bỏ qua generics (cho đến gần đây), kế thừa và ngoại lệ để lấy sự đơn giản. Triết lý "nên có một cách rõ ràng để làm điều đó" của Python giữ cho ngôn ngữ có thể học được.
Kiểm soát so với trừu tượng hóa: C/C++ cho bạn kiểm soát bố cục bộ nhớ, hành vi cache và tương tác phần cứng. Python che giấu tất cả điều này. Đối với huấn luyện ML (nơi tính toán GPU chiếm ưu thế), chi phí của Python là không đáng kể. Đối với suy luận ML (nơi mọi micro giây đều quan trọng), C++ hoặc Rust có thể là cần thiết.
Tốc độ biên dịch so với tốc độ runtime: Go biên dịch trong vài giây (hệ thống kiểu đơn giản, tối ưu hóa tối thiểu). Rust biên dịch trong vài phút (hệ thống kiểu phức tạp, tối ưu hóa mạnh). Đánh đổi là tốc độ phát triển của lập trình viên so với hiệu năng triển khai.
Hệ sinh thái ML phản ánh những sự đánh đổi này: Python cho thử nghiệm và huấn luyện (tính biểu đạt thắng), C++/CUDA cho kernel và suy luận (hiệu năng thắng), Rust cho cơ sở hạ tầng và các hệ thống quan trọng về an toàn (an toàn thắng).

Bài Tập Lập Trình (dùng CoLab hoặc notebook)¶

Khám phá closures và hàm bậc cao (higher-order functions). Hiện thực một xưởng hàm (function factory) đơn giản và xác minh rằng closures nắm bắt môi trường của chúng.

def make_multiplier(factor):
    """Returns a function that multiplies by factor."""
    def multiply(x):
        return x * factor
    return multiply

double = make_multiplier(2)
triple = make_multiplier(3)

print(f"double(5) = {double(5)}")  # 10
print(f"triple(5) = {triple(5)}")  # 15

# Closures capture by reference, not by value
def make_counter():
    count = [0]  # mutable container to allow modification
    def increment():
        count[0] += 1
        return count[0]
    return increment

counter = make_counter()
print(f"counter() = {counter()}")  # 1
print(f"counter() = {counter()}")  # 2
print(f"counter() = {counter()}")  # 3

So sánh hành vi kiểu động và kiểu tĩnh. Chỉ ra cách kiểu động của Python cho phép sự linh hoạt nhưng có thể che giấu lỗi.

def add(a, b):
    return a + b

# Works with different types -- flexible!
print(add(3, 5))           # 8 (int + int)
print(add("hello ", "world"))  # "hello world" (str + str)
print(add([1, 2], [3, 4]))    # [1, 2, 3, 4] (list + list)

# But type errors only surface at runtime:
try:
    print(add("hello", 5))  # TypeError! str + int
except TypeError as e:
    print(f"Runtime error: {e}")
    print("A static type checker would catch this before running")

Đo sự khác biệt hiệu năng giữa Python thông dịch và các phương pháp biên dịch/JIT cho một tác vụ tính toán nặng.

import time
import jax
import jax.numpy as jnp

n = 1_000_000

# Pure Python loop (interpreted)
start = time.time()
total = 0.0
for i in range(n):
    total += i * i
python_time = time.time() - start

# JAX (compiled via XLA)
@jax.jit
def sum_squares_jax(n):
    return jnp.sum(jnp.arange(n, dtype=jnp.float32) ** 2)

_ = sum_squares_jax(10)  # warm up JIT
start = time.time()
result = sum_squares_jax(n)
jax_time = time.time() - start

print(f"Python loop: {python_time:.4f}s")
print(f"JAX (JIT):   {jax_time:.6f}s")
print(f"Speedup:     {python_time / jax_time:.0f}x")