🤖 AI

torch-sla: Differentiable Sparse Linear Algebra with Adjoint Solvers and Sparse Tensor Parallelism for PyTorch

यह शोध पत्र torch-sla को प्रस्तुत करता है, जो एक ओपन-सोर्स PyTorch लाइब्रेरी है जो विविध CPU और GPU बैकएंड्स पर विभिन्न सॉल्वर्स के लिए एक एकल autograd-जागरूक API प्रदान करके डिफरेंशिएबल स्पार्स लीनियर अलजेब्रा को एकीकृत करती है, जबकि यह बैच ऑपरेशन्स और O(1)-ग्राफ एडजॉइंट फ्रेमवर्क के माध्यम से स्केलेबल डिस्ट्रिब्यूटेड मल्टी-GPU निष्पादन का समर्थन करती है।

मूल लेखक: Mingyuan Chi, Shizheng Wen

प्रकाशित 2026-05-07

📖 7 मिनट में पढ़ें🧠 गहराई से पढ़ें

CC BY 4.0

मूल लेखक: Mingyuan Chi, Shizheng Wen

मूल पेपर CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) के तहत लाइसेंस किया गया है। ✨ नीचे दिए गए पेपर की यह व्याख्या AI से तैयार की गई है। इसे लेखकों ने न तो लिखा है, न इसका समर्थन किया है। तकनीकी सटीकता के लिए मूल पेपर देखें। पूरा डिस्क्लेमर पढ़ें

कल्पना कीजिए कि आप एक विशाल, जटिल पहेली को हल करने की कोशिश कर रहे हैं। कंप्यूटर विज्ञान और मशीन लर्निंग की दुनिया में, यह पहेली अक्सर एक "स्पार्स लीनियर सिस्टम" (sparse linear system) होती है—अंकों का एक विशाल ग्रिड जहाँ अधिकांश सेल खाली (शून्य) होते हैं, लेकिन कुछ भरे हुए सेल ही सब कुछ समझने की कुंजी होते हैं।

लंबे समय तक, लोकप्रिय प्रोग्रामिंग टूल PyTorch (जिसका उपयोग कई AI शोधकर्ता करते हैं) के पास एक समस्या थी: यह घने पज़ल्स (जहाँ हर सेल भरा होता है) को संभालने में बहुत अच्छा था, लेकिन इन विशाल, अधिकतर खाली पज़ल्स को हल करने में यह बहुत खराब था—खासकर तब जब इसे अपनी गलतियों से सीखना हो। यदि आप इसे हल करने की कोशिश करते, तो कंप्यूटर या तो क्रैश हो जाता, या अविश्वसनीय रूप से धीमा हो जाता, या फिर परिणाम से सीखना भूल जाता।

यहाँ torch-sla आता है, जो एक नया ओपन-सोर्स लाइब्रेरी है जो इन पहेलियों के लिए एक यूनिवर्सल ट्रांसलेटर और एक सुपर-एफिशिएंट मैकेनिक की तरह काम करता है। यह कैसे काम करता है, इसके लिए सरल उपमाओं का उपयोग किया गया है:

1. "स्मार्ट डिस्पैचर" (यूनिफाइड बैकएंड)

कल्पना कीजिए कि आप एक पैकेज लेकर डिलीवरी ड्राइवर हैं। आपको उसे एक गंतव्य तक पहुँचाना है, लेकिन सड़क की स्थितियाँ बदलती रहती हैं: कभी यह एक छोटा शहर की गली होती है, कभी एक विशाल हाईवे, और कभी आप एक अलग देश में होते हैं।

पहले, आपको हर एक यात्रा के लिए अलग वाहन मैन्युअल रूप से चुनना पड़ता था (छोटे कामों के लिए साइकिल, बड़े कामों के लिए ट्रक, पानी के लिए नाव)। यदि आपने गलत वाहन चुना, तो आपने समय बर्बाद किया या आप फंस गए।

torch-sla एक स्मार्ट डिस्पैचर की तरह है। आप बस कहते हैं, "मुझे इस पहेली को हल करना है," और यह स्वचालित रूप से जाँच करता है:

आप कहाँ हैं? (क्या आपका कंप्यूटर CPU का उपयोग कर रहा है या एक शक्तिशाली ग्राफिक्स कार्ड/GPU का?)
पहेली कितनी बड़ी है? (क्या यह 100 टुकड़ों वाली छोटी पहेली है या 10 करोड़ टुकड़ों वाली एक विशाल पहेली?)
इसका आकार क्या है? (क्या यह पूरी तरह से सममित/सिमेट्रिकल है?)

