1/*! For license information please see 9bec93b0.1e35a413.js.LICENSE.txt */ 2"use strict";(globalThis.webpackChunkblog||=[]).push([[4556],{8352(e,n,i){i.r(n),i.d(n,{assets:()=>h,contentTitle:()=>d,default:()=>p,frontMatter:()=>c,metadata:()=>a,toc:()=>m});var s=i.cw(function(e,n){const s=i.p+"assets/images/jeo_geeflow_processing-c89908630dc90acd7689a6be201a47a3.png";i.d(n,["A",0,s])}),a=i(92053),t=i(74848),r=i(28453),l=i(93),o=i(98979);const c={title:"Accelerating Science with JAX - Simulations, Physics, and Beyond",slug:"accelerating-science-with-jax-simulations-physics-and-beyond",description:"This post explores how JAX is reshaping modern scientific computing by enabling high-performance, differentiable, and hardware-accelerated workflows across a wide range of applications. From modeling complex systems and solving differential equations to training large-scale geospatial and machine learning models, we introduce powerful libraries like Diffrax, fdtdx, jax-md, Equinox, and Jeo. Learn how JAX supports scalable, composable research pipelines that span domains such as engineering, biology, Earth observation, and AI.",authors:["haruiz"],tags:["python","data-science","jax","scientific-computing","machine-learning","simulations","physics","electromagnetics","geosciences","earth-observation"]},d=void 0,h={authorsImageUrls:[void 0]},m=[{value:"1. Differentiable ODEs with <code>Diffrax</code>",id:"1-differentiable-odes-with-diffrax",level:2},{value:"Applications",id:"applications",level:3},{value:"Example: Solve",id:"example-solve",level:3},{value:"2. Electromagnetic Simulation with <code>fdtdx</code>",id:"2-electromagnetic-simulation-with-fdtdx",level:2},{value:"Applications",id:"applications-1",level:3},{value:"3. Molecular Simulation with <code>jax-md</code>",id:"3-molecular-simulation-with-jax-md",level:2},{value:"Applications",id:"applications-2",level:3},{value:"4. Scientific ML with <code>Equinox</code>",id:"4-scientific-ml-with-equinox",level:2},{value:"5. Model training and inference for geospatial remote sensing and Earth Observation with <code>Jeo</code>",id:"5-model-training-and-inference-for-geospatial-remote-sensing-and-earth-observation-with-jeo",level:2},{value:"Applications",id:"applications-3",level:2},{value:"Looking Ahead",id:"looking-ahead",level:2},{value:"Conclusion",id:"conclusion",level:2},{value:"References",id:"references",level:3}];function u(e){const n={a:"a",admonition:"admonition",annotation:"annotation",blockquote:"blockquote",code:"code",h2:"h2",h3:"h3",img:"img",li:"li",math:"math",mfrac:"mfrac",mi:"mi",mn:"mn",mo:"mo",mrow:"mrow",mspace:"mspace",mtext:"mtext",p:"p",pre:"pre",semantics:"semantics",span:"span",strong:"strong",ul:"ul",...(0,r.R)(),...e.components};return(0,t.jsxs)(t.Fragment,{children:[(0,t.jsx)(l.A,{audioSrc:i(44198).A}),"\n","\n",(0,t.jsxs)(n.p,{children:["After working hands-on with ",(0,t.jsx)(n.a,{href:"https://github.com/google/jax",children:"JAX"}),"\u2014particularly through libraries like ",(0,t.jsx)(n.a,{href:"https://github.com/ymahlau/fdtdx",children:(0,t.jsx)(n.code,{children:"fdtdx"})})," for accelerating electromagnetic simulations\u2014I've come to see it not just as a tool for deep learning, but as a core engine for modern scientific computing. What makes JAX stand out is its elegant fusion of NumPy-like syntax, automatic differentiation, and hardware-agnostic acceleration across CPUs, GPUs, and TPUs\u2014all wrapped in a functional programming model that encourages composability and clarity."]}),"\n",(0,t.jsx)(n.p,{children:"Beyond its foundational capabilities, high-level libraries like NNX (see example below) further extend JAX\u2019s flexibility, making it easy to build modular, expressive, and differentiable models with minimal boilerplate. This powerful ecosystem positions JAX as a go-to framework not only for machine learning, but also for developing high-performance simulations across disciplines such as physics, electromagnetics, and scientific optimization\u2014unlocking workflows that demand both numerical precision and differentiable programming."}),"\n",(0,t.jsx)(n.pre,{children:(0,t.jsx)(n.code,{className:"language-python",children:"from flax import nnx\nimport optax\n\n\nclass Model(nnx.Module):\n def __init__(self, din, dmid, dout, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dmid, rngs=rngs)\n self.bn = nnx.BatchNorm(dmid, rngs=rngs)\n self.dropout = nnx.Dropout(0.2, rngs=rngs)\n self.linear_out = nnx.Linear(dmid, dout, rngs=rngs)\n\n def __call__(self, x):\n x = nnx.relu(self.dropout(self.bn(self.linear(x))))\n return self.linear_out(x)\n\nmodel = Model(2, 64, 3, rngs=nnx.Rngs(0)) # eager initialization\noptimizer = nnx.Optimizer(model, optax.adam(1e-3)) # reference sharing\n\[email protected] # automatic state management for JAX transforms\ndef train_step(model, optimizer, x, y):\n def loss_fn(model):\n y_pred = model(x) # call methods directly\n return ((y_pred - y) ** 2).mean()\n\n loss, grads = nnx.value_and_grad(loss_fn)(model)\n optimizer.update(grads) # in-place updates\n\n return loss\n\nnnx.display(optimizer)\n"})}),"\n",(0,t.jsx)(n.p,{children:"Building on these capabilities, in this post I\u2019ll highlight some of the most exciting open-source projects built on JAX and demonstrate how they can be used in scientific computing to model complex physical systems\u2014ranging from solving ordinary and partial differential equations (ODEs/PDEs), to simulating molecular dynamics, electromagnetic wave propagation, fluid dynamics, and geophysical processes. These tools empower researchers to explore phenomena such as heat diffusion, quantum interactions, wave mechanics, subsurface imaging, and environmental transport models with high performance and end-to-e
2nd differentiability. I\u2019ll also discuss where these tools can be effectively applied across domains such as physics, engineering, geosciences, earth and environmental systems, and computational biology to optimize and accelerate research workflows."}),"\n",(0,t.jsx)(n.p,{children:"To run the code examples in this post, you\u2019ll need to install JAX. You can do this by running the following command in your terminal or in the environment where you plan to run the examples:"}),"\n",(0,t.jsx)(o.A,{lines:[{type:"input",value:'pip install -U "jax[cuda12]" -q'}]}),"\n",(0,t.jsxs)(n.admonition,{type:"tip",children:[(0,t.jsx)(n.p,{children:"Use the following code snippet to verify that JAX is installed correctly and to determine which hardware accelerator (CPU, GPU, or TPU) is being used:"}),(0,t.jsx)(n.pre,{children:(0,t.jsx)(n.code,{className:"language-python",children:"import jax\n\n# List all available devices\ndevices = jax.devices()\nprint(devices)\nprint(jax.default_backend())\nprint(jax.lib.xla_bridge.get_backend().platform)\n"})})]}),"\n",(0,t.jsxs)(n.h2,{id:"1-differentiable-odes-with-diffrax",children:["1. Differentiable ODEs with ",(0,t.jsx)(n.a,{href:"https://github.com/patrick-kidger/diffrax",children:(0,t.jsx)(n.code,{children:"Diffrax"})})]}),"\n",(0,t.jsxs)(n.p,{children:["JAX\u2019s powerful automatic differentiation system enables seamless integration of differential equations into loss functions and optimization routines. Building on this foundation, ",(0,t.jsx)(n.strong,{children:"Diffrax"})," provides state-of-the-art ODE solvers that are fully compatible with ",(0,t.jsx)(n.code,{children:"jax.grad"}),", enabling end-to-end differentiability. This extends JAX\u2019s capabilities far beyond neural network training, transforming it into a versatile tool for solving both ordinary differential equations (ODEs) and partial differential equations (PDEs) within a unified, gradient-based framework."]}),"\n",(0,t.jsx)(n.h3,{id:"applications",children:"Applications"}),"\n",(0,t.jsx)(n.p,{children:"Diffrax can be applied across a wide range of scientific and engineering domains, including:"}),"\n",(0,t.jsxs)(n.ul,{children:["\n",(0,t.jsxs)(n.li,{children:[(0,t.jsx)(n.strong,{children:"Physics-informed neural networks (PINNs)"})," \u2013 Solve ODEs and PDEs by incorporating physical laws directly into model training."]}),"\n",(0,t.jsxs)(n.li,{children:[(0,t.jsx)(n.strong,{children:"Parameter estimation and system identification"})," \u2013 Learn unknown parameters of dynamical systems using gradient-based optimization."]}),"\n",(0,t.jsxs)(n.li,{children:[(0,t.jsx)(n.strong,{children:"Optimal control and trajectory planning"})," \u2013 Optimize control inputs for physical systems such as robots or autonomous vehicles."]}),"\n",(0,t.jsxs)(n.li,{children:[(0,t.jsx)(n.strong,{children:"Biological and ecological modeling"})," \u2013 Fit and simulate population dynamics, gene regulation networks, or epidemiological models."]}),"\n",(0,t.jsxs)(n.li,{children:[(0,t.jsx)(n.strong,{children:"Financial modeling"})," \u2013 Solve stochastic differential equations for pricing, forecasting, or risk assessment."]}),"\n"]}),"\n",(0,t.jsx)(n.p,{children:"The following example demonstrates how to solve a simple ODE using Diffrax:"}),"\n",(0,t.jsx)(n.h3,{id:"example-solve",children:"Example: Solve"}),"\n",(0,t.jsx)(n.span,{className:"katex-display",children:(0,t.jsxs)(n.span,{className:"katex",children:[(0,t.jsx)(n.span,{className:"katex-mathml",children:(0,t.jsx)(n.math,{xmlns:"http://www.w3.org/1998/Math/MathML",display:"block",children:(0,t.jsxs)(n.semantics,{children:[(0,t.jsxs)(n.mrow,{children:[(0,t.jsx)(n.mi,{children:"y"}),(0,t.jsx)(n.mo,{stretchy:"false",children:"("}),(0,t.jsx)(n.mn,{children:"0"}),(0,t.jsx)(n.mo,{stretchy:"false",children:")"}),(0,t.jsx)(n.mo,{children:"="}),(0,t.jsx)(n.mn,{children:"1"}),(0,t.jsx)(n.mspace,{width:"1em"}),(0,t.jsxs)(n.mfrac,{children:[(0,t.jsxs)(n.mrow,{children:[(0,t.jsxs)(n.mrow,{children:[(0,t.jsx)(n.mtext,{children:"\xa0"}),(0,t.jsx)(n.mi,{mathvariant:"normal",children:"d"})]}),(0,t.jsx)(n.mi,{children:"y"})]}),(0,t.jsxs)(n.mrow,{children:[(0,t.jsxs)(n.mrow,{children:[(0,t.jsx)(n.mtext,{children:"\xa0"}),(0,t.jsx)(n.mi,{mathvariant:"normal",children:"d"})]}),(0,t.jsx)(n.mi,{children:"t"})]})]}),(0,t.jsx)(n.mo,{stretchy:"false",children:"("}),(0,t.jsx)(n.mi,{children:"t"}),(0,t.jsx)(n.mo,{stretchy:"false",children:")"}),(0,t.jsx)(n.mo,{children:"="}),(0,t.jsx)(n.mo,{children:"\u2212"}),(0,t.jsx)(n.mi,{children:"y"}),(0,t.jsx)(n.mo,{stretchy:"false",children:"("}),(0,t.jsx)(n.mi,{children:"t"}),(0,t.jsx)(n.mo,{stretchy:"false",children:")"})]}),(0,t.jsx)(n.annotation,{encoding:"application/x-tex",children:"y(0)=1 \\quad \\frac{\\mathrm{~d} y}{\\mathrm{~d} t}(t)=-y(t)"})]})})}),(0,t.jsxs)(n.span,{className:"katex-html","aria-hidden":"true",children:[(0,t.jsxs)(n.span,{className:"base",children:[(0,t.jsx)(n.span,{className:"strut",style:{height:"1em",verticalAlign:"-0.25em"}}),(0,t.jsx)(n.span,{className:"mord mathnormal",style:{marginRight:"0.0359em"},children:"y"}),(0,t.jsx)(n.span,{className:"mopen",children:"("}),(0,t.jsx)(n.span,{className:"mord",children:"0"}),(0,t.jsx)(n.span,{className:"mclose",children:")"}),(0,t.jsx)(n.span,{className:"mspace",style:{marginRight:"0.2778em"}}),(0,t.jsx)(n.span,{className:"mrel",children:"="}),(0,t.jsx)(n.span,{className:"mspace",style:{marginRight:"0.2778em"}})]}),(0,t.jsxs)(n.span,{className:"base",children:[(0,t.jsx)(n.span,{className:"strut",style:{height:"2.0574em",verticalAlign:"-0.686em"}}),(0,t.jsx)(n.span,{className:"mord",children:"1"}),(0,t.jsx)(n.span,{className:"mspace",style:{marginRight:"1em"}}),(0,t.jsxs)(n.span,{className:"mord",children:[(0,t.jsx)(n.span,{className:"mopen nulldelimiter"}),(0,t.jsx)(n.span,{className:"mfrac",children:(0,t.jsxs)(n.span,{className:"vlist-t vlist-t2",children:[(0,t.jsxs)(n.span,{className:"vlist-r",children:[(0,t.jsxs)(n.span,{className:"vlist",style:{height:"1.3714em"}
2,children:[(0,t.jsxs)(n.span,{style:{top:"-2.314em"},children:[(0,t.jsx)(n.span,{className:"pstrut",style:{height:"3em"}}),(0,t.jsxs)(n.span,{className:"mord",children:[(0,t.jsxs)(n.span,{className:"mord",children:[(0,t.jsx)(n.span,{className:"mspace nobreak",children:"\xa0"}),(0,t.jsx)(n.span,{className:"mord mathrm",children:"d"})]}),(0,t.jsx)(n.span,{className:"mord mathnormal",children:"t"})]})]}),(0,t.jsxs)(n.span,{style:{top:"-3.23em"},children:[(0,t.jsx)(n.span,{className:"pstrut",style:{height:"3em"}}),(0,t.jsx)(n.span,{className:"frac-line",style:{borderBottomWidth:"0.04em"}})]}),(0,t.jsxs)(n.span,{style:{top:"-3.677em"},children:[(0,t.jsx)(n.span,{className:"pstrut",style:{height:"3em"}}),(0,t.jsxs)(n.span,{className:"mord",children:[(0,t.jsxs)(n.span,{className:"mord",children:[(0,t.jsx)(n.span,{className:"mspace nobreak",children:"\xa0"}),(0,t.jsx)(n.span,{className:"mord mathrm",children:"d"})]}),(0,t.jsx)(n.span,{className:"mord mathnormal",style:{marginRight:"0.0359em"},children:"y"})]})]})]}),(0,t.jsx)(n.span,{className:"vlist-s",children:"\u200b"})]}),(0,t.jsx)(n.span,{className:"vlist-r",children:(0,t.jsx)(n.span,{className:"vlist",style:{height:"0.686em"},children:(0,t.jsx)(n.span,{})})})]})}),(0,t.jsx)(n.span,{className:"mclose nulldelimiter"})]}),(0,t.jsx)(n.span,{className:"mopen",children:"("}),(0,t.jsx)(n.span,{className:"mord mathnormal",children:"t"}),(0,t.jsx)(n.span,{className:"mclose",children:")"}),(0,t.jsx)(n.span,{className:"mspace",style:{marginRight:"0.2778em"}}),(0,t.jsx)(n.span,{className:"mrel",children:"="}),(0,t.jsx)(n.span,{className:"mspace",style:{marginRight:"0.2778em"}})]}),(0,t.jsxs)(n.span,{className:"base",children:[(0,t.jsx)(n.span,{className:"strut",style:{height:"1em",verticalAlign:"-0.25em"}}),(0,t.jsx)(n.span,{className:"mord",children:"\u2212"}),(0,t.jsx)(n.span,{className:"mord mathnormal",style:{marginRight:"0.0359em"},children:"y"}),(0,t.jsx)(n.span,{className:"mopen",children:"("}),(0,t.jsx)(n.span,{className:"mord mathnormal",children:"t"}),(0,t.jsx)(n.span,{className:"mclose",children:")"})]})]})]})}),"\n",(0,t.jsxs)(n.p,{children:["over the interval ",(0,t.jsxs)(n.span,{className:"katex",children:[(0,t.jsx)(n.span,{className:"katex-mathml",children:(0,t.jsx)(n.math,{xmlns:"http://www.w3.org/1998/Math/MathML",children:(0,t.jsxs)(n.semantics,{children:[(0,t.jsxs)(n.mrow,{children:[(0,t.jsx)(n.mo,{stretchy:"false",children:"["}),(0,t.jsx)(n.mn,{children:"0"}),(0,t.jsx)(n.mo,{separator:"true",children:","}),(0,t.jsx)(n.mn,{children:"3"}),(0,t.jsx)(n.mo,{stretchy:"false",children:"]"})]}),(0,t.jsx)(n.annotation,{encoding:"application/x-tex",children:"[0,3]"})]})})}),(0,t.jsx)(n.span,{className:"katex-html","aria-hidden":"true",children:(0,t.jsxs)(n.span,{className:"base",children:[(0,t.jsx)(n.span,{className:"strut",style:{height:"1em",verticalAlign:"-0.25em"}}),(0,t.jsx)(n.span,{className:"mopen",children:"["}),(0,t.jsx)(n.span,{className:"mord",children:"0"}),(0,t.jsx)(n.span,{className:"mpunct",children:","}),(0,t.jsx)(n.span,{className:"mspace",style:{marginRight:"0.1667em"}}),(0,t.jsx)(n.span,{className:"mord",children:"3"}),(0,t.jsx)(n.span,{className:"mclose",children:"]"})]})})]}),"."]}),"\n",(0,t.jsx)(o.A,{lines:[{type:"input",value:"pip install -U diffrax -q"}]}),"\n",(0,t.jsx)(n.pre,{children:(0,t.jsx)(n.code,{className:"language-python",children:"from diffrax import diffeqsolve, Dopri5, ODETerm, SaveAt, PIDController\n\nvector_field = lambda t, y, args: -y\nterm = ODETerm(vector_field)\nsolver = Dopri5()\nsaveat = SaveAt(ts=[0., 1., 2., 3.])\nstepsize_controller = PIDController(rtol=1e-5, atol=1e-5)\n\nsol = diffeqsolve(term, solver, t0=0, t1=3, dt0=0.1, y0=1, saveat=saveat,\n stepsize_controller=stepsize_controller)\n\nprint(sol.ts)\nprint(sol.ys)\n"})}),"\n",(0,t.jsxs)(n.h2,{id:"2-electromagnetic-simulation-with-fdtdx",children:["2. Electromagnetic Simulation with ",(0,t.jsx)(n.a,{href:"https://github.com/ymahlau/fdtdx",children:(0,t.jsx)(n.code,{children:"fdtdx"})})]}),"\n",(0,t.jsxs)(n.p,{children:[(0,t.jsx)(n.a,{href:"https://github.com/ymahlau/fdtdx",children:(0,t.jsx)(n.code,{children:"fdtdx"})})," is a JAX-based solver for simulating electromagnetic wave propagation using the finite-difference time-domain (FDTD) method\u2014a widely adopted numerical approach for solving Maxwell\u2019s equations. The library supports 1
2D, 2D, and 3D simulations with GPU acceleration through JAX's ",(0,t.jsx)(n.code,{children:"jit"})," compilation."]}),"\n",(0,t.jsxs)(n.p,{children:["What makes ",(0,t.jsx)(n.code,{children:"fdtdx"})," particularly powerful is its ability to run ",(0,t.jsx)(n.strong,{children:"differentiable simulations"}),", enabling gradient-based optimization and control. This feature opens the door to integrating electromagnetic simulations directly into learning and optimization pipelines, dramatically accelerating scientific and engineering workflows."]}),"\n",(0,t.jsx)(n.h3,{id:"applications-1",children:"Applications"}),"\n",(0,t.jsxs)(n.p,{children:[(0,t.jsx)(n.code,{children:"fdtdx"})," can be used in a wide range of scientific and engineering domains, including:"]}),"\n",(0,t.jsxs)(n.ul,{children:["\n",(0,t.jsxs)(n.li,{children:[(0,t.jsx)(n.strong,{children:"Antenna design"})," \u2013 Automatically optimize shape, size, and placement to achieve desired radiation patterns."]}),"\n",(0,t.jsxs)(n.li,{children:[(0,t.jsx)(n.strong,{children:"Waveguide and photonic circuit analysis"})," \u2013 Simulate and refine structures for efficient light propagation."]}),"\n",(0,t.jsxs)(n.li,{children:[(0,t.jsx)(n.strong,{children:"Metamaterial and photonic crystal engineering"})," \u2013 Design custom materials with tailored electromagnetic properties."]}),"\n",(0,t.jsxs)(n.li,{children:[(0,t.jsx)(n.strong,{children:"Inverse design problems"})," \u2013 Use gradients to iteratively improve device geometry or material distribution."]}),"\n",(0,t.jsxs)(n.li,{children:[(0,t.jsx)(n.strong,{children:"Differentiable physics-informed ML"})," \u2013 Integrate Maxwell-compliant simulations into neural network training."]}),"\n"]}),"\n",(0,t.jsxs)(n.h2,{id:"3-molecular-simulation-with-jax-md",children:["3. Molecular Simulation with ",(0,t.jsx)(n.a,{href:"https://github.com/jax-md/jax-md",children:(0,t.jsx)(n.code,{children:"jax-md"})})]}),"\n",(0,t.jsx)(n.p,{children:"jax-md is a high-performance, differentiable molecular dynamics engine built with JAX, designed for simulating particles, materials, and custom potentials. Molecular dynamics is a key tool in computational condensed matter physics, used to explore how microscopic interactions give rise to macroscopic behavior."}),"\n",(0,t.jsx)(n.h3,{id:"applications-2",children:"Applications"}),"\n",(0,t.jsxs)(n.p,{children:[(0,t.jsx)(n.code,{children:"jax-md"})," can be applied in various domains, including:"]}),"\n",(0,t.jsxs)(n.ul,{children:["\n",(0,t.jsxs)(n.li,{children:[(0,t.jsx)(n.strong,{children:"Materials science"})," \u2013 Simulate crystal structures, defects, and phase transitions."]}),"\n",(0,t.jsxs)(n.li,{children:[(0,t.jsx)(n.strong,{children:"Biophysics"})," \u2013 Model protein folding, ligand binding, and molecular interactions."]}),"\n",(0,t.jsxs)(n.li,{children:[(0,t.jsx)(n.strong,{children:"Chemistry"})," \u2013 Study reaction dynamics, solvent effects, and catalysis."]}),"\n"]}),"\n",(0,t.jsx)("iframe",{width:"560",height:"315",src:"https://www.youtube.com/embed/Bkm8tGET7-w?si=mshSBt_3T4R_ha9E",title:"YouTube video player",frameborder:"0",allow:"accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture; web-share",referrerpolicy:"strict-origin-when-cross-origin",allowfullscreen:!0}),"\n",(0,t.jsxs)(n.h2,{id:"4-scientific-ml-with-equinox",children:["4. Scientific ML with ",(0,t.jsx)(n.a,{href:"https://github.com/patrick-kidger/equinox",children:(0,t.jsx)(n.code,{children:"Equinox"})})]}),"\n",(0,t.jsx)(n.p,{children:"Equinox is a flexible JAX library designed for building models\u2014including neural networks and physical systems\u2014using a PyTorch-like syntax. It integrates seamlessly with JAX's functional paradigm, offering features like filtered transformation APIs, PyTree manipulation utilities, and support for runtime errors. Equinox models are pure PyTrees, enabling smooth compatibility with JAX transformations such as jit, grad, and vmap. It's a great choice for scientific computing workflows that require combining neural and physics-based models within the same differentiable pipeline"}),"\n",(0,t.jsx)(o.A,{lines:[{type:"input",value:"pip install -U equinox -q"}]}),"\n",(0,t.jsx)(n.pre,{children:(0,t.jsx)(n.code,{className:"language-python",children:"import equinox as eqx\nimport jax\n\nclass Linear(eqx.Module):\n weight: jax.Array\n bias: jax.Array\n\n def __init__(self, in_size, out_size, key):\n wkey, bkey = jax.random.split(key)\n self.weight = jax.random.normal(wkey, (out_size, in_size))\n self.bias = jax.random.normal(bkey, (out_size,))\n\n def __call__(self, x):\n return self.weight @ x + self.bias\n\[email protected]\[email protected]\ndef loss_fn(model, x, y):\n pred_y = j
2ax.vmap(model)(x)\n return jax.numpy.mean((y - pred_y) ** 2)\n\nbatch_size, in_size, out_size = 32, 2, 3\nmodel = Linear(in_size, out_size, key=jax.random.PRNGKey(0))\nx = jax.numpy.zeros((batch_size, in_size))\ny = jax.numpy.zeros((batch_size, out_size))\ngrads = loss_fn(model, x, y)\n"})}),"\n",(0,t.jsxs)(n.h2,{id:"5-model-training-and-inference-for-geospatial-remote-sensing-and-earth-observation-with-jeo",children:["5. Model training and inference for geospatial remote sensing and Earth Observation with ",(0,t.jsx)(n.a,{href:"https://github.com/google-deepmind/jeo",children:(0,t.jsx)(n.code,{children:"Jeo"})})]}),"\n",(0,t.jsx)(n.p,{children:"Jeo is an open-source library developed by Google DeepMind for training machine learning models in geospatial remote sensing and Earth Observation (EO) applications. Built on JAX and Flax, Jeo offers a flexible and scalable framework designed to run seamlessly on CPUs, GPUs, and Google Cloud TPU VMs.\nJeo provides a comprehensive set of tools for building, training, and evaluating models on large-scale geospatial datasets. It includes support for data loading, preprocessing, model definition, training loops, and evaluation metrics\u2014all optimized for high performance and scalability."}),"\n",(0,t.jsx)(n.p,{children:(0,t.jsx)(n.img,{alt:"Jeo",src:s().A+"",width:"960",height:"391"})}),"\n",(0,t.jsx)(n.h2,{id:"applications-3",children:"Applications"}),"\n",(0,t.jsx)(n.p,{children:"Jeo is particularly well-suited for applications in Earth Observation and remote sensing, including:"}),"\n",(0,t.jsxs)(n.ul,{children:["\n",(0,t.jsxs)(n.li,{children:[(0,t.jsx)(n.strong,{children:"Land cover classification"})," \u2013 Classify satellite imagery into different land cover types (e.g., urban, forest, water)."]}),"\n",(0,t.jsxs)(n.li,{children:[(0,t.jsx)(n.strong,{children:"Change detection"})," \u2013 Identify changes in land use or land cover over time using multi-temporal satellite data."]}),"\n",(0,t.jsxs)(n.li,{children:[(0,t.jsx)(n.strong,{children:"Tracking and monitoring"})," \u2013 Monitor environmental changes such as deforestation, urbanization, or natural disasters.","\n",(0,t.jsxs)(n.ul,{children:["\n",(0,t.jsx)(n.li,{children:(0,t.jsx)(n.a,{href:"https://www.climatechange.ai/events/neurips2024",children:"NeurIPS 2024 Workshop on Tackling Climate Change with Machine Learning"})}),"\n"]}),"\n"]}),"\n",(0,t.jsxs)(n.li,{children:[(0,t.jsx)(n.strong,{children:"Object detection"})," \u2013 Detect and classify objects in satellite imagery, such as buildings, roads, or vehicles.","\n",(0,t.jsxs)(n.ul,{children:["\n",(0,t.jsx)(n.li,{children:(0,t.jsx)(n.a,{href:"https://arxiv.org/pdf/2406.18554",children:"Planted: a dataset for planted forest identification from multi-satellite time series"})}),"\n"]}),"\n"]}),"\n"]}),"\n",(0,t.jsx)(n.h2,{id:"looking-ahead",children:"Looking Ahead"}),"\n",(0,t.jsx)(n.p,{children:"JAX is transforming scientific programming by:"}),"\n",(0,t.jsxs)(n.ul,{children:["\n",(0,t.jsx)(n.li,{children:"Unifying simulation, optimization, and differentiation"}),"\n",(0,t.jsx)(n.li,{children:"Supporting GPU/TPU acceleration"}),"\n",(0,t.jsx)(n.li,{children:"Enabling end-to-end differentiable workflows"}),"\n"]}),"\n",(0,t.jsxs)(n.p,{children:[(0,t.jsx)(n.strong,{children:"Coming soon"}),": A beginner-friendly guide on how to get started with JAX\u2014from installation to writing your first ",(0,t.jsx)(n.code,{children:"jit"}),"-compiled function, understanding ",(0,t.jsx)(n.code,{children:"grad"}),", and building your first differentiable simulator."]}),"\n",(0,t.jsxs)(n.blockquote,{children:["\n",(0,t.jsx)(n.p,{children:"\u26a1 The future of research is fast, differentiable, and written in JAX."}),"\n"]}),"\n",(0,t.jsx)(n.h2,{id:"conclusion",children:"Conclusion"}),"\n",(0,t.jsx)(n.p,{children:"AX goes beyond accelerating model training and inference\u2014it redefines what\u2019s possible in the machine learning context by seamlessly integrating scientific computing into ML workflows. With NumPy-like syntax, automatic differentiation, and XLA-compiled execution across CPUs, GPUs, and TPUs, JAX empowers researchers to build high-performance, end-to-e
2nd differentiable systems. Its functional programming paradigm and composable transformations\u2014such as jit, grad, vmap, and pmap\u2014make it easy to define and express complex models, simulators, and optimization routines. Unlike traditional ML libraries focused primarily on neural networks, JAX supports a broader class of models that incorporate physics, dynamics, and numerical solvers\u2014enabling more interpretable, grounded, and scientifically rigorous machine learning."}),"\n",(0,t.jsx)(n.p,{children:"JAX has also played a central role in training large foundation models such as Gemma and Gemini, which were developed using JAX in combination with ML Pathways, Google\u2019s scalable infrastructure for multitask learning. JAX\u2019s ability to leverage modern hardware\u2014especially TPUs\u2014enables faster, more efficient training at massive scale. As detailed in the Gemini research paper, \"the 'single controller' programming model of JAX and Pathways allows a single Python process to orchestrate the entire training run, dramatically simplifying the development workflow.\" This synergy between JAX and Pathways exemplifies how cutting-edge infrastructure and functional design can accelerate the creation of general-purpose AI systems."}),"\n",(0,t.jsx)(n.h3,{id:"references",children:"References"}),"\n",(0,t.jsxs)(n.ul,{children:["\n",(0,t.jsx)(n.li,{children:(0,t.jsx)(n.a,{href:"https://docs.kidger.site/diffrax/",children:"Diffrax Documentation"})}),"\n",(0,t.jsx)(n.li,{children:(0,t.jsx)(n.a,{href:"https://deepmind.google/discover/blog/using-jax-to-accelerate-our-research/",children:"Using JAX to accelerate our research"})}),"\n",(0,t.jsx)(n.li,{children:(0,t.jsx)(n.a,{href:"https://arxiv.org/html/2412.12360v2",children:"A flexible framework for large-scale FDTD simulations: open-source inverse design for 3D nanostructures"})}),"\n",(0,t.jsx)(n.li,{children:(0,t.jsx)(n.a,{href:"https://arxiv.org/abs/1912.04232",children:"JAX, M.D.: A Framework for Differentiable Physics"})}),"\n",(0,t.jsx)(n.li,{children:(0,t.jsx)(n.a,{href:"https://docs.kidger.site/equinox/",children:"Equinox Documentation"})}),"\n",(0,t.jsx)(n.li,{children:(0,t.jsx)(n.a,{href:"https://google-deepmind.github.io/jeo/",children:"Jeo Documentation"})}),"\n",(0,t.jsx)(n.li,{children:(0,t.jsx)(n.a,{href:"https://ai.google.dev/gemma/docs/core/model_card_3#:~:text=to%20operate%20sustainably.-,Software,dramatically%20simplifying%20the%20development%20workflow.%22",children:"Gemma 3 model card"})}),"\n",(0,t.jsx)(n.li,{children:(0,t.jsx)(n.a,{href:"https://www.manning.com/books/deep-learning-with-jax",children:"Deep Learning with JAX"})}),"\n"]})]})}function p(e={}){const{wrapper:n}={...(0,r.R)(),...e.components};return n?(0,t.jsx)(n,{...e,children:(0,t.jsx)(u,{...e})}):u(e)}},7227(e,n,i){i.d(n,{A:()=>o});i(96540);var s=i(34164),a=i(8626);const t="tabItem_Ymn6";var r=i(74848);function l({children:e,className:n,hidden:i}){return(0,r.jsx)("div",{role:"tabpanel",className:(0,s.A)(t,n),hidden:i,children:e})}function o({children:e,className:n,value:i}){const{selectedValue:s,lazy:t}=(0,a.uc)(),o=i===s;return!o&&t?null:(0,r.jsx)(l,{className:n,hidden:!o,children:e})}},45645(e,n,i){i.d(n,{A:()=>p});i(96540);var s=i(34164),a=i(18630),t=i(8626),r=i(24245),l=i(11062);const o="tabList__CuJ",c="tabItem_LNqP";var d=i(74848);function h({className:e}){const{selectedValue:n,selectValue:i,tabValues:a,block:l}=(0,t.uc)(),o=[],{blockElementScrollPositionUntilNextRender:h}=(0,r.a_)(),m=e=>{const s=e.currentTarget,t=o.indexOf(s),r=a[t].value;r!==n&&(h(s),i(r))},u=e=>{let n=null;switch(e.key){case"Enter":m(e);break;case"ArrowRight":{const i=o.indexOf(e.currentTarget)+1;n=o[i]??o[0];break}case"ArrowLeft":{const i=o.indexOf(e.currentTarget)-1;n=o[i]??o[o.length-1];break}}n?.focus()};return(0,d.jsx)("ul",{role:"tablist","aria-orientation":"horizontal",className:(0,s.A)("tabs",{"tabs--block":l},e),children:a.map(({value:e,label:i,attributes:a})=>(0,d.jsx)("li",{role:"tab",tabIndex:n===e?0:-1,"aria-selected":n===e,ref:e=>{o.push(e)},onKeyDown:u,onClick:m,...a,className:(0,s.A)("tabs__item",c,a?.className,{"tabs__item--active":n===e}),children:i??e},e))})}function m({children:e}){return(0,d.jsx)("div",{className:"margin-top--md",children:e})}function u({className:e,children:n}){return(0,d.jsxs)("div",{className:(0,s.A)(a.G.tabs.container,"tabs-container",o),children:[(0,d.jsx)(h,{className:e}),(0,d.jsx)(m,{children:n})]})}function p(e){const n=(0,l.A)(),i=(0,t.OC)(e);return(0,d.jsx)(t.O_,{value:i,children:(0,d.jsx)(u,{className:e.className,children:(0,t.vT)(e.children)})},String(n))}},8626(e,n,i){i.d(n,{OC:()=>p,O_:()=>f,uc:()=>x,vT:()=>d});var s=i(96540),a=i(19413),t=i(36494),r=i(62814),l=i(69900),o=i(45167),c=i(74848);function d(e){return s.Children.toArray(e).filter(e=>"\n"!==e)}function h(e){const{values:n,children:i}=e;return(0,s.useMemo)(()=>{const e=n??function(e){return s.Children.toArray(e).flatMap(e=>{if(!e)return[];if((0,s.isValidElement)(e)&&function(e){const{props:n}=e;return!!n&&"object"==typeof n&&"value"in n}(e))return[e];
2const n="string"==typeof e.type?e.type:e.type.name;throw new Error(`Docusaurus error: Bad <Tabs> child <${n}>: all children of the <Tabs> component should be <TabItem>, and every <TabItem> should have a unique "value" prop.\nIf you do not want to pass on a "value" prop to the direct children of <Tabs>, you can also pass an explicit <Tabs values={...}> prop.`)}).map(({props:{value:e,label:n,attributes:i,default:s}})=>({value:e,label:n,attributes:i,default:s}))}(i);return function(e){const n=(0,o.XI)(e,(e,n)=>e.value===n.value);if(n.length>0)throw new Error(`Docusaurus error: Duplicate values "${n.map(e=>`'${e.value}'`).join(", ")}" found in <Tabs>. Every value needs to be unique.`)}(e),e},[n,i])}function m({value:e,tabValues:n}){return n.some(n=>n.value===e)}function u({queryString:e=!1,groupId:n}){const i=(0,a.W6)(),t=function({queryString:e=!1,groupId:n}){if("string"==typeof e)return e;if(!1===e)return null;if(!0===e&&!n)throw new Error('Docusaurus error: The <Tabs> component groupId prop is required if queryString=true, because this value is used as the search param name. You can also provide an explicit value such as queryString="my-search-param".');return n??null}({queryString:e,groupId:n});return[(0,r.aZ)(t),(0,s.useCallback)(e=>{if(!t)return;const n=new URLSearchParams(i.location.search);n.set(t,e),i.replace({...i.location,search:n.toString()})},[t,i])]}function p(e){const{defaultValue:n,queryString:i=!1,groupId:a}=e,r=h(e),[o,c]=(0,s.useState)(()=>function({defaultValue:e,tabValues:n}){if(0===n.length)throw new Error("Docusaurus error: the <Tabs> component requires at least one <TabItem> children component");if(e){if(!m({value:e,tabValues:n}))throw new Error(`Docusaurus error: The <Tabs> has a defaultValue "${e}" but none of its children has the corresponding value. Available values are: ${n.map(e=>e.value).join(", ")}. If you intend to show no default tab, use defaultValue={null} instead.`);return e}const i=n.find(e=>e.default)??n[0];if(!i)throw new Error("Unexpected error: 0 tabValues");return i.value}({defaultValue:n,tabValues:r})),[d,p]=u({queryString:i,groupId:a}),[g,x]=function({groupId:e}){const n=function(e){return e?`docusaurus.tab.${e}`:null}(e),[i,a]=(0,l.Dv)(n);return[i,(0,s.useCallback)(e=>{n&&a.set(e)},[n,a])]}({groupId:a}),f=(()=>{const e=d??g;return m({value:e,tabValues:r})?e:null})();(0,t.A)(()=>{f&&c(f)},[f]);return{selectedValue:o,selectValue:(0,s.useCallback)(e=>{if(!m({value:e,tabValues:r}))throw new Error(`Can't select invalid tab value=${e}`);c(e),p(e),x(e)},[p,x,r]),tabValues:r,lazy:e.lazy??!1,block:e.block??!1}}const g=(0,s.createContext)(null);function x(){const e=s.useContext(g);if(!e)throw new Error("useTabsContext() must be used within a Tabs component");return e}function f(e){return(0,c.jsx)(g.Provider,{value:e.value,children:e.children})}},93(e,n,i){i.d(n,{A:()=>y});var s=i(96540);const a="audioPlayer_bD77",t="controlButton_18Mu",r="playPauseButton_BOck",l="timeDisplay__nfh",o="progressBar_GLd5",c="volumeControl_jpxX",d="volumeSlider_ecko",h="speedControl_V7oK",m="speedSelect_tSpO",u="loadingContainer_ejfH",p="loadingText_ca0W",g="spinner_NOhs";var x=i(4309),f=i(74848);const j=e=>{if(isNaN(e)||e===1/0)return"00:00";const n=Math.floor(e/60),i=Math.floor(e%60);return`${String(n).padStart(2,"0")}:${String(i).padStart(2,"0")}`},y=({audioSrc:e,title:n="Audio Version"})=>{const[i,y]=(0,s.useState)(!1),[b,v]=(0,s.useState)(0),[w,A]=(0,s.useState)(0),[k,N]=(0,s.useState)(1),[_,D]=(0,s.use
2State)(!1),[T,E]=(0,s.useState)(1),[S,P]=(0,s.useState)(!1),q=(0,s.useRef)(null),C=(0,s.useRef)(null),J=(0,s.useRef)(null);(0,s.useEffect)(()=>{const e=q.current;if(!e)return;const n=()=>{v(e.duration),A(e.currentTime),P(!0)},i=()=>A(e.currentTime),s=()=>{y(!1),A(e.duration),cancelAnimationFrame(J.current)};return y(!1),A(0),v(0),P(!1),J.current&&cancelAnimationFrame(J.current),e.addEventListener("loadedmetadata",n),e.addEventListener("timeupdate",i),e.addEventListener("ended",s),()=>{e.removeEventListener("loadedmetadata",n),e.removeEventListener("timeupdate",i),e.removeEventListener("ended",s),J.current&&cancelAnimationFrame(J.current)}},[e]);const z=(0,s.useCallback)(()=>{if(q.current&&C.current){const e=q.current,n=C.current,i=e.currentTime;A(i),n.value=i,n.style.setProperty("--seek-before-width",i/b*100+"%"),J.current=requestAnimationFrame(z)}},[b]);(0,s.useEffect)(()=>(i?(q.current.play(),J.current=requestAnimationFrame(z)):(q.current.pause(),cancelAnimationFrame(J.current)),()=>cancelAnimationFrame(J.current)),[i,z]);return(0,f.jsxs)("div",{className:a,children:[(0,f.jsx)("audio",{ref:q,src:e,preload:"metadata"}),S?(0,f.jsxs)(f.Fragment,{children:[(0,f.jsx)("button",{onClick:()=>{S&&y(e=>!e)},className:`${t} ${r}`,"aria-label":i?"Pause":"Play",children:i?(0,f.jsx)(x.kwt,{}):(0,f.jsx)(x.gSK,{})}),(0,f.jsx)("span",{className:l,children:j(w)}),(0,f.jsx)("input",{ref:C,type:"range",className:o,value:w,max:b||0,onChange:e=>{if(!S)return;const n=Number(e.target.value);q.current.currentTime=n,A(n),C.current.style.setProperty("--seek-before-width",n/b*100+"%")},"aria-label":"Seek progress",style:{"--seek-before-width":w/b*100+"%"}}),(0,f.jsx)("span",{className:l,children:j(b)}),(0,f.jsxs)("div",{className:c,children:[(0,f.jsx)("button",{onClick:()=>{const e=q.current;if(_){const n=k>0?k:.5;N(n),e.volume=n,D(!1)}else e.volume=0,D(!0),N(0)},className:t,"aria-label":_?"Unmute":"Mute",children:_?(0,f.jsx)(x.FZ2,{}):(0,f.jsx)(x.pPd,{})}),(0,f.jsx)("input",{type:"range",min:"0",max:"1",step:"0.01",value:_?0:k,onChange:e=>{const n=Number(e.target.value);N(n),q.current.volume=n,D(0===n)},className:d,style:{"--volume-before-width":(_?0:100*k)+"%"}})]}),(0,f.jsx)("div",{className:h,children:(0,f.jsxs)("select",{value:T,onChange:e=>{const n=Number(e.target.value);E(n),q.current.playbackRate=n},className:m,"aria-label":"Playback speed",children:[(0,f.jsx)("option",{value:"0.5",children:"0.5x"}),(0,f.jsx)("option",{value:"0.75",children:"0.75x"}),(0,f.jsx)("option",{value:"1",children:"1x"}),(0,f.jsx)("option",{value:"1.25",children:"1.25x"}),(0,f.jsx)("option",{value:"1.5",children:"1.5x"}),(0,f.jsx)("option",{value:"2",children:"2x"})]})})]}):(0,f.jsxs)("div",{className:u,children:[(0,f.jsx)("span",{className:p,children:"Loading audio..."}),(0,f.jsx)("div",{className:g})]})]})}},98979(e,n,i){i.d(n,{A:()=>h});var s=i.cw(function(e,n){e.exports=class{constructor(e="#termynal",n={}){this.container="string"==typeof e?document.querySelector(e):e,this.pfx=`data-${n.prefix||"ty"}`,this.startDelay=n.startDelay||parseFloat(this.container.getAttribute(`${this.pfx}
2-startDelay`))||600,this.typeDelay=n.typeDelay||parseFloat(this.container.getAttribute(`${this.pfx}-typeDelay`))||90,this.lineDelay=n.lineDelay||parseFloat(this.container.getAttribute(`${this.pfx}-lineDelay`))||1500,this.progressLength=n.progressLength||parseFloat(this.container.getAttribute(`${this.pfx}-progressLength`))||40,this.progressChar=n.progressChar||this.container.getAttribute(`${this.pfx}-progressChar`)||"\u2588",this.progressPercent=n.progressPercent||parseFloat(this.container.getAttribute(`${this.pfx}-progressPercent`))||100,this.color=n.color||this.container.getAttribute(`${this.pfx}-color`)||"white",this.cursor=n.cursor||this.container.getAttribute(`${this.pfx}-cursor`)||"\u258b",this.lineData=this.lineDataToElements(n.lineData||[]),n.noInit||this.init(),this.isRunning=!1}init(){this.lines=[...this.container.querySelectorAll(`[${this.pfx}]`)].concat(this.lineData);const e=getComputedStyle(this.container);this.container.style.width="0px"!==e.width?e.width:void 0,this.container.style.minHeight="0px"!==e.height?e.height:void 0,this.container.setAttribute("data-termynal",""),this.container.innerHTML=""}async start(){await this._wait(this.startDelay),this.isRunning=!0;for(let e of this.lines){const n=e.getAttribute(this.pfx),i=e.getAttribute(`${this.pfx}-delay`)||this.lineDelay;"input"===n?(e.setAttribute(`${this.pfx}-cursor`,this.cursor),await this.type(e),await this._wait(i)):"progress"===n?(await this.progress(e),await this._wait(i)):(this.container.appendChild(e),await this._wait(i)),e.removeAttribute(`${this.pfx}-cursor`)}}async type(e){const n=[...e.textContent],i=e.getAttribute(`${this.pfx}-typeDelay`)||this.typeDelay;e.textContent="",this.container.appendChild(e);for(let s of n)await this._wait(i),e.textContent+=s}async progress(e){const n=e.getAttribute(`${this.pfx}-progressLength`)||this.progressLength,i=(e.getAttribute(`${this.pfx}-progressChar`)||this.progressChar).repeat(n),s=e.getAttribute(`${this.pfx}-progressPercent`)||this.progressPercent;e.textContent="",this.container.appendChild(e);for(let a=1;a<i.length+1;
2a++){await this._wait(this.typeDelay);const n=Math.round(a/i.length*100);if(e.textContent=`${i.slice(0,a)} ${n}%`,n>s)break}}_wait(e){return new Promise(n=>setTimeout(n,e))}lineDataToElements(e){return e.map(e=>{let n=document.createElement("div");return n.innerHTML=`<span ${this._attributes(e)} style="color: ${e.color}">${e.value||""}</span>`,n.firstElementChild})}_attributes(e){let n="";for(let i in e)n+=this.pfx,"type"===i?n+=`="${e[i]}" `:"value"!==i&&(n+=`-${i}="${e[i]}" `);return n}}}),a=i(96540),t=i(29813),r=i(45645),l=i(7227),o=i(51491),c=i(92551),d=i(74848);function h(e){const n=(0,a.useRef)(null),{lines:i,...h}=e;return(0,d.jsx)(t.A,{fallback:(0,d.jsx)("div",{children:"Loading..."}),children:()=>{const t=(0,o.A)(),m=s();(0,a.useEffect)(()=>{const i=new m(n.current,{lineData:e.lines}),s=new IntersectionObserver(([e])=>{e.isIntersecting&&(i.isRunning||i.start())});return s.observe(n.current),()=>{s.disconnect()}},[n]);const u=i.filter(e=>"input"===e.type).map(e=>"comment"in e?"# "+e.comment+"\n"+e.value:e.value).join("\n\n");return(0,d.jsxs)(r.A,{groupId:t,children:[(0,d.jsx)(l.A,{value:"bash",label:"Bash",children:(0,d.jsx)(c.A,{language:"bash",showLineNumbers:!1,children:u})}),(0,d.jsx)(l.A,{value:"terminal",label:"terminal",children:(0,d.jsx)("div",{ref:n,...h,style:{marginBottom:20,marginTop:20}})})]})}})}},44198(e,n,i){const s=i.p+"assets/medias/audio-11ff781d9ae0769abd5bf9adad2b40ba.wav";i.d(n,["A",0,s])},51491(e,n,i){i.d(n,{A:()=>r});const s=new Uint8Array(16);const a=[];for(let l=0;l<256;++l)a.push((l+256).toString(16).slice(1));function t(e,n=0){return(a[e[n+0]]+a[e[n+1]]+a[e[n+2]]+a[e[n+3]]+"-"+a[e[n+4]]+a[e[n+5]]+"-"+a[e[n+6]]+a[e[n+7]]+"-"+a[e[n+8]]+a[e[n+9]]+"-"+a[e[n+10]]+a[e[n+11]]+a[e[n+12]]+a[e[n+13]]+a[e[n+14]]+a[e[n+15]]).toLowerCase()}const r=function(e,n,i){return n||e||!crypto.randomUUID?function(e,n,i){e=e||{};const a=e.random??e.rng?.()??crypto.getRandomValues(s);if(a.length<16)throw new Error("Random bytes length must be >= 16");if(a[6]=15&a[6]|64,a[8]=63&a[8]|128,n){if((i=i||0)<0||i+16>n.length)throw new RangeError(`UUID byte range ${i}:${i+15} is out of buffer bounds`);for(let e=0;e<16;++e)n[i+e]=a[e];return n}return t(a)}(e,n,i):crypto.randomUUID()}},92053(e){e.exports=JSON.parse('{"permalink":"/blog/accelerating-science-with-jax-simulations-physics-and-beyond","source":"@site/blog/2025-06-05-accelerating-science-with-jax-simulations-physics-and-beyond/index.mdx","title":"Accelerating Science with JAX - Simulations, Physics, and Beyond","description":"This post explores how JAX is reshaping modern scientific computing by enabling high-performance, differentiable, and hardware-accelerated workflows across a wide range of applications. From modeling complex systems and solving differential equations to training large-scale geospatial and machine learning models, we introduce powerful libraries like Diffrax, fdtdx, jax-md, Equinox, and Jeo. Learn how JAX supports scalable, composable research pipelines that span domains such as engineering, biology, Earth observation, and AI.","date":"2025-06-05T00:00:00.000Z","tags":[{"inline":true,"label":"python","permalink":"/blog/tags/python"},{"inline":true,"label":"data-science","permalink":"/blog/tags/data-science"},{"inline":true,"label":"jax","permalink":"/blog/tags/jax"},{"inline":true,"label":"scientific-computing","permalink":"/blog/tags/scientific-computing"},{"inline":true,"label":"machine-learning","permalink":"/blog/tags/machine-learning"},{"inline":true,"label":"simulations","permalink":"/blog/tags/simulations"},{"inline":true,"label":"physics","permalink":"/blog/tags/physics"},{"inline":true,"label":"electromagnetics","permalink":"/blog/tags/electromagnetics"},{"inline":true,"label":"geosciences","permalink":"/blog/tags/geosciences"},{"inline":true,"label":"earth-observation","permalink":"/blog/tags/earth-observation"}],"readingTime":9.17,"hasTruncateMarker":true,"authors":[{"name":"Henry Ruiz","title":"Blog Author","url":"https://github.com/haruiz","
2imageURL":"https://github.com/haruiz.png","key":"haruiz","page":null}],"frontMatter":{"title":"Accelerating Science with JAX - Simulations, Physics, and Beyond","slug":"accelerating-science-with-jax-simulations-physics-and-beyond","description":"This post explores how JAX is reshaping modern scientific computing by enabling high-performance, differentiable, and hardware-accelerated workflows across a wide range of applications. From modeling complex systems and solving differential equations to training large-scale geospatial and machine learning models, we introduce powerful libraries like Diffrax, fdtdx, jax-md, Equinox, and Jeo. Learn how JAX supports scalable, composable research pipelines that span domains such as engineering, biology, Earth observation, and AI.","authors":["haruiz"],"tags":["python","data-science","jax","scientific-computing","machine-learning","simulations","physics","electromagnetics","geosciences","earth-observation"]},"unlisted":false,"prevItem":{"title":"Building a RAG System with PostgreSQL, pgvector, SQLAlchemy, and Gemma/Gemini Embeddings","permalink":"/blog/building-a-rag-system-with-postgresql-pgvector-sqlalchemy-gemma-gemini-embeddings"},"nextItem":{"title":"Implementing Anthropic\u2019s Agent Design Patterns with Google ADK","permalink":"/blog/implementing-anthropic\'s-agent-design-patterns-with-google-adk"}}')}}]);
Line numbers count LF bytes from the start of the resource, as the search results do. Vendor segments are library code the classifier recognised; they are stored but not indexed. Bytes are shown as Latin1 characters, one per byte.