প্রকল্প সম্পর্কে
JAX, Python এবং NumPy প্রোগ্রামের কম্পোজেবল ফাংশন ট্রান্সফরমেশনের জন্য একটি সিস্টেম প্রদান করে। এটি কম্পাইলেশনের জন্য XLA (Accelerated Linear Algebra) ব্যবহার করে GPU এবং TPU-সহ হার্ডওয়্যার এক্সিলারেটর জুড়ে নিউমেরিক্যাল কম্পিউটিং স্কেল করার জন্য ডিজাইন করা হয়েছে।
এর মূল সক্ষমতাগুলোর মধ্যে রয়েছে:
- অটোমেটিক ডিফারেনশিয়েশন: `jax.grad` ব্যবহার করে, এটি Python লুপ, ব্রাঞ্চ, রিকার্শন এবং ক্লোজারের মাধ্যমে রিভার্স-মোড এবং ফরওয়ার্ড-মোড ডিফারেনশিয়েশন সমর্থন করে, যা হায়ার-অর্ডার ডেরিভেটিভস পেতে সাহায্য করে।
- জাস্ট-ইন-টাইম (JIT) কম্পাইলেশন: `jax.jit` ব্যবহার করে, এটি এক্সিলারেটরে এক্সিকিউশন পারফরম্যান্স উন্নত করতে XLA-এর মাধ্যমে পিওর ফাংশনগুলোকে কম্পাইল করে।
- অটো-ভেক্টরাইজেশন: `jax.vmap` ব্যবহার করে, এটি লুপগুলোকে প্রিমিটিভ অপারেশনে পুশ করে অ্যারে অ্যাক্সিস বরাবর ফাংশনগুলোকে ম্যাপ করে, যার ফলে ম্যানুয়াল ব্যাচ ডাইমেনশন ম্যানেজমেন্টের প্রয়োজন হয় না।
- স্কেলিং এবং প্যারালালাইজেশন: এটি কম্পাইলার-ভিত্তিক অটোমেটিক প্যারালালাইজেশন, এক্সপ্লিসিট শার্ডিং এবং এক্সপ্লিসিট কালেক্টিভস সহ বিভিন্ন স্কেলিং মোড সমর্থন করে।
JAX লিনাক্স, ম্যাকওএস এবং উইন্ডোজসহ একাধিক প্ল্যাটফর্মের সাথে সামঞ্জস্যপূর্ণ, যা CPU, NVIDIA GPU, Google TPU, AMD GPU এবং অ্যাপল ও ইন্টেল GPU-এর জন্য পরীক্ষামূলক সমর্থন প্রদান করে।