इन सुरागों के आधार पर, यह तुरंत पांच अलग-अलग विकल्पों के गैरेज (जिसमें SciPy, NVIDIA और स्वयं PyTorch के टूल्स शामिल हैं) में से सबसे अच्छा "वाहन" (सॉल्वर) चुन लेता है। आपको यह जानने की ज़रूरत नहीं है कि किसे चुनना है; लाइब्रेरी आपके लिए यह काम करती है।

2. "जादुई शॉर्टकट" (एडजॉइंट डिफरेंशिएशन)

यह इस लाइब्रेरी का सबसे जादुई करतब है।

पुराना तरीका (नेइव बैकप्रोपैगेशन):
कल्पना कीजिए कि आप एक अंधेरी भूलभुलैया (maze) में चल रहे हैं, बाहर निकलने का रास्ता खोजने के लिए 1,000 कदम उठा रहे हैं। अपनी यात्रा से सीखने के लिए, आप हर एक कदम, अपने पैरों की सटीक स्थिति और हर क्षण के वायु दबाव को लिखने का निर्णय लेते हैं। जब आप अंत में पहुँचते हैं, तो आपके पास नोट्स की 1,000 पृष्ठों वाली एक नोटबुक होती है। यदि आप वापस जाकर देखना चाहते हैं कि आपने कहाँ गलती की, तो आपको उन 1,000 पृष्ठों को पलटना होगा। यदि भूलभुलैया बड़ी होती है, तो आपकी नोटबुक बहुत बड़ी हो जाएगी, और आपका कंप्यूटर उन नोट्स को रखने के लिए केवल मेमोरी (RAM) की कमी के कारण क्रैश हो जाएगा।

torch-sla का तरीका (एडजॉइंट सॉल्वर):
हर कदम को लिखने के बजाय, torch-sla इम्प्लिसिट फंक्शन थ्योरम नामक गणितीय सिद्धांत पर आधारित एक "जादुई शॉर्टकट" का उपयोग करता है।
इसे ऐसे सोचें: आप भूलभुलैया में एक बार चलते हैं। जब आप बाहर पहुँचते हैं, तो अपने कदमों को दोबारा ट्रेस करने के बजाय, आप अपने दिमाग में तुरंत "रिवर्स मैप" की गणना कर लेते हैं। आपको केवल शुरुआत, अंत और भूलभुलैया के नियमों को याद रखने की आवश्यकता है।

परिणाम: चाहे भूलभुलैया में 10 कदम लगे हों या 10,000, आपकी "नोटबुक" का आकार एक समान छोटा रहता है।
लाभ: शोध पत्र दिखाता है कि यह पुराने तरीके की तुलना में 195 गुना अधिक मेमोरी बचाता है। यह आपको एक एकल कंप्यूटर पर उन विशाल पज़ल्स को हल करने की अनुमति देता है जो पहले सिस्टम को क्रैश कर देते थे।

3. "टीम हडल" (डिस्ट्रीब्यूटेड सॉल्वर्स)

क्या होगा अगर पहेली इतनी बड़ी हो कि एक कंप्यूटर उसे संभाल ही न सके? आपको कंप्यूटरों की एक टीम की आवश्यकता है (जैसे दोस्तों का एक समूह एक विशाल भित्ति चित्र/म्यूरल पर काम कर रहा हो)।

आमतौर पर, जब दोस्त म्यूरल के अलग-अलग हिस्सों पर काम करते हैं, तो उन्हें यह सुनिश्चित करने के लिए आपस में चिल्लाना पड़ता है कि किनारे आपस में मिल रहे हैं या नहीं। कंप्यूटर की भाषा में, इसे "हेलो एक्सचेंज" (halo exchange) कहा जाता है।

समस्या: जब आप अंतिम चित्र से सीखने की कोशिश करते हैं, तो पुराने टूल्स यह नहीं जानते थे कि सुधारों को सही लोगों तक वापस कैसे "चिल्लाकर" भेजा जाए।
torch-sla समाधान: यह कंप्यूटरों को सुधारों को उलटे क्रम में पास करना सिखाता है। यदि कंप्यूटर A, कंप्यूटर B को डेटा का एक हिस्सा भेजता है, तो torch-sla यह सुनिश्चित करता है कि जब लर्निंग (सीखना) होती है, तो कंप्यूटर B सुधार को वापस कंप्यूटर A को बिल्कुल सही तरीके से भेजता है। यह लाइब्रेरी को तीन शक्तिशाली GPU पर 40 करोड़ अज्ञात (unknowns) तक स्केल करने की अनुमति देता है।

