Scale SciPy with jax.shard_map - jax.shard_mapで SciPy をスケール

SciPyData2026でLT登壇した内容です. https://scipydata.connpass.com/event/364718/

Transcript

1. Scale SciPy with jax.shard_map jax.shard_mapで SciPy をスケール 2026.01.26
2. Self-Introduction © Sho Tanaka 3 tsho / 田中 翔 (Sho Tanaka) Snowflake, Lead Developer Advocate Demonstrating & Developing: ML/AI, Data projects Private - Google Developer Expert AI/ML - MLOps community organizer - AI/ML Advisor
3. 0. Context for Today’s Topic Let’s give it a try!!! © Sho Tanaka 3 MLOps/LLMOps professionalist however didn’t use JAX so far…
4. 0. Today’s Topic © Sho Tanaka 3 *MathJax, GPJax etc. are NUMFOCUS’s project. ** JAX 0.9 was released 3 days ago
5. 1. Shift from pmap © Sho Tanaka 3 Available for distributed and parallel execution since the initial JAX release. An API that requires specifying the number of devices upfront for sharding, which began to conflict with the JAX trends jax.pmap
6. 1. Shift from pmap to shard_map: (Recommend migration JAX 0.8+) © Sho Tanaka 3
7. 1. Shift to shard_map © Sho Tanaka 3 No need to define device num Always use global array, explicitly define divided num, easily extended to other API jax.shard_map Special page for migration of jap.pmap v0.8.0+
8. 2. Example: Distributed Gaussian Process Regressor © Sho Tanaka 3 f(x) = sin(x**5) + cos(x ** 2) + (0.5 * x)
9. © Sho Tanaka 3 pmap - Define distribution number from device (CPU/GPU/TPU） - Insert jax.pmap into as var of predict_pmap 2. pmap: Distributed Gaussian Process Regressor Prep - Define Kernel - Scipy: Cholesky decomposition - Scipy: Solve triangular - Scipy: α=(K**-1) y
10. 2. shard_map: Distributed Gaussian Process Regressor © Sho Tanaka 3 shard_map Prep, same as pmap - Define Kernel - Scipy: Cholesky decomposition - Scipy: Solve triangular - Scipy: α=(K**-1) y Decorator with jax.shard_map and calculates Gaussian process Really readable!
11. Thank you! © Sho Tanaka 2