4. यह वास्तव में क्या कर सकता है?

शोध पत्र यह प्रदर्शित करता है कि यह लाइब्रेरी निम्नलिखित को संभाल सकती है:

डायरेक्ट सॉल्वर्स (Direct Solvers): मध्यम आकार की पहेलियों के लिए कोड को तुरंत क्रैक करना।
इटरेटिव सॉल्वर्स (Iterative Solvers): विशाल पहेलियों के लिए चरण-दर-चरण अनुमान लेना।
नॉनलीनियर और आइजनवैल्यू सॉल्वर्स (Nonlinear & Eigenvalue Solvers): उन पहेलियों को संभालना जहाँ नियम हल करते समय बदलते रहते हैं, या डेटा में विशिष्ट "कंपनों" (आइजनवैल्यू) को खोजना।
बैचिंग (Batching): एक ही समय में हजारों पहेलियों को हल करना (जैसे एक साथ 1,000 अलग-अलग भूलभुलैया को हल करना)।

वास्तविक दुनिया का परीक्षण

यह साबित करने के लिए कि यह काम करता है, लेखकों ने "इनवर्स कोएफिशिएंट लर्निंग" नामक एक परीक्षण चलाया।
कल्पना कीजिए कि आप एक अंतिम छाया (समाधान) देखते हैं और आप पता लगाना चाहते हैं कि किस वस्तु (इनपुट) ने वह छाया बनाई है।

उन्होंने छिपी हुई वस्तु के आकार को उसकी छाया को देखकर "सीखने" के लिए torch-sla का उपयोग किया।
लाइब्रेरी ने एक मिनट से भी कम समय में 99.77% सटीकता के साथ छिपी हुई वस्तु का सफलतापूर्वक पता लगा लिया।
महत्वपूर्ण बात यह है कि शोधकर्ता को यह सीखने की प्रक्रिया शुरू करने के लिए कोई जटिल गणितीय कोड नहीं लिखना पड़ा; उन्होंने बस एक सरल कमांड कॉल की: A.solve(f)।

सारांश

torch-sla PyTorch इकोसिस्टम में एक बड़े अंतर को भरता है। यह विशाल, स्पार्स गणितीय पहेलियों को हल करने के कठिन, मेमोरी-गहन कार्य को लेता है और इसे बनाता है:

स्वचालित: यह काम के लिए सबसे अच्छा टूल खुद चुनता है।
कुशल: यह मेमोरी बचाने के लिए "जादुई शॉर्टकट" का उपयोग करता है।
स्केलेबल: यह कंप्यूटरों की टीमों को निर्बाध रूप से मिलकर काम करने की अनुमति देता है।
सीखने योग्य: यह AI मॉडल को बिना टूटे इन जटिल गणितीय समस्याओं से सीधे सीखने की अनुमति देता है।

यह अनिवार्य रूप से वह लापता कड़ी है जो PyTorch उपयोगकर्ताओं को उन प्रकार की वैज्ञानिक और इंजीनियरिंग समस्याओं से निपटने में सक्षम बनाती है जो पहले अन्य, अधिक विशिष्ट प्रोग्रामिंग वातावरणों में ही संभव थीं।

तकनीकी सारांश: torch-sla

समस्या विवरण

डिफरेंशिएबल स्पार्स लीनियर अल्जेब्रा (Differentiable sparse linear algebra) वैज्ञानिक मशीन लर्निंग के लिए एक मौलिक आवश्यकता है, जिसमें ग्राफ न्यूरल नेटवर्क (GNNs), अनस्ट्रक्चर्ड मेश पर न्यूरल ऑपरेटर्स, और डिफरेंशिएबल सिमुलेशन पाइपलाइन शामिल हैं। हालाँकि, PyTorch इकोसिस्टम में इन वर्कलोड्स को सपोर्ट करने के लिए कोई एकीकृत लाइब्रेरी का अभाव है। नेटिव PyTorch सपोर्ट खंडित है: torch.sparse लो-लेवल कर्नेल और एक नॉन-डिफरेंशिएबल, केवल CPU-आधारित spsolve (एक SuperLU रैपर) प्रदान करता है, जबकि torch.linalg घने मैट्रिसेस (dense matrices) तक सीमित है।

मौजूदा वर्कअराउंड बड़े पैमाने या जटिल प्रशिक्षण परिदृश्यों के लिए अपर्याप्त हैं:

ग्राफ को तोड़ना (Breaking the Graph): बाहरी सॉल्वर (जैसे, cuDSS, PETSc) को कॉल करने से ऑटोग्राफ (autograd graph) टूट जाता है, जिससे उपयोगकर्ताओं को या तो ग्रेडिएंट छोड़ना पड़ता है या हर समस्या के लिए कस्टम torch.autograd.Function रैपर लिखने पड़ते हैं।
अकुशल अवकलन (Inefficient Differentiation): एक इटरेटिव सॉल्वर के $k$ इटरेशन के माध्यम से नाइवली बैकप्रोपैगेट करने से $O(k)$ -नोड का कम्प्यूटेशनल ग्राफ बन जाता है। $10^6$ डिग्री ऑफ फ्रीडम (DOF) वाली समस्या के लिए, 1,000 इटरेशन के साथ, इंटरमीडिएट वेक्टर्स को स्टोर करने के लिए लगभग 80 GB मेमोरी की आवश्यकता होगी, जो बड़े पैमाने के प्रशिक्षण को असंभव बना देता है।
स्केलेबिलिटी सीमाएँ: $10^7$ – $10^9$ अज्ञात (unknowns) वाली समस्याओं के लिए डिस्ट्रिब्यूटेड डोमेन डीकंपोजिशन की आवश्यकता होती है। मौजूदा PyTorch डिस्ट्रिब्यूटेड प्रिमिटिव्स स्पार्स-ऑपरेटर एडजॉइंट्स के बैकवर्ड पास के लिए आवश्यक ट्रांसपोज़्ड हेलो एक्सचेंज (transposed halo exchange) को नेटिव रूप से सपोर्ट नहीं करते हैं।
बैकएंड विखंडन (Backend Fragmentation): डायरेक्ट सॉल्वर (जैसे, cuDSS) छोटे प्रब्लम्स ( $<10^5$ DOF) के लिए तेज़ होते हैं लेकिन बड़े सिस्टम के लिए $O(n^{1.5})$ फिल-इन और मेमोरी थकावट का शिकार होते हैं। इटरेटिव सॉल्वर बेहतर स्केल करते हैं लेकिन छोटे प्रब्लम्स पर कर्नेल-लॉन्च ओवरहेड का सामना करते हैं। डिवाइस और समस्या के आकार के आधार पर विभिन्न हेट्रोजेनियस बैकएंड्स के बीच डिस्पैच करने के लिए कोई एकीकृत इंटरफ़ेस मौजूद नहीं है।

कार्यप्रणाली (Methodology)

लेखक torch-sla प्रस्तुत करते हैं, जो तीन मुख्य घटकों के माध्यम से इस अंतर को भरता है:

1. एकीकृत बैकएंड एब्स्ट्रैक्शन (Unified Backend Abstraction)

torch-sla एक एकल ऑटोग्राड-अवेयर API (.solve, .matvec, .eigsh, .det) एक्सपोज़ करता है जो डिवाइस और समस्या के आकार के आधार पर पांच इंटरचेंजेबल बैकएंड्स में डिस्पैच करता है:

CPU: SciPy (SuperLU, UMFPACK, CG, BiCGStab) और Eigen (CG, BiCGStab)।
GPU: NVIDIA cuDSS (LU, Cholesky, LDLT), CuPy (इटरेटिव और डायरेक्ट), और एक PyTorch-नेटिव इटरेटिव सॉल्वर।
ऑटो-डिस्पैच: सिस्टम स्वचालित रूप से बैकएंड चुनता है: छोटे-से-मध्यम समस्याओं (जैसे, GPU पर $<2 \times 10^6$ DOF) के लिए डायरेक्ट सॉल्वर और मेमोरी ओवरफ्लो से बचने के लिए बड़ी समस्याओं के लिए इटरेटिव सॉल्वर। सिमिट्री और पॉजिटिव-डेफिनेटनेस को अनुकूलन (जैसे, LU को Cholesky में अपग्रेड करना) के लिए डिटेक्ट किया जाता है।

लाइब्रेरी चार स्पार्स टेंसर प्रकारों का समर्थन करती है जो लोकैलिटी और बैचिंग द्वारा व्यवस्थित हैं: SparseTensor (सिंगल मैट्रिक्स/शेयर्ड पैटर्न), SparseTensorList (प्रत्येक बैच के लिए अलग पैटर्न), और उनके डिस्ट्रिब्यूटेड वेरिएंट DSparseTensor और DSparseTensorList|

2. एडजॉइंट डिफरेंशिएशन फ्रेमवर्क (Adjoint Differentiation Framework)

नाइव बैकप्रोपैगेशन के $O(k)$ मेमोरी और ग्राफ जटिलता से बचने के लिए, torch-sla एक इम्प्लिसिट फंक्शन थ्योरम दृष्टिकोण अपनाता है।

लीनियर सिस्टम्स: $Ax=b $के लिए, बैकवर्ड पास एक एकल एडजॉइंट सिस्टम$ A^\top \lambda = \partial L / \partial x $को हल करता है। मैट्रिक्स के सापेक्ष ग्रेडिएंट केवल स्पार्सिटी पैटर्न पर$ -\lambda_i x_j$ के रूप में कंप्यूट किया जाता है।
नॉनलीनियर सिस्टम्स: $F(u, \theta)=0$ के लिए, बैकवर्ड पास में एक एडजॉइंट लीनियर सॉल्व ( $J^\top \lambda = \partial L / \partial u$ ) और एक वेक्टर-जैकबियन प्रोडक्ट की आवश्यकता होती है, जो नॉनलीनियर इटरेशन की संख्या से स्वतंत्र है।
आइजनवैल्यू प्रॉब्लम्स: आइजनवैल्यू के लिए ग्रेडिएंट्स हेलमैन-फायनमैन प्रमेय ( $\partial \lambda / \partial A_{ij} = v_i v_j$ ) का उपयोग करते हैं, जिसके लिए स्पार्सिटी पैटर्न पर केवल एक आउटर प्रोडक्ट की आवश्यकता होती है।
जटिलता: यह फ्रेमवर्क ऑटोग्राफ को $O(1)$ नोड्स में और मेमोरी को $O(n + \text{nnz})$ में बदल देता है, जो सॉल्वर इटरेशन ( $k$ ) की संख्या से स्वतंत्र है।

3. हेलो एक्सचेंज के साथ डिस्ट्रिब्यूटेड लेयर (Distributed Layer with Halo Exchange)

मल्टी-GPU निष्पादन के लिए, torch-sla डोमेन डीकंपोजिशन लागू करता है जहाँ प्रत्येक प्रोसेस एक रो ब्लॉक (row block) और हेलो मेटाडेटा का स्वामित्व रखती है।

फॉरवर्ड पास: प्रोसेस लोकल स्पार्स मैट्रिक्स-वेक्टर (SpMV) ऑपरेशन्स से पहले हेलो वैल्यूज का आदान-प्रदान करते हैं।
बैकवर्ड पास: लाइब्रेरी एक ट्रांसपोज़्ड हेलो एक्सचेंज लागू करती है। जबकि फॉरवर्ड पास स्वामित्व वाले मानों को पड़ोसी हेलो में स्कैटर करता है, बैकवर्ड पास इनकमिंग हेलो ग्रेडिएंट्स को स्वामित्व वाले प्रोसेस के बाउंड्री एंट्रीज में वापस सम करता है।
कार्यान्वयन: फॉरवर्ड और बैकवर्ड दोनों संचारों को torch.distributed का उपयोग करके torch.autograd.Function प्रिमिटिव्स के रूप में रैप किया गया है, जो एंड-टू-एंड डिफरेंशिएबिलिटी सुनिश्चित करता है।

मुख्य योगदान (Key Contributions)

एकीकृत API: पांच बैकएंड्स (SciPy, Eigen, cuDSS, CuPy, PyTorch-native) में डायरेक्ट, इटरेटिव, नॉनलीनियर और आइजनवैल्यू सॉल्वर के लिए ऑटोग्राड-अवेयर इंटरफेस प्रदान करने वाली एक एकल लाइब्रेरी।
स्केलेबल एडजॉइंट फ्रेमवर्क: एक $O(1)$ -ग्राफ डिफरेंशिएशन विधि जो इटरेटिव सॉल्वर बैकप्रोपैगेशन की मेमोरी बाधा को समाप्त करती है, जिससे लाखों अज्ञात (unknowns) वाली समस्याओं पर प्रशिक्षण संभव होता है।
डिस्ट्रिब्यूटेड ऑटोग्राड: PyTorch में पहला डिस्ट्रिब्यूटेड स्पार्स सॉल्वर कार्यान्वयन जो ट्रांसपोज़्ड हेलो एक्सचेंज के माध्यम से एंड-टू-एंड डिफरेंशिएबिलिटी को नेटिव रूप से सपोर्ट करता है।
विस्तार क्षमता (Extensibility): एक बैकएंड एब्स्ट्रैक्शन जो एक सामान्य इंटरफ़ेस को लागू करके बाहरी लाइब्रेरीज़ (जैसे, PETSc, Trilinos, hypre) या लर्नड प्रीकंडीशनर्स को एकीकृत करने की अनुमति देता है।

प्रयोगात्मक परिणाम (Experimental Results)

लेखकों ने NVIDIA H200 GPUs का उपयोग करके 2D पॉइसन बेंचमार्क पर torch-sla का मूल्यांकन किया:

स्केलेबिलिटी: लाइब्रेरी एक सिंगल GPU पर 169M DOF और तीन GPUs पर 400M DOF तक स्केल करती है।
प्रदर्शन: 1M DOF की समस्या के लिए, torch-sla CG बैकएंड (190 ms), cuDSS (7.3 s) की तुलना में 38 $\times$ बेहतर प्रदर्शन करता है। 2M DOF से ऊपर, डायरेक्ट सॉल्वर मेमोरी खत्म (OOM) हो जाते हैं, जबकि इटरेटिव सॉल्वर लगभग लीनियर स्केलिंग बनाए रखते हैं।
मेमोरी दक्षता: 1,000 CG इटरेशन के लिए नाइव बैकप्रोपैगेशन बनाम एडजॉइंट मेथड की तुलना में, नाइव दृष्टिकोण को 64 GB की आवश्यकता थी (2,000 इटरेशन पर OOM), जबकि एडजॉइंट मेथड ने निरंतर **328 MB** का उपयोग बनाए रखा, जिससे 195 $\times$ मेमोरी रिडक्शन प्राप्त हुआ।
ग्रेडिएंट वेरिफिकेशन: नॉनलीनियर और आइजनवैल्यू समस्याओं के लिए ग्रेडिएंट्स को फाईनाइट डिफरेंस के विरुद्ध सत्यापित किया गया, जो $10^{-5}$ से कम का रिलेटिव एरर दिखाते हैं।
एंड-टू-एंड एप्लीकेशन: एक इनवर्स कोएफिशिएंट लर्निंग टास्क (पॉइसन समीकरण में स्थानिक रूप से परिवर्तनशील चालकता को रिकवर करना) एक मानक PyTorch Adam लूप का उपयोग करके एक मिनट से कम समय में 0.23% रिलेटिव L2 एरर के साथ सफलतापूर्वक अभिसरित (converge) हुआ।

महत्व (Significance)

पेपर का दावा है कि torch-sla PyTorch इकोसिस्टम में एक लंबे समय से चले आ रहे अंतर को भरता है, जो JAX के jax.lax.custom_linear_solve और jax.scipy.sparse.linalg के समान क्षमताएं प्रदान करता है, लेकिन एक एकीकृत PyTorch लाइब्रेरी के भीतर। यह शोधकर्ताओं को सक्षम बनाता है:

बिना ग्रेडिएंट ड्रॉप किए या कोड को JAX में पोर्ट किए स्पार्स लीनियर सॉल्व शामिल करने वाले मॉडल को प्रशिक्षित करना।
वितरित GPU क्लस्टर्स पर सैकड़ों मिलियन डिग्री ऑफ फ्रीडम वाले वैज्ञानिक मशीन लर्निंग वर्कलोड को स्केल करना।
स्पार्स सॉल्वर्स को डिफरेंशिएबल लेयर्स के रूप में उपयोग करना, जिससे लर्नड प्रीकंडिशनर्स और फिजिकल सिमुलेशन के एंड-टू-एंड ऑप्टिमाइजेशन के अनुसंधान को सुगम बनाया जा सके।

लेखक सीमाओं का उल्लेख करते हैं, जिसमें वर्तमान में PyTorch-नेटिव बैकएंड के लिए जैकोबी प्रीकंडीशनिंग (Jacobi preconditioning) पर निर्भरता (जो बहुत बड़े सिस्टम के लिए कन्वर्जेंस स्पीड को सीमित करती है) और 2D पॉइसन बेंचमार्क पर ध्यान केंद्रित करना शामिल है, जिसे भविष्य के कार्य के रूप में 3D PDEs और इंडेफिनेट सिस्टम्स के साथ पहचाना गया है।