Deploy the app

assets/engine.worker-Ds1p3O3X.js

@@ -0,0 +1,635 @@
+const e=e=>e.a,t=e=>e.a,n=e=>1/(1+Math.exp(-e));function r(e,t,r,i){switch(e){case`silu`:return t*n(t);case`sigmoid`:return n(t);case`exp`:return Math.exp(t);case`square`:return t*t;case`sqrt`:return Math.sqrt(t);case`neg`:return-t;case`tanh`:return Math.tanh(t);case`logclamp`:return Math.log(Math.max(t,r));case`clamp`:return Math.min(Math.max(t,r),i);case`acos`:return Math.acos(t);case`cos`:return Math.cos(t);case`pow`:return t**+r;case`celu`:return t>0?t:r*Math.expm1(t/r);case`erf`:return o(t);case`switch`:return s(t,r,i)[0]}}function i(e,t,r,i,a){switch(e){case`silu`:{let e=n(t);return e*(1+t*(1-e))}case`sigmoid`:return r*(1-r);case`exp`:return r;case`square`:return 2*t;case`sqrt`:return .5/r;case`neg`:return-1;case`tanh`:return 1-r*r;case`logclamp`:return t>=i?1/t:0;case`clamp`:return+(t>=i&&t<=a);case`acos`:return-1/Math.sqrt(1-t*t);case`cos`:return-Math.sin(t);case`pow`:return i*t**(i-1);case`celu`:return t>0?1:Math.exp(t/i);case`erf`:return 2/Math.sqrt(Math.PI)*Math.exp(-t*t);case`switch`:return s(t,i,a)[1]}}const a=(e,t,n)=>e===`full`?t:e===`scalar`?0:e===`row`?t/n|0:t%n;function o(e){let t=Math.sign(e),n=1/(1+.3275911*Math.abs(e));return t*(1-((((1.061405429*n-1.453152027)*n+1.421413741)*n-.284496736)*n+.254829592)*n*Math.exp(-e*e))}function s(e,t,n){let r=n-t,i=(e-t)/r;if(i<5.960464477539063e-8)return[0,0];if(i>.9999999403953552)return[1,0];let a=Math.sqrt(3)/2,o=Math.exp(a*(1/i-1/(1-i))),s=1/(1+o);return[s,o*a*(1/(i*i)+1/((1-i)*(1-i)))*s*s/r]}const c=(e,t,n)=>e===`add`?t+n:e===`sub`?t-n:e===`mul`?t*n:t/n;function l(e,t,n,r){let i=(t-(n-r))/r;if(e===`cosine`){let e=Math.min(Math.max(i,0),1);return[.5*(1+Math.cos(Math.PI*e)),i>=0&&i<=1?-.5*Math.PI*Math.sin(Math.PI*e):0]}let a=1e-6,o=1-1e-6,s=Math.min(Math.max(i,a),o),c=Math.sin(Math.PI*s),l=Math.cos(Math.PI*s)/c,u=Math.tanh(l);return[.5*(1+u),i>=a&&i<=o?-.5*Math.PI*(1-u*u)/(c*c):0]}var u=class{name=`cpu`;matmulNT=null;zeros(e){return{n:e,a:new Float32Array(Math.max(e,1))}}upload(e){return{n:e.length,a:e.slice()}}uploadI32(e){return{n:e.length,i32:!0,a:e.slice()}}async read(t){return e(t).slice(0,t.n)}readSync(t){return e(t).subarray(0,t.n)}free(e){}async sync(){}fill(t,n,r){e(t).fill(n,0,r)}scale(t,n,r,i,a){let o=e(t),s=e(n);if(a)for(let e=0;e<i;e++)s[e]+=r*o[e];else for(let e=0;e<i;e++)s[e]=r*o[e]}matmul(t,n,r,i,a,o,s,c,l,u){let d=e(t),f=e(n),p=e(r);if(l||p.fill(0,0,i*a),!(!s&&c&&this.matmulNT?.(d,f,p,i,a,o,!0))){if(!s&&c)for(let e=0;e<i;e++){let t=e*o;for(let n=0;n<a;n++){let r=n*o,i=0;for(let e=0;e<o;e++)i+=d[t+e]*f[r+e];p[e*a+n]+=i}}else if(!s&&!c)for(let e=0;e<i;e++){let t=e*a;for(let n=0;n<o;n++){let r=d[e*o+n];if(r===0)continue;let i=n*a;for(let e=0;e<a;e++)p[t+e]+=r*f[i+e]}}else if(s&&!c)for(let e=0;e<o;e++){let t=e*i,n=e*a;for(let e=0;e<i;e++){let r=d[t+e];if(r===0)continue;let i=e*a;for(let e=0;e<a;e++)p[i+e]+=r*f[n+e]}}else for(let e=0;e<i;e++)for(let t=0;t<a;t++){let n=0;for(let r=0;r<o;r++)n+=d[r*i+e]*f[t*o+r];p[e*a+t]+=n}}if(u){let t=e(u);for(let e=0;e<i;e++)for(let n=0;n<a;n++)p[e*a+n]+=t[n]}}unary(t,n,i,a,o,s){let c=e(n),l=e(i);for(let e=0;e<a;e++)l[e]=r(t,c[e],o,s)}unaryGrad(t,n,r,a,o,s,c,l){let u=e(n),d=e(r),f=e(a),p=e(o);for(let e=0;e<s;e++)p[e]+=i(t,u[e],d[e],c,l)*f[e]}binary(t,n,r,i,o,s,l){let u=e(n),d=e(r),f=e(i);for(let e=0;e<o;e++){let n=d[a(s,e,l)];f[e]=c(t,u[e],n)}}binaryGrad(t,n,r,i,o,s,c,l,u){let d=e(n),f=e(r),p=e(i),m=o&&e(o),h=s&&e(s);for(let e=0;e<c;e++){let n=a(l,e,u),r=d[e],i=f[n],o=p[e];m&&(m[e]+=t===`mul`?o*i:t===`div`?o/i:o),h&&(h[n]+=t===`mul`?o*r:t===`div`?-o*r/(i*i):t===`sub`?-o:o)}}gather(n,r,i,a,o,s){let c=e(n),l=e(i),u=t(r);for(let e=0;e<a;e++){let t=u[e],n=e*o;if(t<0){s||l.fill(0,n,n+o);continue}let r=t*o;if(s)for(let e=0;e<o;e++)l[n+e]+=c[r+e];else for(let e=0;e<o;e++)l[n+e]=c[r+e]}}segment(n,r,i,a,o,s,c){let l=e(n),u=e(a),d=t(r),f=t(i);for(let e=0;e<o;e++){let t=e*s;c||u.fill(0,t,t+s);for(let n=d[e];n<d[e+1];n++){let e=f[n]*s;for(let n=0;n<s;n++)u[t+n]+=l[e+n]}}}copyCols(t,n,r,i,a,o,s,c,l){let u=e(t),d=e(i);for(let e=0;e<s;e++){let t=e*n+r,i=e*a+o;if(l)for(let e=0;e<c;e++)d[i+e]+=u[t+e];else for(let e=0;e<c;e++)d[i+e]=u[t+e]}}colSum(t,n,r,i,a){let o=e(t),s=e(n);a||s.fill(0,0,i);for(let e=0;e<r;e++)for(let t=0;t<i;t++)s[t]+=o[e*i+t]}rowSum(t,n,r,i,a){let o=e(t),s=e(n);for(let e=0;e<r;e++){let t=0;for(let n=0;n<i;n++)t+=o[e*i+n];s[e]=a?s[e]+t:t}}norm(t,n,r,i,a,o,s,c,l){let u=e(n),d=e(r),f=i&&e(i),p=e(a),m=e(o);for(let e=0;e<s;e++){let n=e*c,r=0;if(t===`layer`){for(let e=0;e<c;e++)r+=u[n+e];r/=c}let i=0;for(let e=0;e<c;e++){let t=u[n+e]-r;i+=t*t}let a=1/Math.sqrt(i/c+l);m[2*e]=r,m[2*e+1]=a;for(let e=0;e<c;e++)p[n+e]=(u[n+e]-r)*a*d[e]+(f?f[e]:0)}}normGrad(t,n,r,i,a,o,s,c,l,u){let d=e(n),f=e(r),p=e(i),m=e(a),h=o&&e(o),g=s&&e(s),_=c&&e(c);for(let e=0;e<l;e++){let n=e*u,r=p[2*e],i=p[2*e+1],a=0,o=0;for(let e=0;e<u;e++){let t=(d[n+e]-r)*i,s=m[n+e]*f[e];a+=s,o+=s*t,g&&(g[e]+=m[n+e]*t),_&&(_[e]+=m[n+e])}if(h){t===`rms`&&(a=0);for(let e=0;e<u;e++){let t=(d[n+e]-r)*i;h[n+e]+=i*(m[n+e]*f[e]-a/u-t*o/u)}}}}attention(n,r,i,a,o,s,c,l){let u=e(n),d=e(r),f=e(a),p=e(o),m=t(i.off),h=t(i.poff),g=s*c,_=3*g,v=/* @__PURE__ */ new Float64Array(256);for(let e=0;e<i.nSeq;e++){let t=m[e],n=m[e+1]-t,r=v.length>=n?v:new Float64Array(n);for(let i=0;i<s;i++)for(let a=0;a<n;a++){let o=(t+a)*_+i*c,s=-1/0;for(let e=0;e<n;e++){let n=(t+e)*_+g+i*c,a=0;for(let e=0;e<c;e++)a+=u[o+e]*u[n+e];a=a*l+d[t+e],r[e]=a,a>s&&(s=a)}let m=0;for(let e=0;e<n;e++)r[e]=Math.exp(r[e]-s),m+=r[e];let v=h[e]+(i*n+a)*n,y=(t+a)*g+i*c;for(let e=0;e<c;e++)f[y+e]=0;for(let e=0;e<n;e++){let n=r[e]/m;p[v+e]=n;let a=(t+e)*_+2*g+i*c;for(let e=0;e<c;e++)f[y+e]+=n*u[a+e]}}}}attentionGrad(n,r,i,a,o,s,c,l,u){let d=e(n),f=e(r),p=e(i),m=e(o),h=s&&e(s),g=t(a.off),_=t(a.poff),v=c*l,y=3*v;for(let e=0;e<a.nSeq;e++){let t=g[e],n=g[e+1]-t,r=new Float64Array(n);for(let i=0;i<c;i++)for(let a=0;a<n;a++){let o=_[e]+(i*n+a)*n,s=(t+a)*v+i*l,c=0;for(let e=0;e<n;e++){let n=(t+e)*y+2*v+i*l,a=0;for(let e=0;e<l;e++)a+=p[s+e]*d[n+e];r[e]=a,c+=a*f[o+e];let u=f[o+e];for(let e=0;e<l;e++)m[n+e]+=u*p[s+e]}let g=(t+a)*y+i*l;for(let e=0;e<n;e++){let n=f[o+e]*(r[e]-c);h&&(h[t+e]+=n);let a=(t+e)*y+v+i*l;for(let e=0;e<l;e++)m[g+e]+=n*u*d[a+e],m[a+e]+=n*u*d[g+e]}}}}cutoff(t,n,r,i,a,o){let s=e(n),c=e(r),u=e(i);for(let e=0;e<a;e++)u[e]=l(t,s[e],c[e],o)[0]}cutoffGrad(t,n,r,i,a,o,s,c){let u=e(n),d=e(r),f=e(i),p=a&&e(a),m=o&&e(o);for(let e=0;e<s;e++){let n=l(t,u[e],d[e],c)[1]*f[e]/c;p&&(p[e]+=n),m&&(m[e]-=n)}}spline(t,n,r,i,a,o,s,c){let l=e(t),u=e(n),f=e(r),p=e(i);for(let e=0;e<a;e++)for(let t=0;t<o;t++)p[e*o+t]=d(u,f,t,s,c,l[e])[0]}splineGrad(t,n,r,i,a,o,s,c,l){let u=e(t),f=e(n),p=e(r),m=e(i),h=e(a);for(let e=0;e<o;e++){let t=0;for(let n=0;n<s;n++)t+=m[e*s+n]*d(f,p,n,c,l,u[e])[1];h[e]+=t}}sph(t,n,r,i){let a=e(t),o=e(n),s=(i+1)**2;for(let e=0;e<r;e++){let t=f(a[3*e],a[3*e+1],a[3*e+2],i,!1);for(let n=0;n<s;n++)o[e*s+n]=t.y[n]}}sphGrad(t,n,r,i,a){let o=e(t),s=e(n),c=e(r),l=(a+1)**2;for(let e=0;e<i;e++){let t=f(o[3*e],o[3*e+1],o[3*e+2],a,!0);for(let n=0;n<l;n++){let r=s[e*l+n];c[3*e]+=r*t.dx[n],c[3*e+1]+=r*t.dy[n],c[3*e+2]+=r*t.dz[n]}}}power(t,n,r,i,a){let o=e(t),s=e(n),c=(a+1)**2,l=a+1,u=i*(i+1)/2;for(let e=0;e<r;e++){let t=0;for(let n=0;n<i;n++)for(let r=n;r<i;r++,t++){let a=n===r?1:Math.SQRT2,d=(e*i+n)*c,f=(e*i+r)*c;for(let n=0;n<l;n++){let r=0;for(let e=n*n;e<(n+1)*(n+1);e++)r+=o[d+e]*o[f+e];s[(e*u+t)*l+n]=a*Math.PI*Math.sqrt(8/(2*n+1))*r}}}}powerGrad(t,n,r,i,a,o){let s=e(t),c=e(n),l=e(r),u=(o+1)**2,d=o+1,f=a*(a+1)/2;for(let e=0;e<i;e++){let t=0;for(let n=0;n<a;n++)for(let r=n;r<a;r++,t++){let i=n===r?2:Math.SQRT2,o=(e*a+n)*u,p=(e*a+r)*u;for(let a=0;a<d;a++){let u=c[(e*f+t)*d+a]*Math.PI*Math.sqrt(8/(2*a+1));for(let e=a*a;e<(a+1)*(a+1);e++)n===r?l[o+e]+=i*u*s[o+e]:(l[o+e]+=i*u*s[p+e],l[p+e]+=i*u*s[o+e])}}}}adam(t,n,r,i,a,o,s,c,l,u){let d=e(t),f=e(n),p=e(r),m=e(i),h=1-s**u,g=1-c**u;for(let e=0;e<a;e++)p[e]=s*p[e]+(1-s)*f[e],m[e]=c*m[e]+(1-c)*f[e]*f[e],d[e]-=o*(p[e]/h)/(Math.sqrt(m[e]/g)+l)}async thumb(t,n,r,i,a){return p(e(t),n,r,i,a)}};function d(e,t,n,r,i,a){let o=a/i,s=Math.floor(o);if(s<0||s>=r-1)return[0,0];let c=o-s,l=n*r+s,u=c*c,d=u*c,f=2*d-3*u+1,p=d-2*u+c,m=-2*d+3*u,h=d-u;return[f*e[l]+p*i*t[l]+m*e[l+1]+h*i*t[l+1],((6*u-6*c)*e[l]+(3*u-4*c+1)*i*t[l]+(-6*u+6*c)*e[l+1]+(3*u-2*c)*i*t[l+1])/i]}function f(e,t,n,r,i){let a=(r+1)**2,o={y:new Float64Array(a),dx:new Float64Array(a),dy:new Float64Array(a),dz:new Float64Array(a)},s=[1],c=[0],l=[0],u=[0],d=[0],f=[0];for(let n=0;n<r;n++)s.push(e*s[n]-t*c[n]),c.push(e*c[n]+t*s[n]),l.push(s[n]+e*l[n]-t*d[n]),u.push(e*u[n]-c[n]-t*f[n]),d.push(c[n]+e*d[n]+t*l[n]),f.push(e*f[n]+s[n]+t*u[n]);for(let e=0;e<=r;e++){let t=0,a=0,p=1;for(let t=1;t<=2*e-1;t+=2)p*=t;let m=0;for(let h=e;h<=r;h++){if(h>e){let r=h===e+1?(2*e+1)*n*p:((2*h-1)*n*p-(h+e-1)*t)/(h-e),i=h===e+1?(2*e+1)*p:((2*h-1)*(p+n*m)-(h+e-1)*a)/(h-e);t=p,a=m,p=r,m=i}let r=1;for(let t=h-e+1;t<=h+e;t++)r*=t;let g=Math.sqrt((2*h+1)/(4*Math.PI)/r)*(e?Math.SQRT2:1),_=h*h+h+e,v=h*h+h-e;o.y[_]=g*p*s[e],e&&(o.y[v]=g*p*c[e]),i&&(o.dx[_]=g*p*l[e],o.dy[_]=g*p*u[e],o.dz[_]=g*m*s[e],e&&(o.dx[v]=g*p*d[e],o.dy[v]=g*p*f[e],o.dz[v]=g*m*c[e]))}}return o}function p(e,t,n,r,i){let a=new Float32Array(r*i),o=0,s=0,c=1/0,l=-1/0;for(let o=0;o<r;o++){let s=Math.floor(o*t/r),c=Math.max(s+1,Math.floor((o+1)*t/r));for(let t=0;t<i;t++){let r=Math.floor(t*n/i),l=Math.max(r+1,Math.floor((t+1)*n/i)),u=0;for(let t=s;t<c;t++)for(let i=r;i<l;i++)u+=e[t*n+i];a[o*i+t]=u/((c-s)*(l-r))}}let u=t*n;for(let t=0;t<u;t++){let n=e[t];o+=n,s+=n*n,n<c&&(c=n),n>l&&(l=n)}let d=u?o/u:0;return{rows:r,cols:i,data:a,mean:d,std:u?Math.sqrt(Math.max(s/u-d*d,0)):0,min:u?c:0,max:u?l:0,absmax:u?Math.max(Math.abs(c),Math.abs(l)):0}}const m=(e,t=1)=>({[e]:t}),h=m(`Å`),g=m(`eV`),_=m(`e`),v=m(`Ha`),y=e=>Object.fromEntries(Object.entries(e).filter(([,e])=>Math.abs(e)>1e-9)),b=(e,t,n=1)=>{let r={...e};for(let[e,i]of Object.entries(t))r[e]=(r[e]??0)+n*i;return y(r)},x=(e,t)=>y(Object.fromEntries(Object.entries(e).map(([e,n])=>[e,n*t]))),S=e=>!!e&&Object.keys(e).every(e=>e===`rad`),C={"-":`⁻`,".":`·`,0:`⁰`,1:`¹`,2:`²`,3:`³`,4:`⁴`,5:`⁵`,6:`⁶`,7:`⁷`,8:`⁸`,9:`⁹`},w=e=>e===1?``:Number.isInteger(e)?[...String(e)].map(e=>C[e]).join(``):`^${+e.toFixed(3)}`;function T(e){if(!e)return``;let t=[`eV`,`Ha`,`e`,`Å`,`rad`],n=Object.keys(e).sort((e,n)=>t.indexOf(e)-t.indexOf(n)),r=n.filter(t=>e[t]>0).map(t=>t+w(e[t])).join(`·`),i=n.filter(t=>e[t]<0).map(t=>t+w(-e[t])).join(`·`);return i?`${r||`1`}/${i}`:r}const E={keep:(e=>{let t=e.filter(e=>e!==void 0);if(t.length)return t.some(e=>e===null)?null:t.every(e=>T(e)===T(t[0]))?t[0]:null}),mul:(e=>e.every(e=>e)?e.reduce((e,t)=>b(e,t)):null),div:(([e,t])=>e&&t?b(e,t,-1):null),pow:e=>([t])=>t?x(t,e):null,fn:(([e])=>S(e)?{}:null),step:(([e])=>e?{}:null),acos:(([e])=>S(e)?m(`rad`):null),none:(()=>null)},D={square:()=>E.pow(2),sqrt:()=>E.pow(.5),pow:e=>E.pow(e),neg:()=>E.keep,clamp:()=>E.keep,acos:()=>E.acos,switch:()=>E.step};var O=class e{shape;buf;requiresGrad;name;static nextId=0;id=e.nextId++;grad=null;node=null;kind;unit;constructor(e,t,n=!1,r=``){this.shape=e,this.buf=t,this.requiresGrad=n,this.name=r}get size(){return this.shape.reduce((e,t)=>e*t,1)}withUnit(e){return this.unit=e,this}get rows(){return this.shape.length>1?this.shape.slice(0,-1).reduce((e,t)=>e*t,1):this.shape[0]??1}get cols(){return this.shape.length>1?this.shape[this.shape.length-1]:1}},k=class{host;nSrc;kind;idx;inv;constructor(e,t,n,r={}){this.host=t,this.nSrc=n,this.kind=r,this.idx=e.uploadI32(t);let i=new Int32Array(n+1);for(let e of t)e>=0&&i[e+1]++;for(let e=0;e<n;e++)i[e+1]+=i[e];let a=new Int32Array(i[n]),o=i.slice(0,n);t.forEach((e,t)=>{e>=0&&(a[o[e]++]=t)}),this.inv={off:e.uploadI32(i),src:e.uploadI32(a)}}get n(){return this.host.length}},A=class{be;tape=[];scopes=[];owned=[];tracer=null;recording=!0;constructor(e){this.be=e}scope(e,t){this.scopes.push(e);try{return t()}finally{this.scopes.pop()}}get scopeName(){return this.scopes.join(`/`)}tensor(e,t,n=``){let r=e.reduce((e,t)=>e*t,1),i=t??this.be.zeros(r);return this.owned.push(i),new O(e,i,!1,n)}constant(e,t,n=``,r){let i=this.tensor(t,this.be.upload(e),n);return i.kind=r,i}index(e,t,n={}){let r=new k(this.be,e,t,n);return this.owned.push(r.idx,r.inv.off,r.inv.src),r}own(e){return this.owned.push(e),e}record(e,t,n,r,i,a=E.none){let o={id:n.id,op:e,scope:this.scopeName,inputs:t,out:n,aux:i};return n.kind??=t.find(e=>e.kind&&e.rows===n.rows)?.kind,n.unit===void 0&&(n.unit=a(t.map(e=>e.unit))),n.requiresGrad=this.recording&&t.some(e=>e.requiresGrad),n.requiresGrad&&(o.backward=r,n.node=o),this.tape.push(o),this.tracer?.forward?.(o),n}g(e){return e.requiresGrad?(e.grad||=this.own(this.be.zeros(e.size)),e.grad):null}linear(e,t,n,r=`linear`){let[i,a]=[e.rows,e.cols],o=t.shape[0],s=this.tensor([...e.shape.slice(0,-1),o]);return this.be.matmul(e.buf,t.buf,s.buf,i,o,a,!1,!0,!1,n?.buf),this.record(r,n?[e,t,n]:[e,t],s,()=>{let r=s.grad,c=this.g(e),l=this.g(t),u=n&&this.g(n);c&&this.be.matmul(r,t.buf,c,i,a,o,!1,!1,!0),l&&this.be.matmul(r,e.buf,l,o,a,i,!0,!1,!0),u&&this.be.colSum(r,u,i,o,!0)})}unary(e,t,n=0,r=0){let i=this.tensor(t.shape);return this.be.unary(e,t.buf,i.buf,t.size,n,r),this.record(e,[t],i,()=>{this.be.unaryGrad(e,t.buf,i.buf,i.grad,this.g(t),t.size,n,r)},void 0,D[e]?.(n)??E.fn)}silu(e){return this.unary(`silu`,e)}sigmoid(e){return this.unary(`sigmoid`,e)}binary(e,t,n,r=e){let i=n.shape.length===2&&n.shape[0]===1&&n.size===t.cols&&t.rows>1,a=n.size===t.size?`full`:n.size===1?`scalar`:i?`col`:`row`;if(a===`row`&&n.size!==t.rows)throw Error(`${e}: cannot broadcast ${n.shape} onto ${t.shape}`);let o=this.tensor(t.shape);return this.be.binary(e,t.buf,n.buf,o.buf,t.size,a,t.cols),this.record(r,[t,n],o,()=>{this.be.binaryGrad(e,t.buf,n.buf,o.grad,this.g(t),this.g(n),t.size,a,t.cols)},void 0,e===`mul`?E.mul:e===`div`?E.div:E.keep)}add(e,t,n=`add`){return this.binary(`add`,e,t,n)}sub(e,t,n=`sub`){return this.binary(`sub`,e,t,n)}mul(e,t,n=`mul`){return this.binary(`mul`,e,t,n)}div(e,t,n=`div`){return this.binary(`div`,e,t,n)}scale(e,t,n=`scale`,r){let i=this.tensor(e.shape);return this.be.scale(e.buf,i.buf,t,e.size,!1),this.record(n,[e],i,()=>this.be.scale(i.grad,this.g(e),t,e.size,!0),void 0,r?([e])=>e?b(e,r):null:E.keep)}gather(e,t,n=`gather`){let r=e.cols,i=this.tensor(e.shape.length>1?[t.n,r]:[t.n]);return i.kind=t.kind.out,this.be.gather(e.buf,t.idx,i.buf,t.n,r,!1),this.record(n,[e],i,()=>{this.be.segment(i.grad,t.inv.off,t.inv.src,this.g(e),t.nSrc,r,!0)},void 0,E.keep)}segmentSum(e,t,n=`segment_sum`){let r=e.cols,i=this.tensor(e.shape.length>1?[t.nSrc,r]:[t.nSrc]);return i.kind=t.kind.src,this.be.segment(e.buf,t.inv.off,t.inv.src,i.buf,t.nSrc,r,!1),this.record(n,[e],i,()=>{this.be.gather(i.grad,t.idx,this.g(e),t.n,r,!0)},void 0,E.keep)}concatCols(e,t=`concat`){let n=e[0].rows,r=e.reduce((e,t)=>e+t.cols,0),i=this.tensor([n,r]),a=0;for(let t of e)this.be.copyCols(t.buf,t.cols,0,i.buf,r,a,n,t.cols,!1),a+=t.cols;return this.record(t,e,i,()=>{let t=0;for(let a of e){let e=this.g(a);e&&this.be.copyCols(i.grad,r,t,e,a.cols,0,n,a.cols,!0),t+=a.cols}},void 0,E.keep)}sliceCols(e,t,n,r=`slice`){let i=e.rows,a=this.tensor([i,n]);return this.be.copyCols(e.buf,e.cols,t,a.buf,n,0,i,n,!1),this.record(r,[e],a,()=>{this.be.copyCols(a.grad,n,0,this.g(e),e.cols,t,i,n,!0)},void 0,E.keep)}concatRows(e,t,n=`concat_rows`){let r=e.cols,i=this.tensor([e.rows+t.rows,r]);return i.kind=`stack`,this.be.scale(e.buf,i.buf,1,e.size,!1),this.be.copyCols(t.buf,t.size,0,i.buf,i.size,e.size,1,t.size,!1),this.record(n,[e,t],i,()=>{let n=this.g(e),r=this.g(t);n&&this.be.scale(i.grad,n,1,e.size,!0),r&&this.be.copyCols(i.grad,i.size,e.size,r,t.size,0,1,t.size,!0)},void 0,E.keep)}reshape(e,t,n,r=`reshape`){let i=this.tensor(t);return i.kind=n,this.be.scale(e.buf,i.buf,1,e.size,!1),this.record(r,[e],i,()=>this.be.scale(i.grad,this.g(e),1,e.size,!0),void 0,E.keep)}repeatCols(e,t,n=`repeat`){return this.linear(e.shape.length>1?e:this.reshape(e,[e.size,1],e.kind,`column`),this.constant(new Float32Array(t).fill(1),[t,1],`ones`),void 0,n).withUnit(e.unit)}sumRows(e,t=`sum`){let n=this.tensor([e.rows]);return this.be.rowSum(e.buf,n.buf,e.rows,e.cols,!1),this.record(t,[e],n,()=>{let t=this.own(this.be.upload(new Float32Array(e.cols).fill(1)));this.be.matmul(n.grad,t,this.g(e),e.rows,e.cols,1,!1,!1,!0)},void 0,E.keep)}sumAll(e,t=`sum`){let n=this.tensor([1]);return this.be.rowSum(e.buf,n.buf,1,e.size,!1),this.record(t,[e],n,()=>{let t=this.own(this.be.upload(new Float32Array(e.size).fill(1)));this.be.matmul(t,n.grad,this.g(e),e.size,1,1,!1,!1,!0)},void 0,E.keep)}norm(e,t,n,r,i){let a=t.rows,o=t.cols,s=this.tensor(t.shape),c=this.own(this.be.zeros(2*a));return this.be.norm(e,t.buf,n.buf,r?.buf??null,s.buf,c,a,o,i),this.record(e===`layer`?`layer_norm`:`rms_norm`,r?[t,n,r]:[t,n],s,()=>{this.be.normGrad(e,t.buf,n.buf,c,s.grad,this.g(t),this.g(n),r?this.g(r):null,a,o)})}attention(e,t,n,r,i){let a=e.cols/3,o=a/r,s=1/(Math.sqrt(o)*i),c=this.tensor([n.nTok,a]),l=this.own(this.be.zeros(n.nProb));return this.be.attention(e.buf,t.buf,n,c.buf,l,r,o,s),this.record(`attention`,[e,t],c,()=>{this.be.attentionGrad(e.buf,l,c.grad,n,this.g(e),this.g(t),r,o,s)},{probs:{buf:l,shape:[n.nProb]}})}cutoff(e,t,n,r){let i=this.tensor(t.shape);return this.be.cutoff(e,t.buf,n.buf,i.buf,t.size,r),this.record(`cutoff_${e}`,[t,n],i,()=>{this.be.cutoffGrad(e,t.buf,n.buf,i.grad,this.g(t),this.g(n),t.size,r)},void 0,E.step)}spline(e,t,n,r,i=`spline`){let[a,o]=t.shape,s=e.size,c=this.tensor([s,a]);return this.be.spline(e.buf,t.buf,n.buf,c.buf,s,a,o,r),this.record(i,[e,t,n],c,()=>this.be.splineGrad(e.buf,t.buf,n.buf,c.grad,this.g(e),s,a,o,r))}sph(e,t,n=`spherical_harmonics`){let r=e.rows,i=this.tensor([r,(t+1)**2]);return this.be.sph(e.buf,i.buf,r,t),this.record(n,[e],i,()=>this.be.sphGrad(e.buf,i.grad,this.g(e),r,t),void 0,E.fn)}power(e,t,n,r,i=`power_spectrum`){let a=this.tensor([t,n*(n+1)/2*(r+1)]);return this.be.power(e.buf,a.buf,t,n,r),this.record(i,[e],a,()=>this.be.powerGrad(e.buf,a.grad,this.g(e),t,n,r))}transpose(e,t=`transpose`){let[n,r]=[e.rows,e.cols],i=this.tensor([r,n]),a=this.own(this.be.upload(Float32Array.from({length:n*n},(e,t)=>+(t%(n+1)===0))));return this.be.matmul(e.buf,a,i.buf,r,n,n,!0,!1,!1),this.record(t,[e],i,()=>this.be.matmul(a,i.grad,this.g(e),n,r,n,!1,!0,!0),void 0,E.keep)}opaque(e,t,n,r){let i=this.tensor(r,this.be.upload(n)),a=this.recording;this.recording=!1;try{return this.record(e,t,i,()=>{})}finally{this.recording=a}}rowNorm(e,t=1e-15){let n=this.unary(`square`,e),r=this.sumRows(n,`sum_sq`),i=this.binary(`add`,r,this.constant(new Float32Array([t]),[1]));return this.unary(`sqrt`,i)}backward(e,t){if(e.requiresGrad){e.grad=this.own(this.be.upload(t??new Float32Array(e.size).fill(1)));for(let e=this.tape.length-1;e>=0;e--){let t=this.tape[e];t.backward&&t.out.grad&&(this.tracer?.backward?.(t),t.backward())}}}release(){for(let e of this.owned)this.be.free(e);this.owned=[],this.tape=[]}};const j=e=>e.g,M={silu:0,sigmoid:1,exp:2,square:3,sqrt:4,neg:5,tanh:6,logclamp:7,clamp:8,acos:9,cos:10,pow:11,celu:12,erf:13,switch:14},N={add:0,sub:1,mul:2,div:3},P={full:0,scalar:1,row:2,col:3},F=`
+fn gidx(g: vec3u, n: vec3u) -> u32 { return g.x + g.y * n.x * 256u; }`,I={fill:`
+struct P { n: u32, v: f32 }
+@group(0) @binding(0) var<storage, read_write> y: array<f32>;
+@group(0) @binding(1) var<uniform> p: P;
+${F}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let i = gidx(g, nw); if (i < p.n) { y[i] = p.v; }
+}`,scale:`
+struct P { n: u32, a: f32, acc: u32 }
+@group(0) @binding(0) var<storage, read> x: array<f32>;
+@group(0) @binding(1) var<storage, read_write> y: array<f32>;
+@group(0) @binding(2) var<uniform> p: P;
+${F}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let i = gidx(g, nw); if (i >= p.n) { return; }
+  var v = p.a * x[i]; if (p.acc != 0u) { v += y[i]; } y[i] = v;
+}`,matmul:`
+struct P { M: u32, N: u32, K: u32, tA: u32, tB: u32, acc: u32, bias: u32 }
+@group(0) @binding(0) var<storage, read> A: array<f32>;
+@group(0) @binding(1) var<storage, read> B: array<f32>;
+@group(0) @binding(2) var<storage, read_write> C: array<f32>;
+@group(0) @binding(3) var<storage, read> bias: array<f32>;
+@group(0) @binding(4) var<uniform> p: P;
+const T = 16u;
+const R = 4u; // outputs per thread along each axis: 64x64 tile per workgroup
+var<workgroup> As: array<array<f32, 64>, 16>; // [k][m]
+var<workgroup> Bs: array<array<f32, 64>, 16>; // [k][n]
+fn a(m: u32, k: u32) -> f32 {
+  if (m >= p.M || k >= p.K) { return 0.0; }
+  if (p.tA != 0u) { return A[k * p.M + m]; } return A[m * p.K + k];
+}
+fn b(k: u32, n: u32) -> f32 {
+  if (k >= p.K || n >= p.N) { return 0.0; }
+  if (p.tB != 0u) { return B[n * p.K + k]; } return B[k * p.N + n];
+}
+@compute @workgroup_size(16, 16)
+fn main(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_id) l: vec3u) {
+  let m0 = wg.y * 64u; let n0 = wg.x * 64u;
+  var s: array<array<f32, 4>, 4>;
+  let lid = l.y * 16u + l.x;
+  for (var k0 = 0u; k0 < p.K; k0 += T) {
+    // 16 x 64 elements of each tile, 4 per thread
+    for (var r = 0u; r < 4u; r++) {
+      let e = lid + r * 256u; let kk = e / 64u; let mm = e % 64u;
+      As[kk][mm] = a(m0 + mm, k0 + kk);
+      Bs[kk][mm] = b(k0 + kk, n0 + mm);
+    }
+    workgroupBarrier();
+    for (var k = 0u; k < T; k++) {
+      var av: array<f32, 4>; var bv: array<f32, 4>;
+      for (var i = 0u; i < R; i++) { av[i] = As[k][l.y + 16u * i]; bv[i] = Bs[k][l.x + 16u * i]; }
+      for (var i = 0u; i < R; i++) { for (var j = 0u; j < R; j++) { s[i][j] += av[i] * bv[j]; } }
+    }
+    workgroupBarrier();
+  }
+  for (var i = 0u; i < R; i++) {
+    let m = m0 + l.y + 16u * i; if (m >= p.M) { continue; }
+    for (var j = 0u; j < R; j++) {
+      let n = n0 + l.x + 16u * j; if (n >= p.N) { continue; }
+      var v = s[i][j];
+      if (p.bias != 0u) { v += bias[n]; }
+      if (p.acc != 0u) { v += C[m * p.N + n]; }
+      C[m * p.N + n] = v;
+    }
+  }
+}`,unary:`
+struct P { n: u32, op: u32, a: f32, b: f32 }
+@group(0) @binding(0) var<storage, read> x: array<f32>;
+@group(0) @binding(1) var<storage, read_write> y: array<f32>;
+@group(0) @binding(2) var<uniform> p: P;
+${F}
+fn sig(v: f32) -> f32 { return 1.0 / (1.0 + exp(-v)); }
+fn erf_(x: f32) -> f32 {
+  let t = 1.0 / (1.0 + 0.3275911 * abs(x));
+  let y = 1.0 - ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t * exp(-x * x);
+  return sign(x) * y;
+}
+fn sw(x: f32, x0: f32, x1: f32) -> vec2f {
+  let w = x1 - x0; let s = (x - x0) / w; let eps = 5.960464477539063e-8;
+  if (s < eps) { return vec2f(0.0, 0.0); }
+  if (s > 1.0 - eps) { return vec2f(1.0, 0.0); }
+  let c = 0.8660254037844386; let g = exp(c * (1.0 / s - 1.0 / (1.0 - s)));
+  let f = 1.0 / (1.0 + g);
+  return vec2f(f, g * c * (1.0 / (s * s) + 1.0 / ((1.0 - s) * (1.0 - s))) * f * f / w);
+}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let i = gidx(g, nw); if (i >= p.n) { return; }
+  let v = x[i]; var r = 0.0;
+  switch p.op {
+    case 0u: { r = v * sig(v); }
+    case 1u: { r = sig(v); }
+    case 2u: { r = exp(v); }
+    case 3u: { r = v * v; }
+    case 4u: { r = sqrt(v); }
+    case 5u: { r = -v; }
+    case 6u: { r = tanh(clamp(v, -20.0, 20.0)); }
+    case 7u: { r = log(max(v, p.a)); }
+    case 8u: { r = min(max(v, p.a), p.b); }
+    case 9u: { r = acos(v); }
+    case 10u: { r = cos(v); }
+    case 11u: { r = pow(v, p.a); }
+    case 12u: { r = select(p.a * (exp(v / p.a) - 1.0), v, v > 0.0); }
+    case 13u: { r = erf_(v); }
+    default: { r = sw(v, p.a, p.b).x; }
+  }
+  y[i] = r;
+}`,unaryGrad:`
+struct P { n: u32, op: u32, a: f32, b: f32 }
+@group(0) @binding(0) var<storage, read> x: array<f32>;
+@group(0) @binding(1) var<storage, read> y: array<f32>;
+@group(0) @binding(2) var<storage, read> dy: array<f32>;
+@group(0) @binding(3) var<storage, read_write> dx: array<f32>;
+@group(0) @binding(4) var<uniform> p: P;
+${F}
+fn sig(v: f32) -> f32 { return 1.0 / (1.0 + exp(-v)); }
+fn erf_(x: f32) -> f32 {
+  let t = 1.0 / (1.0 + 0.3275911 * abs(x));
+  let y = 1.0 - ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t * exp(-x * x);
+  return sign(x) * y;
+}
+fn sw(x: f32, x0: f32, x1: f32) -> vec2f {
+  let w = x1 - x0; let s = (x - x0) / w; let eps = 5.960464477539063e-8;
+  if (s < eps) { return vec2f(0.0, 0.0); }
+  if (s > 1.0 - eps) { return vec2f(1.0, 0.0); }
+  let c = 0.8660254037844386; let g = exp(c * (1.0 / s - 1.0 / (1.0 - s)));
+  let f = 1.0 / (1.0 + g);
+  return vec2f(f, g * c * (1.0 / (s * s) + 1.0 / ((1.0 - s) * (1.0 - s))) * f * f / w);
+}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let i = gidx(g, nw); if (i >= p.n) { return; }
+  let v = x[i]; let o = y[i]; var d = 0.0;
+  switch p.op {
+    case 0u: { let s = sig(v); d = s * (1.0 + v * (1.0 - s)); }
+    case 1u: { d = o * (1.0 - o); }
+    case 2u: { d = o; }
+    case 3u: { d = 2.0 * v; }
+    case 4u: { d = 0.5 / o; }
+    case 5u: { d = -1.0; }
+    case 6u: { d = 1.0 - o * o; }
+    case 7u: { d = select(0.0, 1.0 / v, v >= p.a); }
+    case 8u: { d = select(0.0, 1.0, v >= p.a && v <= p.b); }
+    case 9u: { d = -1.0 / sqrt(1.0 - v * v); }
+    case 10u: { d = -sin(v); }
+    case 11u: { d = p.a * pow(v, p.a - 1.0); }
+    case 12u: { d = select(exp(v / p.a), 1.0, v > 0.0); }
+    case 13u: { d = 1.1283791670955126 * exp(-v * v); }
+    default: { d = sw(v, p.a, p.b).y; }
+  }
+  dx[i] += d * dy[i];
+}`,binary:`
+struct P { n: u32, op: u32, mode: u32, inner: u32 }
+@group(0) @binding(0) var<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+@group(0) @binding(2) var<storage, read_write> y: array<f32>;
+@group(0) @binding(3) var<uniform> p: P;
+${F}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let i = gidx(g, nw); if (i >= p.n) { return; }
+  var j = i; if (p.mode == 1u) { j = 0u; } else if (p.mode == 2u) { j = i / p.inner; } else if (p.mode == 3u) { j = i % p.inner; }
+  let u = a[i]; let v = b[j]; var r = 0.0;
+  switch p.op { case 0u: { r = u + v; } case 1u: { r = u - v; } case 2u: { r = u * v; } default: { r = u / v; } }
+  y[i] = r;
+}`,binaryGradA:`
+struct P { n: u32, op: u32, mode: u32, inner: u32, wantA: u32, wantB: u32 }
+@group(0) @binding(0) var<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+@group(0) @binding(2) var<storage, read> dy: array<f32>;
+@group(0) @binding(3) var<storage, read_write> da: array<f32>;
+@group(0) @binding(4) var<storage, read_write> db: array<f32>;
+@group(0) @binding(5) var<uniform> p: P;
+${F}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let i = gidx(g, nw); if (i >= p.n) { return; }
+  var j = i; if (p.mode == 1u) { j = 0u; } else if (p.mode == 2u) { j = i / p.inner; } else if (p.mode == 3u) { j = i % p.inner; }
+  let u = a[i]; let v = b[j]; let gr = dy[i];
+  if (p.wantA != 0u) {
+    var d = gr; if (p.op == 2u) { d = gr * v; } else if (p.op == 3u) { d = gr / v; }
+    da[i] += d;
+  }
+  if (p.wantB != 0u && p.mode == 0u) {
+    var d = gr; if (p.op == 1u) { d = -gr; } else if (p.op == 2u) { d = gr * u; } else if (p.op == 3u) { d = -gr * u / (v * v); }
+    db[i] += d;
+  }
+}`,binaryGradB:`
+struct P { nb: u32, op: u32, count: u32, stride: u32, step: u32 }
+@group(0) @binding(0) var<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+@group(0) @binding(2) var<storage, read> dy: array<f32>;
+@group(0) @binding(3) var<storage, read_write> db: array<f32>;
+@group(0) @binding(4) var<uniform> p: P;
+${F}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let j = gidx(g, nw); if (j >= p.nb) { return; }
+  let v = b[j]; var s = 0.0;
+  for (var k = 0u; k < p.count; k++) {
+    let i = j * p.stride + k * p.step; let gr = dy[i];
+    var d = gr; if (p.op == 1u) { d = -gr; } else if (p.op == 2u) { d = gr * a[i]; } else if (p.op == 3u) { d = -gr * a[i] / (v * v); }
+    s += d;
+  }
+  db[j] += s;
+}`,gather:`
+struct P { nOut: u32, d: u32, acc: u32 }
+@group(0) @binding(0) var<storage, read> x: array<f32>;
+@group(0) @binding(1) var<storage, read> idx: array<i32>;
+@group(0) @binding(2) var<storage, read_write> y: array<f32>;
+@group(0) @binding(3) var<uniform> p: P;
+${F}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let e = gidx(g, nw); if (e >= p.nOut * p.d) { return; }
+  let i = e / p.d; let c = e % p.d; let j = idx[i];
+  var v = 0.0; if (j >= 0) { v = x[u32(j) * p.d + c]; }
+  if (p.acc != 0u) { v += y[e]; } y[e] = v;
+}`,segment:`
+struct P { nOut: u32, d: u32, acc: u32 }
+@group(0) @binding(0) var<storage, read> x: array<f32>;
+@group(0) @binding(1) var<storage, read> off: array<i32>;
+@group(0) @binding(2) var<storage, read> src: array<i32>;
+@group(0) @binding(3) var<storage, read_write> y: array<f32>;
+@group(0) @binding(4) var<uniform> p: P;
+${F}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let e = gidx(g, nw); if (e >= p.nOut * p.d) { return; }
+  let t = e / p.d; let c = e % p.d;
+  var v = 0.0;
+  for (var k = off[t]; k < off[t + 1u]; k++) { v += x[u32(src[k]) * p.d + c]; }
+  if (p.acc != 0u) { v += y[e]; } y[e] = v;
+}`,copyCols:`
+struct P { rows: u32, w: u32, xCols: u32, xOff: u32, yCols: u32, yOff: u32, acc: u32 }
+@group(0) @binding(0) var<storage, read> x: array<f32>;
+@group(0) @binding(1) var<storage, read_write> y: array<f32>;
+@group(0) @binding(2) var<uniform> p: P;
+${F}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let e = gidx(g, nw); if (e >= p.rows * p.w) { return; }
+  let r = e / p.w; let c = e % p.w;
+  let yi = r * p.yCols + p.yOff + c;
+  var v = x[r * p.xCols + p.xOff + c]; if (p.acc != 0u) { v += y[yi]; } y[yi] = v;
+}`,reduce:`
+struct P { nOut: u32, len: u32, outStride: u32, inStride: u32, acc: u32 }
+@group(0) @binding(0) var<storage, read> x: array<f32>;
+@group(0) @binding(1) var<storage, read_write> y: array<f32>;
+@group(0) @binding(2) var<uniform> p: P;
+${F}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let i = gidx(g, nw); if (i >= p.nOut) { return; }
+  var s = 0.0;
+  for (var k = 0u; k < p.len; k++) { s += x[i * p.outStride + k * p.inStride]; }
+  if (p.acc != 0u) { s += y[i]; } y[i] = s;
+}`,norm:`
+struct P { rows: u32, d: u32, rms: u32, hasB: u32, eps: f32 }
+@group(0) @binding(0) var<storage, read> x: array<f32>;
+@group(0) @binding(1) var<storage, read> w: array<f32>;
+@group(0) @binding(2) var<storage, read> b: array<f32>;
+@group(0) @binding(3) var<storage, read_write> y: array<f32>;
+@group(0) @binding(4) var<storage, read_write> saved: array<f32>;
+@group(0) @binding(5) var<uniform> p: P;
+${F}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let r = gidx(g, nw); if (r >= p.rows) { return; }
+  let o = r * p.d;
+  var mean = 0.0;
+  if (p.rms == 0u) { for (var c = 0u; c < p.d; c++) { mean += x[o + c]; } mean /= f32(p.d); }
+  var v = 0.0;
+  for (var c = 0u; c < p.d; c++) { let t = x[o + c] - mean; v += t * t; }
+  let rstd = 1.0 / sqrt(v / f32(p.d) + p.eps);
+  saved[2u * r] = mean; saved[2u * r + 1u] = rstd;
+  for (var c = 0u; c < p.d; c++) {
+    var bb = 0.0; if (p.hasB != 0u) { bb = b[c]; }
+    y[o + c] = (x[o + c] - mean) * rstd * w[c] + bb;
+  }
+}`,normGradX:`
+struct P { rows: u32, d: u32, rms: u32 }
+@group(0) @binding(0) var<storage, read> x: array<f32>;
+@group(0) @binding(1) var<storage, read> w: array<f32>;
+@group(0) @binding(2) var<storage, read> saved: array<f32>;
+@group(0) @binding(3) var<storage, read> dy: array<f32>;
+@group(0) @binding(4) var<storage, read_write> dx: array<f32>;
+@group(0) @binding(5) var<uniform> p: P;
+${F}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let r = gidx(g, nw); if (r >= p.rows) { return; }
+  let o = r * p.d; let mean = saved[2u * r]; let rstd = saved[2u * r + 1u];
+  var sg = 0.0; var sgx = 0.0;
+  for (var c = 0u; c < p.d; c++) {
+    let xh = (x[o + c] - mean) * rstd; let gg = dy[o + c] * w[c];
+    sg += gg; sgx += gg * xh;
+  }
+  if (p.rms != 0u) { sg = 0.0; }
+  let d = f32(p.d);
+  for (var c = 0u; c < p.d; c++) {
+    let xh = (x[o + c] - mean) * rstd;
+    dx[o + c] += rstd * (dy[o + c] * w[c] - sg / d - xh * sgx / d);
+  }
+}`,normGradW:`
+struct P { rows: u32, d: u32, wantW: u32, wantB: u32 }
+@group(0) @binding(0) var<storage, read> x: array<f32>;
+@group(0) @binding(1) var<storage, read> saved: array<f32>;
+@group(0) @binding(2) var<storage, read> dy: array<f32>;
+@group(0) @binding(3) var<storage, read_write> dw: array<f32>;
+@group(0) @binding(4) var<storage, read_write> db: array<f32>;
+@group(0) @binding(5) var<uniform> p: P;
+${F}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let c = gidx(g, nw); if (c >= p.d) { return; }
+  var sw = 0.0; var sb = 0.0;
+  for (var r = 0u; r < p.rows; r++) {
+    let i = r * p.d + c;
+    sw += dy[i] * (x[i] - saved[2u * r]) * saved[2u * r + 1u]; sb += dy[i];
+  }
+  if (p.wantW != 0u) { dw[c] += sw; }
+  if (p.wantB != 0u) { db[c] += sb; }
+}`,attention:`
+struct P { nTok: u32, H: u32, dh: u32, scale: f32 }
+@group(0) @binding(0) var<storage, read> qkv: array<f32>;
+@group(0) @binding(1) var<storage, read> bias: array<f32>;
+@group(0) @binding(2) var<storage, read> off: array<i32>;
+@group(0) @binding(3) var<storage, read> poff: array<i32>;
+@group(0) @binding(4) var<storage, read> tok2seq: array<i32>;
+@group(0) @binding(5) var<storage, read_write> outp: array<f32>;
+@group(0) @binding(6) var<storage, read_write> probs: array<f32>;
+@group(0) @binding(7) var<uniform> p: P;
+${F}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let e = gidx(g, nw); if (e >= p.nTok * p.H) { return; }
+  let t = e / p.H; let h = e % p.H;
+  let a = u32(tok2seq[t]); let t0 = u32(off[a]); let n = u32(off[a + 1u]) - t0; let i = t - t0;
+  let D = p.H * p.dh; let D3 = 3u * D;
+  let po = u32(poff[a]) + (h * n + i) * n;
+  let qo = t * D3 + h * p.dh;
+  var mx = -3.4e38;
+  for (var j = 0u; j < n; j++) {
+    let ko = (t0 + j) * D3 + D + h * p.dh;
+    var s = 0.0;
+    for (var c = 0u; c < p.dh; c++) { s += qkv[qo + c] * qkv[ko + c]; }
+    s = s * p.scale + bias[t0 + j];
+    probs[po + j] = s; mx = max(mx, s);
+  }
+  var z = 0.0;
+  for (var j = 0u; j < n; j++) { let v = exp(probs[po + j] - mx); probs[po + j] = v; z += v; }
+  let oo = t * D + h * p.dh;
+  for (var c = 0u; c < p.dh; c++) { outp[oo + c] = 0.0; }
+  for (var j = 0u; j < n; j++) {
+    let pr = probs[po + j] / z; probs[po + j] = pr;
+    let vo = (t0 + j) * D3 + 2u * D + h * p.dh;
+    for (var c = 0u; c < p.dh; c++) { outp[oo + c] += pr * qkv[vo + c]; }
+  }
+}`,attnGradQ:`
+struct P { nTok: u32, H: u32, dh: u32, scale: f32 }
+@group(0) @binding(0) var<storage, read> qkv: array<f32>;
+@group(0) @binding(1) var<storage, read> probs: array<f32>;
+@group(0) @binding(2) var<storage, read> dout: array<f32>;
+@group(0) @binding(3) var<storage, read> off: array<i32>;
+@group(0) @binding(4) var<storage, read> poff: array<i32>;
+@group(0) @binding(5) var<storage, read> tok2seq: array<i32>;
+@group(0) @binding(6) var<storage, read_write> dqkv: array<f32>;
+@group(0) @binding(7) var<storage, read_write> rowdot: array<f32>;
+@group(0) @binding(8) var<uniform> p: P;
+${F}
+fn dP(t: u32, j: u32, h: u32) -> f32 {
+  let D = p.H * p.dh; let oo = t * D + h * p.dh; let vo = j * 3u * D + 2u * D + h * p.dh;
+  var s = 0.0; for (var c = 0u; c < p.dh; c++) { s += dout[oo + c] * qkv[vo + c]; } return s;
+}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let e = gidx(g, nw); if (e >= p.nTok * p.H) { return; }
+  let t = e / p.H; let h = e % p.H;
+  let a = u32(tok2seq[t]); let t0 = u32(off[a]); let n = u32(off[a + 1u]) - t0; let i = t - t0;
+  let D = p.H * p.dh; let D3 = 3u * D;
+  let po = u32(poff[a]) + (h * n + i) * n;
+  var dot = 0.0;
+  for (var j = 0u; j < n; j++) { dot += probs[po + j] * dP(t, t0 + j, h); }
+  rowdot[e] = dot;
+  let qo = t * D3 + h * p.dh;
+  for (var j = 0u; j < n; j++) {
+    let gs = probs[po + j] * (dP(t, t0 + j, h) - dot) * p.scale;
+    let ko = (t0 + j) * D3 + D + h * p.dh;
+    for (var c = 0u; c < p.dh; c++) { dqkv[qo + c] += gs * qkv[ko + c]; }
+  }
+}`,attnGradKV:`
+struct P { nTok: u32, H: u32, dh: u32, scale: f32 }
+@group(0) @binding(0) var<storage, read> qkv: array<f32>;
+@group(0) @binding(1) var<storage, read> probs: array<f32>;
+@group(0) @binding(2) var<storage, read> dout: array<f32>;
+@group(0) @binding(3) var<storage, read> off: array<i32>;
+@group(0) @binding(4) var<storage, read> poff: array<i32>;
+@group(0) @binding(5) var<storage, read> tok2seq: array<i32>;
+@group(0) @binding(6) var<storage, read_write> dqkv: array<f32>;
+@group(0) @binding(7) var<storage, read_write> rowdot: array<f32>;
+@group(0) @binding(8) var<uniform> p: P;
+${F}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let e = gidx(g, nw); if (e >= p.nTok * p.H) { return; }
+  let t = e / p.H; let h = e % p.H;
+  let a = u32(tok2seq[t]); let t0 = u32(off[a]); let n = u32(off[a + 1u]) - t0; let j = t - t0;
+  let D = p.H * p.dh; let D3 = 3u * D;
+  let ko = t * D3 + D + h * p.dh; let vo = t * D3 + 2u * D + h * p.dh;
+  for (var i = 0u; i < n; i++) {
+    let q = t0 + i;
+    let pr = probs[u32(poff[a]) + (h * n + i) * n + j];
+    let oo = q * D + h * p.dh;
+    var dp = 0.0; for (var c = 0u; c < p.dh; c++) { dp += dout[oo + c] * qkv[vo + c]; }
+    let gs = pr * (dp - rowdot[q * p.H + h]);
+    let qo = q * D3 + h * p.dh;
+    for (var c = 0u; c < p.dh; c++) {
+      dqkv[ko + c] += gs * p.scale * qkv[qo + c];
+      dqkv[vo + c] += pr * dout[oo + c];
+    }
+  }
+}`,attnGradBias:`
+struct P { nTok: u32, H: u32, dh: u32, scale: f32 }
+@group(0) @binding(0) var<storage, read> qkv: array<f32>;
+@group(0) @binding(1) var<storage, read> probs: array<f32>;
+@group(0) @binding(2) var<storage, read> dout: array<f32>;
+@group(0) @binding(3) var<storage, read> off: array<i32>;
+@group(0) @binding(4) var<storage, read> poff: array<i32>;
+@group(0) @binding(5) var<storage, read> tok2seq: array<i32>;
+@group(0) @binding(6) var<storage, read_write> dbias: array<f32>;
+@group(0) @binding(7) var<storage, read_write> rowdot: array<f32>;
+@group(0) @binding(8) var<uniform> p: P;
+${F}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let t = gidx(g, nw); if (t >= p.nTok) { return; }
+  let a = u32(tok2seq[t]); let t0 = u32(off[a]); let n = u32(off[a + 1u]) - t0; let j = t - t0;
+  let D = p.H * p.dh; let D3 = 3u * D;
+  var db = 0.0;
+  for (var h = 0u; h < p.H; h++) {
+    let vo = t * D3 + 2u * D + h * p.dh;
+    for (var i = 0u; i < n; i++) {
+      let q = t0 + i;
+      let pr = probs[u32(poff[a]) + (h * n + i) * n + j];
+      let oo = q * D + h * p.dh;
+      var dp = 0.0; for (var c = 0u; c < p.dh; c++) { dp += dout[oo + c] * qkv[vo + c]; }
+      db += pr * (dp - rowdot[q * p.H + h]);
+    }
+  }
+  dbias[t] += db;
+}`,cutoff:`
+struct P { n: u32, bump: u32, w: f32, grad: u32, wantD: u32, wantR: u32 }
+@group(0) @binding(0) var<storage, read> d: array<f32>;
+@group(0) @binding(1) var<storage, read> rc: array<f32>;
+@group(0) @binding(2) var<storage, read> dy: array<f32>;
+@group(0) @binding(3) var<storage, read_write> y: array<f32>;   // value, or dd
+@group(0) @binding(4) var<storage, read_write> drc: array<f32>;
+@group(0) @binding(5) var<uniform> p: P;
+${F}
+const PI = 3.14159265358979;
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let i = gidx(g, nw); if (i >= p.n) { return; }
+  let s = (d[i] - (rc[i] - p.w)) / p.w;
+  var f = 0.0; var dfds = 0.0;
+  if (p.bump == 0u) {
+    let c = clamp(s, 0.0, 1.0);
+    f = 0.5 * (1.0 + cos(PI * c));
+    dfds = select(0.0, -0.5 * PI * sin(PI * c), s >= 0.0 && s <= 1.0);
+  } else {
+    let lo = 1e-6; let hi = 1.0 - 1e-6;
+    let c = clamp(s, lo, hi);
+    let sn = sin(PI * c); let ct = cos(PI * c) / sn; let th = tanh(clamp(ct, -20.0, 20.0));
+    f = 0.5 * (1.0 + th);
+    dfds = select(0.0, -0.5 * PI * (1.0 - th * th) / (sn * sn), s >= lo && s <= hi);
+  }
+  if (p.grad == 0u) { y[i] = f; return; }
+  let gg = dfds * dy[i] / p.w;
+  if (p.wantD != 0u) { y[i] += gg; }
+  if (p.wantR != 0u) { drc[i] -= gg; }
+}`,adam:`
+struct P { n: u32, lr: f32, b1: f32, b2: f32, eps: f32, c1: f32, c2: f32 }
+@group(0) @binding(0) var<storage, read_write> w: array<f32>;
+@group(0) @binding(1) var<storage, read> gr: array<f32>;
+@group(0) @binding(2) var<storage, read_write> m: array<f32>;
+@group(0) @binding(3) var<storage, read_write> v: array<f32>;
+@group(0) @binding(4) var<uniform> p: P;
+${F}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let i = gidx(g, nw); if (i >= p.n) { return; }
+  m[i] = p.b1 * m[i] + (1.0 - p.b1) * gr[i];
+  v[i] = p.b2 * v[i] + (1.0 - p.b2) * gr[i] * gr[i];
+  w[i] -= p.lr * (m[i] / p.c1) / (sqrt(v[i] / p.c2) + p.eps);
+}`,spline:`
+struct P { n: u32, C: u32, K: u32, h: f32, grad: u32 }
+@group(0) @binding(0) var<storage, read> x: array<f32>;
+@group(0) @binding(1) var<storage, read> V: array<f32>;
+@group(0) @binding(2) var<storage, read> D: array<f32>;
+@group(0) @binding(3) var<storage, read> dy: array<f32>;
+@group(0) @binding(4) var<storage, read_write> y: array<f32>; // values, or dx
+@group(0) @binding(5) var<uniform> p: P;
+${F}
+fn herm(c: u32, xv: f32) -> vec2f {
+  let t = xv / p.h; let kf = floor(t);
+  if (kf < 0.0 || kf >= f32(p.K - 1u)) { return vec2f(0.0, 0.0); }
+  let k = u32(kf); let s = t - kf; let o = c * p.K + k;
+  let s2 = s * s; let s3 = s2 * s;
+  let v = (2.0 * s3 - 3.0 * s2 + 1.0) * V[o] + (s3 - 2.0 * s2 + s) * p.h * D[o] + (-2.0 * s3 + 3.0 * s2) * V[o + 1u] + (s3 - s2) * p.h * D[o + 1u];
+  let dv = ((6.0 * s2 - 6.0 * s) * V[o] + (3.0 * s2 - 4.0 * s + 1.0) * p.h * D[o] + (-6.0 * s2 + 6.0 * s) * V[o + 1u] + (3.0 * s2 - 2.0 * s) * p.h * D[o + 1u]) / p.h;
+  return vec2f(v, dv);
+}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  if (p.grad == 0u) {
+    let i = gidx(g, nw); if (i >= p.n * p.C) { return; }
+    y[i] = herm(i % p.C, x[i / p.C]).x;
+  } else {
+    let e = gidx(g, nw); if (e >= p.n) { return; }
+    var s = 0.0;
+    for (var c = 0u; c < p.C; c++) { s += dy[e * p.C + c] * herm(c, x[e]).y; }
+    y[e] += s;
+  }
+}`,sph:`
+struct P { n: u32, L: u32, grad: u32 }
+@group(0) @binding(0) var<storage, read> u: array<f32>;
+@group(0) @binding(1) var<storage, read> dy: array<f32>;
+@group(0) @binding(2) var<storage, read_write> y: array<f32>; // Y, or du
+@group(0) @binding(3) var<uniform> p: P;
+${F}
+const PI = 3.14159265358979;
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let e = gidx(g, nw); if (e >= p.n) { return; }
+  let x = u[3u * e]; let yy = u[3u * e + 1u]; let z = u[3u * e + 2u];
+  let M = (p.L + 1u) * (p.L + 1u);
+  var C: array<f32, 7>; var S: array<f32, 7>; var Cx: array<f32, 7>; var Cy: array<f32, 7>; var Sx: array<f32, 7>; var Sy: array<f32, 7>;
+  C[0] = 1.0; S[0] = 0.0; Cx[0] = 0.0; Cy[0] = 0.0; Sx[0] = 0.0; Sy[0] = 0.0;
+  for (var m = 0u; m < p.L; m++) {
+    C[m + 1u] = x * C[m] - yy * S[m]; S[m + 1u] = x * S[m] + yy * C[m];
+    Cx[m + 1u] = C[m] + x * Cx[m] - yy * Sx[m]; Cy[m + 1u] = x * Cy[m] - S[m] - yy * Sy[m];
+    Sx[m + 1u] = S[m] + x * Sx[m] + yy * Cx[m]; Sy[m + 1u] = x * Sy[m] + C[m] + yy * Cy[m];
+  }
+  var gx = 0.0; var gy = 0.0; var gz = 0.0;
+  for (var m = 0u; m <= p.L; m++) {
+    var q = 1.0; var d = 0.0; var qPrev = 0.0; var dPrev = 0.0;
+    for (var k = 1u; k + 1u <= 2u * m; k += 2u) { q *= f32(k); }
+    for (var l = m; l <= p.L; l++) {
+      if (l > m) {
+        var qn = 0.0; var dn = 0.0;
+        if (l == m + 1u) { qn = f32(2u * m + 1u) * z * q; dn = f32(2u * m + 1u) * q; }
+        else {
+          qn = (f32(2u * l - 1u) * z * q - f32(l + m - 1u) * qPrev) / f32(l - m);
+          dn = (f32(2u * l - 1u) * (q + z * d) - f32(l + m - 1u) * dPrev) / f32(l - m);
+        }
+        qPrev = q; dPrev = d; q = qn; d = dn;
+      }
+      var fact = 1.0;
+      for (var k = l - m + 1u; k <= l + m; k++) { fact *= f32(k); }
+      var N = sqrt(f32(2u * l + 1u) / (4.0 * PI) / fact);
+      if (m > 0u) { N *= 1.41421356237; }
+      let ip = l * l + l + m; let im = l * l + l - m;
+      if (p.grad == 0u) {
+        y[e * M + ip] = N * q * C[m];
+        if (m > 0u) { y[e * M + im] = N * q * S[m]; }
+      } else {
+        let gp = dy[e * M + ip];
+        gx += gp * N * q * Cx[m]; gy += gp * N * q * Cy[m]; gz += gp * N * d * C[m];
+        if (m > 0u) {
+          let gm = dy[e * M + im];
+          gx += gm * N * q * Sx[m]; gy += gm * N * q * Sy[m]; gz += gm * N * d * S[m];
+        }
+      }
+    }
+  }
+  if (p.grad != 0u) { y[3u * e] += gx; y[3u * e + 1u] += gy; y[3u * e + 2u] += gz; }
+}`,power:`
+struct P { B: u32, A: u32, L: u32 }
+@group(0) @binding(0) var<storage, read> x: array<f32>;
+@group(0) @binding(1) var<storage, read_write> pw: array<f32>;
+@group(0) @binding(2) var<uniform> p: P;
+${F}
+const PI = 3.14159265358979;
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let NP = p.A * (p.A + 1u) / 2u; let L1 = p.L + 1u; let M = L1 * L1;
+  let i = gidx(g, nw); if (i >= p.B * NP * L1) { return; }
+  let l = i % L1; let k = (i / L1) % NP; let b = i / (L1 * NP);
+  // pair index k -> (a, a2) with a <= a2
+  var a = 0u; var start = 0u;
+  loop { let rowLen = p.A - a; if (k < start + rowLen) { break; } start += rowLen; a++; }
+  let a2 = a + (k - start);
+  let ro = (b * p.A + a) * M; let r2 = (b * p.A + a2) * M;
+  var s = 0.0;
+  for (var q = l * l; q < (l + 1u) * (l + 1u); q++) { s += x[ro + q] * x[r2 + q]; }
+  var w = 1.0; if (a != a2) { w = 1.41421356237; }
+  pw[i] = w * PI * sqrt(8.0 / f32(2u * l + 1u)) * s;
+}`,powerGrad:`
+struct P { B: u32, A: u32, L: u32 }
+@group(0) @binding(0) var<storage, read> x: array<f32>;
+@group(0) @binding(1) var<storage, read> dp: array<f32>;
+@group(0) @binding(2) var<storage, read_write> dx: array<f32>;
+@group(0) @binding(3) var<uniform> p: P;
+${F}
+const PI = 3.14159265358979;
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let L1 = p.L + 1u; let M = L1 * L1; let NP = p.A * (p.A + 1u) / 2u;
+  let i = gidx(g, nw); if (i >= p.B * p.A * M) { return; }
+  let q = i % M; let row = i / M; let a = row % p.A; let b = row / p.A;
+  var l = 0u; loop { if ((l + 1u) * (l + 1u) > q) { break; } l++; }
+  let f = PI * sqrt(8.0 / f32(2u * l + 1u));
+  var s = 0.0;
+  for (var a2 = 0u; a2 < p.A; a2++) {
+    let lo = min(a, a2); let hi = max(a, a2);
+    let k = lo * p.A - (lo * (lo - 1u)) / 2u + (hi - lo);
+    var w = 1.41421356237; if (a == a2) { w = 2.0; }
+    s += w * f * dp[(b * NP + k) * L1 + l] * x[(b * p.A + a2) * M + q];
+  }
+  dx[i] += s;
+}`,thumb:`
+struct P { R: u32, C: u32, rows: u32, cols: u32, o: u32 }
+@group(0) @binding(0) var<storage, read> x: array<f32>;
+@group(0) @binding(1) var<storage, read_write> y: array<f32>;
+@group(0) @binding(2) var<uniform> p: P;
+${F}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3u, @builtin(num_workgroups) nw: vec3u) {
+  let e = gidx(g, nw); if (e >= p.rows * p.cols) { return; }
+  let r = e / p.cols; let c = e % p.cols;
+  let r0 = (r * p.R) / p.rows; let r1 = max(r0 + 1u, ((r + 1u) * p.R) / p.rows);
+  let c0 = (c * p.C) / p.cols; let c1 = max(c0 + 1u, ((c + 1u) * p.C) / p.cols);
+  var s = 0.0;
+  for (var i = r0; i < r1; i++) { for (var j = c0; j < c1; j++) { s += x[i * p.C + j]; } }
+  y[p.o + e] = s / f32((r1 - r0) * (c1 - c0));
+}`,stats:`
+struct P { n: u32, o: u32 }
+@group(0) @binding(0) var<storage, read> x: array<f32>;
+@group(0) @binding(1) var<storage, read_write> y: array<f32>;
+@group(0) @binding(2) var<uniform> p: P;
+var<workgroup> s1: array<f32, 256>; var<workgroup> s2: array<f32, 256>;
+var<workgroup> mn: array<f32, 256>; var<workgroup> mx: array<f32, 256>;
+@compute @workgroup_size(256) fn main(@builtin(local_invocation_id) l: vec3u) {
+  var a = 0.0; var b = 0.0; var lo = 3.4e38; var hi = -3.4e38;
+  for (var i = l.x; i < p.n; i += 256u) { let v = x[i]; a += v; b += v * v; lo = min(lo, v); hi = max(hi, v); }
+  s1[l.x] = a; s2[l.x] = b; mn[l.x] = lo; mx[l.x] = hi;
+  workgroupBarrier();
+  for (var k = 128u; k > 0u; k >>= 1u) {
+    if (l.x < k) {
+      s1[l.x] += s1[l.x + k]; s2[l.x] += s2[l.x + k];
+      mn[l.x] = min(mn[l.x], mn[l.x + k]); mx[l.x] = max(mx[l.x], mx[l.x + k]);
+    }
+    workgroupBarrier();
+  }
+  if (l.x == 0u) { y[p.o] = s1[0]; y[p.o + 1u] = s2[0]; y[p.o + 2u] = mn[0]; y[p.o + 3u] = mx[0]; }
+}`},L=e=>({f:e});var R=class e{device;adapterName;name=`webgpu`;pipes=/* @__PURE__ */ new Map;pool=/* @__PURE__ */ new Map;enc=null;uni;uniData=/* @__PURE__ */ new ArrayBuffer(1<<20);uniUsed=0;uniAlign;dummy;dispatches=0;owners=[];static async create(t){let n=t??globalThis.navigator?.gpu;if(!n)throw Error(`WebGPU is not available in this browser`);let r=await n.requestAdapter({powerPreference:`high-performance`});if(!r)throw Error(`no WebGPU adapter`);let i=r.limits,a=await r.requestDevice({requiredLimits:{maxStorageBuffersPerShaderStage:Math.min(10,i.maxStorageBuffersPerShaderStage),maxStorageBufferBindingSize:i.maxStorageBufferBindingSize,maxBufferSize:i.maxBufferSize}}),o=new e(a,r.info?.description||r.info?.vendor||`GPU`);return o.owners=[n,r],o}constructor(e,t=`GPU`){this.device=e,this.adapterName=t,this.uniAlign=e.limits.minUniformBufferOffsetAlignment,this.uni=e.createBuffer({size:this.uniData.byteLength,usage:GPUBufferUsage.UNIFORM|GPUBufferUsage.COPY_DST}),this.dummy=e.createBuffer({size:16,usage:GPUBufferUsage.STORAGE})}alloc(e){let t=1<<Math.max(4,Math.ceil(Math.log2(Math.max(e,1)*4))),n=this.pool.get(t);return n?.length?n.pop():this.device.createBuffer({size:t,usage:GPUBufferUsage.STORAGE|GPUBufferUsage.COPY_SRC|GPUBufferUsage.COPY_DST})}wrap(e,t,n=!1){return{n:t,i32:n,g:e,cap:e.size}}zeros(e){let t=this.alloc(e);return this.encoder().clearBuffer(t),this.wrap(t,e)}uploadRaw(e,t){let n=e.length,r=this.device.createBuffer({size:1<<Math.max(4,Math.ceil(Math.log2(Math.max(n,1)*4))),usage:GPUBufferUsage.STORAGE|GPUBufferUsage.COPY_SRC|GPUBufferUsage.COPY_DST,mappedAtCreation:!0});return(t?new Int32Array(r.getMappedRange()):new Float32Array(r.getMappedRange())).set(e),r.unmap(),this.wrap(r,n,t)}upload(e){return this.uploadRaw(e,!1)}uploadI32(e){return this.uploadRaw(e,!0)}free(e){let t=j(e),n=this.pool.get(t.size)??[];n.push(t),this.pool.set(t.size,n)}encoder(){return this.enc??=this.device.createCommandEncoder()}flush(){this.enc&&(this.uniUsed&&this.device.queue.writeBuffer(this.uni,0,this.uniData,0,this.uniUsed),this.device.queue.submit([this.enc.finish()]),this.enc=null,this.uniUsed=0)}async sync(){this.flush(),await this.device.queue.onSubmittedWorkDone()}async read(e){return(await this.readMany([e]))[0]}async readMany(e){let t=e.reduce((e,t)=>e+Math.max(t.n,1)*4,0),n=this.device.createBuffer({size:Math.max(t,4),usage:GPUBufferUsage.MAP_READ|GPUBufferUsage.COPY_DST}),r=this.encoder(),i=0;for(let t of e)t.n&&r.copyBufferToBuffer(j(t),0,n,i,t.n*4),i+=Math.max(t.n,1)*4;this.flush(),await n.mapAsync(GPUMapMode.READ);let a=n.getMappedRange();i=0;let o=e.map(e=>{let t=e.i32?Float32Array.from(new Int32Array(a,i,e.n)):new Float32Array(a.slice(i,i+e.n*4));return i+=Math.max(e.n,1)*4,t});return n.unmap(),n.destroy(),o}pipe(e){let t=this.pipes.get(e);if(!t){let n=this.device.createShaderModule({code:I[e],label:e});t=this.device.createComputePipeline({layout:`auto`,compute:{module:n,entryPoint:`main`},label:e}),this.pipes.set(e,t)}return t}run(e,t,n,r){this.uniUsed+this.uniAlign>this.uniData.byteLength&&this.flush();let i=this.uniUsed,a=new DataView(this.uniData,i,64);n.forEach((e,t)=>typeof e==`number`?a.setUint32(4*t,e,!0):a.setFloat32(4*t,e.f,!0)),this.uniUsed+=Math.max(this.uniAlign,64);let o=this.pipe(e),s=t.map((e,t)=>({binding:t,resource:{buffer:e===null?this.dummy:`g`in e?j(e):e}}));s.push({binding:t.length,resource:{buffer:this.uni,offset:i,size:Math.ceil(n.length*4/16)*16||16}});let c=this.device.createBindGroup({layout:o.getBindGroupLayout(0),entries:s}),l=this.encoder().beginComputePass();l.setPipeline(o),l.setBindGroup(0,c),l.dispatchWorkgroups(r[0],r[1]??1,r[2]??1),l.end(),this.dispatches++}grid(e){let t=Math.max(1,Math.ceil(e/256));return t<=65535?[t,1]:[65535,Math.ceil(t/65535)]}fill(e,t,n){this.run(`fill`,[e],[n,L(t)],this.grid(n))}scale(e,t,n,r,i){this.run(`scale`,[e,t],[r,L(n),+i],this.grid(r))}matmul(e,t,n,r,i,a,o,s,c,l){if(r&&i){if(!a)throw Error(`matmul with K = 0`);this.run(`matmul`,[e,t,n,l??null],[r,i,a,+o,+s,+c,+!!l],[Math.ceil(i/64),Math.ceil(r/64)])}}unary(e,t,n,r,i,a){r&&this.run(`unary`,[t,n],[r,M[e],L(i),L(a)],this.grid(r))}unaryGrad(e,t,n,r,i,a,o,s){a&&this.run(`unaryGrad`,[t,n,r,i],[a,M[e],L(o),L(s)],this.grid(a))}binary(e,t,n,r,i,a,o){i&&this.run(`binary`,[t,n,r],[i,N[e],P[a],o],this.grid(i))}binaryGrad(e,t,n,r,i,a,o,s,c){if(o&&((i||a&&s===`full`)&&this.run(`binaryGradA`,[t,n,r,i,s===`full`?a:null],[o,N[e],P[s],c,+!!i,+(!!a&&s===`full`)],this.grid(o)),a&&s!==`full`)){let[i,l,u,d]=s===`scalar`?[1,o,0,1]:s===`row`?[o/c,c,c,1]:[c,o/c,1,c];this.run(`binaryGradB`,[t,n,r,a],[i,N[e],l,u,d],this.grid(i))}}gather(e,t,n,r,i,a){r*i&&this.run(`gather`,[e,t,n],[r,i,+a],this.grid(r*i))}segment(e,t,n,r,i,a,o){i*a&&this.run(`segment`,[e,t,n,r],[i,a,+o],this.grid(i*a))}copyCols(e,t,n,r,i,a,o,s,c){o*s&&this.run(`copyCols`,[e,r],[o,s,t,n,i,a,+c],this.grid(o*s))}colSum(e,t,n,r,i){r&&this.run(`reduce`,[e,t],[r,n,1,r,+i],this.grid(r))}rowSum(e,t,n,r,i){n&&this.run(`reduce`,[e,t],[n,r,r,1,+i],this.grid(n))}norm(e,t,n,r,i,a,o,s,c){o&&this.run(`norm`,[t,n,r,i,a],[o,s,+(e===`rms`),+!!r,L(c)],this.grid(o))}normGrad(e,t,n,r,i,a,o,s,c,l){c&&(a&&this.run(`normGradX`,[t,n,r,i,a],[c,l,+(e===`rms`)],this.grid(c)),(o||s)&&this.run(`normGradW`,[t,r,i,o,s],[c,l,+!!o,+!!s],this.grid(l)))}attention(e,t,n,r,i,a,o,s){n.nTok&&this.run(`attention`,[e,t,n.off,n.poff,n.tok2seq,r,i],[n.nTok,a,o,L(s)],this.grid(n.nTok*a))}attentionGrad(e,t,n,r,i,a,o,s,c){if(!r.nTok)return;let l=this.zeros(r.nTok*o),u=[e,t,n,r.off,r.poff,r.tok2seq],d=[r.nTok,o,s,L(c)];this.run(`attnGradQ`,[...u,i,l],d,this.grid(r.nTok*o)),this.run(`attnGradKV`,[...u,i,l],d,this.grid(r.nTok*o)),a&&this.run(`attnGradBias`,[...u,a,l],d,this.grid(r.nTok)),this.free(l)}cutoff(e,t,n,r,i,a){i&&this.run(`cutoff`,[t,n,t,r,null],[i,+(e===`bump`),L(a),0,0,0],this.grid(i))}cutoffGrad(e,t,n,r,i,a,o,s){o&&(i||a)&&this.run(`cutoff`,[t,n,r,i,a],[o,+(e===`bump`),L(s),1,+!!i,+!!a],this.grid(o))}spline(e,t,n,r,i,a,o,s){i*a&&this.run(`spline`,[e,t,n,e,r],[i,a,o,L(s),0],this.grid(i*a))}splineGrad(e,t,n,r,i,a,o,s,c){a&&this.run(`spline`,[e,t,n,r,i],[a,o,s,L(c),1],this.grid(a))}sph(e,t,n,r){if(r>6)throw Error(`sph: lmax > 6 is not supported on WebGPU`);n&&this.run(`sph`,[e,e,t],[n,r,0],this.grid(n))}sphGrad(e,t,n,r,i){r&&this.run(`sph`,[e,t,n],[r,i,1],this.grid(r))}power(e,t,n,r,i){let a=n*(r*(r+1)/2)*(i+1);a&&this.run(`power`,[e,t],[n,r,i],this.grid(a))}powerGrad(e,t,n,r,i,a){let o=r*i*(a+1)**2;o&&this.run(`powerGrad`,[e,t,n],[r,i,a],this.grid(o))}adam(e,t,n,r,i,a,o,s,c,l){this.run(`adam`,[e,t,n,r],[i,L(a),L(o),L(s),L(c),L(1-o**l),L(1-s**l)],this.grid(i))}async thumb(e,t,n,r,i){return(await this.thumbs([{x:e,R:t,C:n,rows:r,cols:i}]))[0]}async thumbs(e){let t=0,n=e.map(e=>{let n=t;return t+=e.rows*e.cols+4,n}),r=this.zeros(t);e.forEach((e,t)=>{e.R*e.C&&(this.run(`thumb`,[e.x,r],[e.R,e.C,e.rows,e.cols,n[t]],this.grid(e.rows*e.cols)),this.run(`stats`,[e.x,r],[e.R*e.C,n[t]+e.rows*e.cols],[1]))});let i=await this.read(r);return this.free(r),e.map((e,t)=>{let r=e.R*e.C,a=n[t]+e.rows*e.cols,[o,s,c,l]=i.subarray(a,a+4),u=r?o/r:0;return{rows:e.rows,cols:e.cols,data:i.slice(n[t],a),mean:u,std:r?Math.sqrt(Math.max(s/r-u*u,0)):0,min:r?c:0,max:r?l:0,absmax:r?Math.max(Math.abs(c),Math.abs(l)):0}})}};const z=e=>e[0][0]*(e[1][1]*e[2][2]-e[1][2]*e[2][1])-e[0][1]*(e[1][0]*e[2][2]-e[1][2]*e[2][0])+e[0][2]*(e[1][0]*e[2][1]-e[1][1]*e[2][0]);function B(e){let t=z(e),n=(t,n)=>{let r=[0,1,2].filter(e=>e!==t),i=[0,1,2].filter(e=>e!==n);return(e[r[0]][i[0]]*e[r[1]][i[1]]-e[r[0]][i[1]]*e[r[1]][i[0]])*((t+n)%2?-1:1)};return[0,1,2].map(e=>[0,1,2].map(r=>n(r,e)/t))}const V=(e,t)=>[e[1]*t[2]-e[2]*t[1],e[2]*t[0]-e[0]*t[2],e[0]*t[1]-e[1]*t[0]],ee=e=>Math.hypot(e[0],e[1],e[2]),H=e=>!!e.cell&&!!e.pbc?.some(Boolean)&&Math.abs(z(e.cell))>1e-10,te=e=>H(e)?Math.abs(z(e.cell)):1/0;function U(e,t){let n=e.numbers.length,r=H(e),i=r?e.cell:[[0,0,0],[0,0,0],[0,0,0]],a=r?e.pbc:[!1,!1,!1],o=e.positions.map(()=>[0,0,0]),s=e.positions.map(e=>e.slice()),c=[0,0,0];if(r){let n=B(i);e.positions.forEach((e,t)=>{for(let r=0;r<3;r++){if(!a[r])continue;let i=e[0]*n[0][r]+e[1]*n[1][r]+e[2]*n[2][r];o[t][r]=Math.floor(i)}for(let e=0;e<3;e++)s[t][e]-=o[t][0]*i[0][e]+o[t][1]*i[1][e]+o[t][2]*i[2][e]});let r=Math.abs(z(i));c=[0,1,2].map(e=>a[e]?Math.ceil(t/(r/ee(V(i[(e+1)%3],i[(e+2)%3])))):0)}let l=[],u=[],d=[],f=[],p=[],m=t*t;for(let e=0;e<n;e++)for(let t=-c[0];t<=c[0];t++)for(let r=-c[1];r<=c[1];r++)for(let a=-c[2];a<=c[2];a++){let c=[0,1,2].map(e=>t*i[0][e]+r*i[1][e]+a*i[2][e]);for(let h=0;h<n;h++){if(e===h&&t===0&&r===0&&a===0)continue;let n=s[h][0]+c[0]-s[e][0],g=s[h][1]+c[1]-s[e][1],_=s[h][2]+c[2]-s[e][2],v=n*n+g*g+_*_;if(v>=m)continue;let y=[t-o[h][0]+o[e][0],r-o[h][1]+o[e][1],a-o[h][2]+o[e][2]];l.push(e),u.push(h),d.push(...y),f.push(...[0,1,2].map(e=>y[0]*i[0][e]+y[1]*i[1][e]+y[2]*i[2][e])),p.push(Math.sqrt(v))}}return{center:Int32Array.from(l),neighbor:Int32Array.from(u),shift:Int32Array.from(d),shiftVec:Float64Array.from(f),dist:Float64Array.from(p)}}function W(e){let t=new DataView(e),n=Number(t.getBigUint64(0,!0)),r=JSON.parse(new TextDecoder().decode(new Uint8Array(e,8,n))),i=8+n,a=/* @__PURE__ */ new Map;for(let[t,n]of Object.entries(r)){if(t===`__metadata__`)continue;if(n.dtype!==`F32`)throw Error(`safetensors: ${t} has dtype ${n.dtype}, expected F32`);let[r,o]=n.data_offsets,s=new Float32Array(e.slice(i+r,i+o));a.set(t,{shape:n.shape,data:s})}return a}function G(e){let t=e.hypers;if(e.architecture!==`pet`)throw Error(`not a PET model: ${e.architecture}`);if(t.zbl)throw Error(`ZBL models are not supported yet`);if(t.long_range_enabled)throw Error(`long-range models are not supported yet`);if(t.num_neighbors_adaptive!=null&&t.adaptive_cutoff_method!==`solver`)throw Error(`adaptive_cutoff_method=${t.adaptive_cutoff_method} is not supported yet (only "solver")`)}const K=(e,t=0)=>Int32Array.from({length:e},(e,n)=>t+n),q=e=>({label:`atom`,owner:K(e),atom:K(e),all:!0}),J=e=>({label:`edge`,owner:e,edge:K(e.length)});function ne(e){let t=Array.from({length:e},()=>Array(e).fill(0)),n=0;for(let r=0;r<e;r++)for(let i=r;i<e;i++)t[r][i]=t[i][r]=n++;return t}var re=class{be;meta;kind=`ani`;params=/* @__PURE__ */ new Map;get elements(){return this.meta.elements}constructor(e,t,n){this.be=e,this.meta=t;for(let[t,r]of n)this.params.set(t,new O(r.shape,e.upload(r.data),!1,t))}forward(e,t,n={}){let r=this.meta,i=r.elements.length,a=i*(i+1)/2,o=r.radial,s=r.angular,c=o.shifts.length,l=s.shifts.length,u=s.sections.length,d=t.numbers.length,f=Int32Array.from(t.numbers,e=>r.elements.indexOf(e)),p=U(t,o.cutoff),m=p.center.length,_=(t,n,r,i)=>e.index(Int32Array.from(t),n,{out:r,src:i}),y=(t,n,r,i)=>e.constant(Float32Array.from(t),n,r,i),S=e.constant(Float32Array.from(t.positions.flat()),[d,3],`positions`,`atom`).withUnit(h);S.requiresGrad=!!n.forces;let{v:C,d:w}=e.scope(`geometry`,()=>{let t=e.add(e.sub(e.gather(S,_(p.neighbor,d,`pair`,`atom`),`r_j`),e.gather(S,_(p.center,d,`pair`,`atom`),`r_i`)),y(Float32Array.from(p.shiftVec),[m,3],`shift`,`pair`));return{v:t,d:e.rowNorm(t)}}),T=e.scope(`radial`,()=>{let t=e.cutoff(`cosine`,w,y(Array(m).fill(o.cutoff),[m],`cutoff`,`pair`),o.cutoff),n=e.sub(e.repeatCols(w,c,`distances`),y(o.shifts,[1,c],`shifts`)),r=e.mul(e.scale(e.unary(`exp`,e.scale(e.unary(`square`,n),-o.eta,`scale`,x(h,-2))),.25),t),a=_(Array.from(p.center,(e,t)=>e*i+f[p.neighbor[t]]),d*i,`pair`,`atom_species`);return e.reshape(e.segmentSum(r,a,`sum_per_element`),[d,i*c],`atom`,`radial_aev`)}),E=[],D=Array.from({length:d},()=>[]);p.dist.forEach((e,t)=>{e<s.cutoff&&D[p.center[t]].push(t)}),D.forEach((e,t)=>{for(let n=0;n<e.length;n++)for(let r=n+1;r<e.length;r++)E.push([t,e[n],e[r]])});let O=E.length,k=ne(i),A=e.scope(`angular`,()=>{if(!O)return e.constant(new Float32Array(d*a*l*u),[d,a*l*u],`angular_aev`,`atom`);let t=_(E.map(e=>e[1]),m,`triplet`,`pair`),n=_(E.map(e=>e[2]),m,`triplet`,`pair`),r=e.gather(C,t,`r_ij`),i=e.gather(C,n,`r_ik`),o=e.gather(w,t,`|r_ij|`),c=e.gather(w,n,`|r_ik|`),g=e.div(e.sumRows(e.mul(r,i),`dot`),e.unary(`clamp`,e.mul(o,c),1e-10,1/0)),v=e.unary(`acos`,e.scale(g,.95)),b=e.unary(`cos`,e.sub(e.repeatCols(v,u,`angles`),y(s.sections,[1,u],`sections`))),S=e.scale(e.unary(`pow`,e.unary(`clamp`,e.scale(e.add(b,y([1],[1],`one`)),.5),0,1/0),s.zeta),2),T=e.scale(e.add(o,c),.5),D=e.unary(`exp`,e.scale(e.unary(`square`,e.sub(e.repeatCols(T,l,`mean_distances`),y(s.shifts,[1,l],`shifts`))),-s.eta,`scale`,x(h,-2))),A=new Float32Array(l*u*l),j=new Float32Array(l*u*u);for(let e=0;e<l;e++)for(let t=0;t<u;t++)A[(e*u+t)*l+e]=1,j[(e*u+t)*u+t]=1;let M=e.mul(e.linear(D,y(A,[l*u,l],`expand`),void 0,`expand_radial`),e.linear(S,y(j,[l*u,u],`expand`),void 0,`expand_angle`)),N=y(Array(O).fill(s.cutoff),[O],`cutoff`,`triplet`),P=e.mul(M,e.mul(e.cutoff(`cosine`,o,N,s.cutoff),e.cutoff(`cosine`,c,N,s.cutoff))),F=_(E.map(([e,t,n])=>e*a+k[f[p.neighbor[t]]][f[p.neighbor[n]]]),d*a,`triplet`,`atom_pairs`);return e.reshape(e.segmentSum(P,F,`sum_per_element_pair`),[d,a*l*u],`atom`,`angular_aev`)}),j=e.scope(`aev`,()=>e.concatCols([T,A],`aev`)),M={atom:q(d),pair:{...J(p.center),label:`pair`},triplet:{label:`angular triplet`,owner:Int32Array.from(E,e=>e[0])},atom_species:{label:`neighbour element`,owner:Int32Array.from({length:d*i},(e,t)=>t/i|0)},atom_pairs:{label:`element pair`,owner:Int32Array.from({length:d*a},(e,t)=>t/a|0)}},N=null;e.scope(`networks`,()=>{r.symbols.forEach((t,n)=>{let i=K(d).filter(e=>f[e]===n);if(!i.length)return;let a=`atom_${t}`;M[a]={label:`${t} atom`,owner:i,atom:i,all:!0},e.scope(t,()=>{let n=e.gather(j,_(i,d,a,`atom`),`${t}_atoms`),o=null;for(let i=0;i<r.members;i++){let a=e.scope(`member${i}`,()=>{let a=r.layers[t].length-1,o=n;for(let n=0;n<a;n++)o=e.linear(o,this.params.get(`member${i}.${t}.${n}.weight`),this.params.get(`member${i}.${t}.${n}.bias`),n<a-1?`layer${n}`:`output`).withUnit(n<a-1?null:v),n<a-1&&(o=e.unary(`celu`,o,r.activation.alpha));return o});o=o?e.add(o,a):a}let s=e.scale(o,1/r.members,`ensemble_mean`),c=e.segmentSum(s,_(i,d,a,`atom`),`to_atoms`);N=N?e.add(N,c):c})})});let P=e.scope(`energy`,()=>e.scale(e.add(e.sumRows(N),y(Array.from(f,e=>r.self_energies[e]),[d],`self_energies`,`atom`)),r.hartree_to_ev,`to_eV`,b(g,v,-1)));return{energy:e.scope(`energy`,()=>e.sumAll(P,`total_energy`)),perAtom:P,positions:S,virialVectors:C,rows:M,graph:{center:p.center,neighbor:p.neighbor,shift:Float32Array.from(p.shiftVec),label:`symmetry-function pair`}}}};const ie={cutoff:4,cutoff_width:.5,sigma:.3,n_max:6,l_max:4,zeta:2};function ae(e,t){let n=new Float64Array(t+1);if(e<1e-12)return n[0]=1,n;if(e<20){let r=Math.exp(-e),i=1;for(let a=0;a<=t;a++){a>0&&(i*=e/(2*a+1));let t=1,o=1;for(let n=1;n<200&&(t*=e*e/2/(n*(2*a+2*n+1)),o+=t,!(t<1e-17*o));n++);n[a]=r*i*o}return n}let r=Math.exp(-2*e);n[0]=(1-r)/(2*e),t>=1&&(n[1]=(1+r)/(2*e)-(1-r)/(2*e*e));for(let r=1;r<t;r++)n[r+1]=n[r-1]-(2*r+1)/e*n[r];return n}function oe(e,t,n){let r=e.n_max,i=e.cutoff/r,a=Array.from({length:r},(n,a)=>{let o=e.cutoff*a/(r-1);return t.map(e=>Math.exp(-((e-o)**2)/(2*i*i)))}),o=Array.from({length:r},(e,i)=>Array.from({length:r},(e,r)=>a[i].reduce((e,i,o)=>e+i*a[r][o]*t[o]*t[o]*n[o],0))),s=Array.from({length:r},()=>new Float64Array(r));for(let e=0;e<r;e++)for(let t=0;t<=e;t++){let n=o[e][t];for(let r=0;r<t;r++)n-=s[e][r]*s[t][r];s[e][t]=e===t?Math.sqrt(n):n/s[t][t]}let c=[];for(let e=0;e<r;e++){let t=a[e].slice();for(let n=0;n<e;n++)for(let r=0;r<t.length;r++)t[r]-=s[e][n]*c[n][r];c.push(t.map(t=>t/s[e][e]))}return c}function se(e,t=160){let n=e.cutoff,r=n/600,i=Float64Array.from({length:601},(e,t)=>t*r),a=Float64Array.from({length:601},(e,t)=>r/3*(t===0||t===600?1:t%2?4:2)),o=oe(e,i,a),s=e.l_max,c=e.n_max*(s+1),l=e.sigma*e.sigma,u=t=>{let n=new Float64Array(c);for(let r=0;r<=600;r++){let c=i[r],u=Math.exp(-((c-t)**2)/(2*l)),d=ae(c*t/l,s),f=4*Math.PI*a[r]*c*c*u;for(let t=0;t<e.n_max;t++)for(let e=0;e<=s;e++)n[t*(s+1)+e]+=f*o[t][r]*d[e]}return n},d=n/(t-1),f=1e-5,p=new Float32Array(c*t),m=new Float32Array(c*t);for(let e=0;e<t;e++){let n=e*d,r=u(n),i=u(n+f),a=u(Math.max(n-f,0));for(let o=0;o<c;o++)p[o*t+e]=r[o],m[o*t+e]=(i[o]-a[o])/(n+f-Math.max(n-f,0))}return{V:p,D:m,K:t,C:c,h:d}}var ce=class{be;meta;kind=`krr`;params=/* @__PURE__ */ new Map;h;elements=[];report=null;tables;V;D;sparse=/* @__PURE__ */ new Float32Array;sparseZ=/* @__PURE__ */ new Int32Array;trainX=/* @__PURE__ */ new Float32Array;trainZ=/* @__PURE__ */ new Int32Array;trainOf=/* @__PURE__ */ new Int32Array;trainY=/* @__PURE__ */ new Float32Array;alpha=/* @__PURE__ */ new Float32Array;lambda=0;baseline=/* @__PURE__ */ new Map;F=0;constructor(e,t){this.be=e,this.meta=t,this.h={...ie,...t.soap},this.tables=se(this.h)}async fit(e,t,n,r){let i=performance.now(),a=this.meta,o=a.n_train??96,s=a.n_test??24,[c,l]=a.rattle??[.01,.08];this.elements=[...new Set(e.numbers)].sort((e,t)=>e-t);let u=this.elements.filter(e=>!t.elements.includes(e));if(u.length)throw Error(`the teacher (${n}) has no parameters for Z = ${u.join(`, `)}`);let d=12345,f=()=>(d=(d*1103515245+12345)%2147483648)/2147483648,p=()=>Math.sqrt(-2*Math.log(f()+1e-12))*Math.cos(2*Math.PI*f()),m=Array.from({length:o+s},(t,n)=>{let r=n===0?0:c+(l-c)*f();return{...e,positions:e.positions.map(e=>e.map(e=>e+r*p()))}}),h=[];for(let e=0;e<m.length;e++){e%8==0&&r(`KRR: labelling structure ${e+1}/${m.length} with ${n}…`);let i=new A(this.be),a=t.forward(i,m[e],{forces:!1});h.push((await this.be.read(a.energy.buf))[0]),i.release()}r(`KRR: SOAP power spectra of the training set…`);let g=[];for(let e of m){let t=new A(this.be),n=this.describe(t,e,!1).p;g.push(await this.be.read(n.buf)),this.F=n.cols,t.release()}let _=e.numbers.length,v=this.F,y=Int32Array.from(e.numbers),b=h.slice(0,o).reduce((e,t)=>e+t,0)/o/_;this.baseline=new Map(this.elements.map(e=>[e,b]));let x=a.sparse_per_element??24,S=[],C=[];for(let e of this.elements){let t=[];for(let n=0;n<o;n++)for(let r=0;r<_;r++)y[r]===e&&t.push(g[n].subarray(r*v,(r+1)*v));let n=new Float64Array(t.length).fill(1/0),r=0;for(let i=0;i<Math.min(x,t.length);i++){S.push(...t[r]),C.push(e);let i=-1;for(let e=0;e<t.length;e++){let a=0;for(let n=0;n<v;n++)a+=(t[e][n]-t[r][n])**2;n[e]=Math.min(n[e],a),(i<0||n[e]>n[i])&&(i=e)}r=i}}this.sparse=Float32Array.from(S),this.sparseZ=Int32Array.from(C);let w=C.length,T=(e,t)=>{let n=0;for(let r=0;r<v;r++)n+=e[r]*t[r];return n**this.h.zeta},E=e=>this.sparse.subarray(e*v,(e+1)*v),D=e=>Float64Array.from({length:w},(t,n)=>{let r=0;for(let t=0;t<_;t++)y[t]===C[n]&&(r+=T(g[e].subarray(t*v,(t+1)*v),E(n)));return r}),O=Array.from({length:o+s},(e,t)=>D(t)),k=Array.from({length:w},(e,t)=>Float64Array.from({length:w},(e,n)=>C[t]===C[n]?T(E(t),E(n)):0)),j=h.map(e=>e-_*b),M=Array.from({length:w},(e,t)=>Float64Array.from({length:w},(e,n)=>{let r=0;for(let e=0;e<o;e++)r+=O[e][t]*O[e][n];return r})),N=M.reduce((e,t,n)=>e+t[n],0)/w,P=k.reduce((e,t,n)=>e+t[n],0)/w,F=Float64Array.from({length:w},(e,t)=>{let n=0;for(let e=0;e<o;e++)n+=O[e][t]*j[e];return n}),I=e=>le(M.map((t,n)=>t.map((t,r)=>t+e*k[n][r]+(n===r?1e-10*N:0))),F),L=(e,t,n)=>Math.sqrt(K(n-t,t).reduce((t,n)=>t+((O[n].reduce((t,n,r)=>t+n*e[r],0)-j[n])/_)**2,0)/(n-t)),R=a.regularization==null?[1e-8,1e-7,1e-6,1e-5,1e-4,.001,.01,.1,1]:[a.regularization],z={rel:R[0],alpha:new Float64Array(w),err:1/0};for(let e of R)try{let t=I(N/P*e),n=L(t,o,o+s);n<z.err&&(z={rel:e,alpha:t,err:n})}catch{}this.lambda=z.rel*(N/P);let B=z.alpha;this.alpha=Float32Array.from(B);let V=(e,t)=>L(B,e,t);this.trainX=new Float32Array(o*_*v);for(let e=0;e<o;e++)this.trainX.set(g[e],e*_*v);this.trainZ=Int32Array.from({length:o*_},(e,t)=>y[t%_]),this.trainOf=Int32Array.from({length:o*_},(e,t)=>t/_|0),this.trainY=Float32Array.from(j.slice(0,o)),this.report={teacher:n,nTrain:o,nTest:s,nSparse:w,rmseTrain:V(0,o),rmseTest:V(o,o+s),lambda:this.lambda,lambdaRelative:z.rel,ms:performance.now()-i},this.params.clear(),this.params.set(`sparse_points`,this.constParam(`sparse_points`,this.sparse,[w,v])),this.params.set(`alpha`,this.constParam(`alpha`,this.alpha,[1,w])),r(`KRR: fitted ${w} sparse environments on ${o} structures; test RMSE ${(1e3*this.report.rmseTest).toFixed(2)} meV/atom`)}constParam(e,t,n){let r=new A(this.be).tensor(n,this.be.upload(t),e);return r.name=e,r}describe(e,t,n){let r=this.h,i=t.numbers.length,a=r.l_max,o=(a+1)**2,s=r.n_max,c=this.elements.length?this.elements:[...new Set(t.numbers)].sort((e,t)=>e-t),l=c.length,u=Int32Array.from(t.numbers,e=>c.indexOf(e)),d=U(t,r.cutoff),f=d.center.length,p=(t,n,r,i)=>e.index(Int32Array.from(t),n,{out:r,src:i}),m=(t,n,r,i)=>e.constant(Float32Array.from(t),n,r,i);this.V??=new A(this.be).tensor([this.tables.C,this.tables.K],this.be.upload(this.tables.V),`radial_integrals`),this.D??=new A(this.be).tensor([this.tables.C,this.tables.K],this.be.upload(this.tables.D),`radial_slopes`);let g=m(t.positions.flat(),[i,3],`positions`,`atom`).withUnit(h);g.requiresGrad=n;let{v:_,d:v}=e.scope(`geometry`,()=>{let t=e.add(e.sub(e.gather(g,p(d.neighbor,i,`pair`,`atom`),`r_j`),e.gather(g,p(d.center,i,`pair`,`atom`),`r_i`)),m(d.shiftVec,[f,3],`shift`,`pair`));return{v:t,d:e.rowNorm(t)}}),y=e.scope(`radial`,()=>{let t=e.cutoff(`cosine`,v,m(Array(f).fill(r.cutoff),[f],`cutoff`,`pair`),r.cutoff_width);return e.mul(e.spline(v,this.V,this.D,this.tables.h,`radial_integrals`),t)}),b=e.scope(`angular`,()=>e.sph(e.div(_,v),a)),x=e.scope(`density`,()=>{let t=new Float32Array(s*o*s*(a+1)),n=new Float32Array(s*o*o);for(let e=0;e<s;e++)for(let r=0;r<=a;r++)for(let i=r*r;i<(r+1)**2;i++)t[(e*o+i)*s*(a+1)+e*(a+1)+r]=1,n[(e*o+i)*o+i]=1;let r=e.mul(e.linear(y,m(t,[s*o,s*(a+1)],`expand`),void 0,`per_lm`),e.linear(b,m(n,[s*o,o],`expand`),void 0,`per_n`)),c=p(Array.from(d.center,(e,t)=>e*l+u[d.neighbor[t]]),i*l,`pair`,`atom_element`);return e.reshape(e.segmentSum(r,c,`sum_per_element`),[i*l*s,o],`atom_channel`,`coefficients`)});return{positions:g,v:_,p:e.scope(`power_spectrum`,()=>{let t=e.power(x,i,l*s,a);return t.kind=`atom`,e.div(t,e.rowNorm(t),`normalise`)}),rows:{atom:q(i),pair:{...J(d.center),label:`pair`},atom_element:{label:`neighbour element`,owner:Int32Array.from({length:i*l},(e,t)=>t/l|0)},atom_channel:{label:`element / radial channel`,owner:Int32Array.from({length:i*l*s},(e,t)=>t/(l*s)|0)}},spec:u,graph:{center:d.center,neighbor:d.neighbor,shift:Float32Array.from(d.shiftVec),label:`SOAP neighbour`}}}pow(e,t){let n=this.h.zeta;return n===1?t:n===2?e.unary(`square`,t):n===4?e.unary(`square`,e.unary(`square`,t)):e.unary(`pow`,t,n)}forward(e,t,n={}){if(!this.report)throw Error(`the KRR model has not been fitted`);let r=t.numbers.length,i=this.F,a=this.sparseZ.length,o=(t,n,r,i)=>e.constant(Float32Array.from(t),n,r,i),s=this.params.get(`sparse_points`),c=e=>o(Array.from({length:e.length*a},(t,n)=>+(e[n/a|0]===this.sparseZ[n%a])),[e.length,a],`same_element`),l=e.scope(`fit`,()=>{let t=e.mul(this.pow(e,e.linear(s,s,void 0,`sparse_overlap`)),c(this.sparseZ),`K_MM`),n=o(this.trainX,[this.trainZ.length,i],`training_environments`),r=e.mul(this.pow(e,e.linear(n,s,void 0,`overlap`)),c(this.trainZ),`atomic_kernels`),l=e.segmentSum(r,e.index(this.trainOf,this.trainY.length),`K_NM`),u=e.transpose(l),d=e.add(e.linear(u,u,void 0,`K_NM^T K_NM`),e.scale(t,this.lambda,`lambda_K_MM`),`normal_matrix`),f=e.linear(u,o(this.trainY,[1,this.trainY.length],`energies`),void 0,`K_NM^T y`);return e.opaque(`cholesky_solve`,[d,f],this.alpha,[1,a])}),u=e.scope(`soap`,()=>this.describe(e,t,!!n.forces)),d=e.scope(`kernel`,()=>e.mul(this.pow(e,e.linear(u.p,s,void 0,`overlap`)),c(t.numbers),`kernel`)),f=e.scope(`regression`,()=>{let n=e.sumRows(e.linear(d,l,void 0,`weighted_sum`)).withUnit(g);return e.add(n,o(t.numbers.map(e=>this.baseline.get(e)??0),[r],`baseline`,`atom`),`atomic_energy`)});return{energy:e.scope(`energy`,()=>e.sumAll(f,`total_energy`)),perAtom:f,positions:u.positions,virialVectors:u.v,rows:u.rows,graph:u.graph}}};function le(e,t){let n=t.length,r=Array.from({length:n},()=>new Float64Array(n));for(let t=0;t<n;t++)for(let n=0;n<=t;n++){let i=e[t][n];for(let e=0;e<n;e++)i-=r[t][e]*r[n][e];if(t===n){if(i<=0)throw Error(`KRR: the normal matrix is not positive definite; raise the regularisation`);r[t][t]=Math.sqrt(i)}else r[t][n]=i/r[n][n]}let i=new Float64Array(n);for(let e=0;e<n;e++){let n=t[e];for(let t=0;t<e;t++)n-=r[e][t]*i[t];i[e]=n/r[e][e]}let a=new Float64Array(n);for(let e=n-1;e>=0;e--){let t=i[e];for(let i=e+1;i<n;i++)t-=r[i][e]*a[i];a[e]=t/r[e][e]}return a}var ue=class{be;meta;kind=`physnet`;params=/* @__PURE__ */ new Map;c;get elements(){return this.meta.elements}constructor(e,t,n){this.be=e,this.meta=t;let r=t.config;if(r.max_degree!==0)throw Error(`PhysNet with max_degree = ${r.max_degree} is not supported (only the invariant max_degree = 0)`);if(this.c={...r,charges:r.charges??!1,n_refinement_blocks:r.n_refinement_blocks??r.n_res??3,zbl:r.zbl??!0,zbl_cuton:r.zbl_cuton??.1,zbl_cutoff:r.zbl_cutoff??.6,use_energy_bias:r.use_energy_bias??!1,include_electrostatics:r.include_electrostatics??!0,electrostatics_damping_sigma:r.electrostatics_damping_sigma??4,switch_start:r.switch_start??1,switch_end:r.switch_end??10,electrostatics_off_start:r.electrostatics_off_start??8,electrostatics_off_end:r.electrostatics_off_end??10},this.c.n_refinement_blocks<1)throw Error(`PhysNet needs at least one refinement block`);for(let[t,r]of n)this.params.set(t,new O(r.shape,e.upload(r.data),!1,t));let i=(e,t)=>Array.from(n.get(`repulsion.${e}`)?.data??t,Math.abs);this.zbl={aCoef:i(`a_coefficient`,[.529177249*.8854])[0],aExp:i(`a_exponent`,[.23])[0],phiC:i(`phi_coefficients`,[.18175,.50986,.28022,.02817]),phiE:i(`phi_exponents`,[3.1998,.94229,.4029,.20162])}}zbl;p(e){let t=this.params.get(e);if(!t)throw Error(`PhysNet: missing parameter ${e}`);return t}forward(e,t,n={}){if(H(t))throw Error(`periodic structures are not supported by this PhysNet`);let r=this.c,i=t.numbers.length;r.features;let a=r.num_basis_functions,o=r.charges&&r.include_electrostatics,s=U(t,Math.max(r.cutoff,o?r.electrostatics_off_end:0,r.zbl?r.zbl_cutoff:0)),c=s.center.length,l=e=>K(c).filter(e),u=l(e=>s.dist[e]<r.cutoff),d=o?l(e=>s.dist[e]<r.electrostatics_off_end):/* @__PURE__ */ new Int32Array,f=r.zbl?l(e=>s.dist[e]<r.zbl_cutoff):/* @__PURE__ */ new Int32Array,p=(t,n,r,i)=>e.index(Int32Array.from(t),n,{out:r,src:i}),m=(t,n,r,i)=>e.constant(Float32Array.from(t),n,r,i),v=m([1],[1],`one`),y=p(t.numbers,r.max_atomic_number+1,`atom`),S=m(t.positions.flat(),[i,3],`positions`,`atom`).withUnit(h);S.requiresGrad=!!n.forces;let C=e.scope(`geometry`,()=>e.add(e.sub(e.gather(S,p(s.neighbor,i,`pair`,`atom`),`r_j`),e.gather(S,p(s.center,i,`pair`,`atom`),`r_i`)),m(s.shiftVec,[c,3],`shift`,`pair`))),w=e.scope(`basis`,()=>{let t=e.gather(C,p(u,c,`mp_pair`,`pair`),`pairs_in_cutoff`),n=e.rowNorm(t),i=e.unary(`acos`,e.add(e.scale(e.unary(`exp`,e.unary(`neg`,n)),2),m([-1],[1],`minus_one`))),o=e.unary(`cos`,e.mul(e.repeatCols(i,a,`angles`),m(K(a),[1,a],`orders`)),0,0),s=e.scale(e.unary(`square`,n),1/(r.cutoff*r.cutoff),`scale`,x(h,-2)),l=e.unary(`exp`,e.add(e.unary(`neg`,e.unary(`pow`,e.add(e.unary(`neg`,s),v),-1)),v),0,0);return e.mul(o,l,`basis`)}),T=p(Array.from(u,e=>s.center[e]),i,`mp_pair`,`atom`),E=p(Array.from(u,e=>s.neighbor[e]),i,`mp_pair`,`atom`),D=e.scope(`embedding`,()=>e.gather(this.p(`Embed_0.embedding`),y,`embedding`)),O=0,k=(t,n,r=!0)=>e.linear(t,this.p(`${n}.kernel`),r&&this.params.has(`${n}.bias`)?this.p(`${n}.bias`):void 0,n.split(`.`)[0]);for(let t=0;t<r.num_iterations;t++)D=e.scope(`interaction.${t}`,()=>{let n=e.scope(`message`,()=>{let n=k(w,`MessagePass_${t}.filter.0+`,!1),r=e.gather(D,E,`x_j`),i=e.mul(e.mul(n,r,`filter_x_j`),this.p(`MessagePass_${t}.tensor.kernel`),`tensor_product`);return e.segmentSum(i,T,`sum_messages`)});return e.scope(`refinement`,()=>{let t=n,i=n;for(let n=0;n<r.n_refinement_blocks;n++)i=k(e.add(t,e.silu(t)),`Dense_${O++}.0+`),t=e.add(t,i,`residual`);return e.silu(k(i,`Dense_${O++}.0+`))})});let A=(t,n,r)=>{let i=k(k(D,`Dense_${O++}.0+`,!1),`Dense_${O++}`,!1),a=e.sumRows(i,`per_atom`).withUnit(t===`charges`?_:g);return r?e.add(a,e.gather(this.p(n),y,`${n}`)):a},j=e.scope(`energy_head`,()=>A(`energy`,`energy_bias`,r.use_energy_bias)),M={atom:q(i),pair:{label:`pair`,owner:s.center},mp_pair:{label:`pair`,owner:Int32Array.from(u,e=>s.center[e]),edge:K(u.length)}},N={};if(r.charges){let t=e.scope(`charges`,()=>A(`charges`,`charge_bias`,!0));if(N.charge=t,o){M.es_pair={label:`electrostatic pair`,owner:Int32Array.from(d,e=>s.center[e])};let n=e.scope(`electrostatics`,()=>{let n=e.gather(C,p(d,c,`es_pair`,`pair`),`pairs`),a=e.sumRows(e.unary(`square`,n),`r2`),o=e.unary(`sqrt`,e.unary(`clamp`,a,1e-4,1/0)),l=e.unary(`switch`,o,r.switch_start,r.switch_end),u=e.add(e.unary(`neg`,e.unary(`switch`,o,r.electrostatics_off_start,r.electrostatics_off_end)),v,`off`),f=e.div(l,e.unary(`sqrt`,e.add(a,v)),`short_range`),y=e.div(e.add(e.unary(`neg`,l),v),e.add(o,m([1e-6],[1],`eps`)),`long_range`),S=e.add(f,y,`1/r`);r.electrostatics_damping_sigma>0&&(S=e.mul(S,e.unary(`erf`,e.scale(o,1/r.electrostatics_damping_sigma,`scale`,x(h,-1))),`damped`));let w=e.add(e.scale(o,1/r.switch_end**2,`scale`,x(h,-2)),m([1e-6/r.switch_end**2-2/r.switch_end],[1],`shift`),`eshift`),T=e.unary(`clamp`,e.gather(t,p(Array.from(d,e=>s.center[e]),i,`es_pair`,`atom`),`q_i`),-10,10),E=e.unary(`clamp`,e.gather(t,p(Array.from(d,e=>s.neighbor[e]),i,`es_pair`,`atom`),`q_j`),-10,10),D=e.mul(e.scale(e.mul(T,E),7.199822675975274,`scale`,b(b(g,h),x(_,-2))),e.mul(e.add(S,w),u),`pair_energy`);return e.segmentSum(D,p(Array.from(d,e=>s.center[e]),i,`es_pair`,`atom`),`per_atom`)});j=e.add(j,n,`with_electrostatics`)}}if(f.length){M.zbl_pair={label:`short-range pair`,owner:Int32Array.from(f,e=>s.center[e])};let n=e.scope(`repulsion`,()=>{let{aCoef:n,aExp:a,phiC:o,phiE:l}=this.zbl,u=Math.hypot(...o),d=Array.from(f,e=>t.numbers[s.center[e]]),_=Array.from(f,e=>t.numbers[s.neighbor[e]]),v=d.map((e,t)=>n/(e**a+_[t]**a)),y=e.gather(C,p(f,c,`zbl_pair`,`pair`),`pairs`),S=e.unary(`clamp`,e.rowNorm(y),1e-8,1/0),w=e.div(S,m(v,[f.length],`screening_length`,`zbl_pair`).withUnit(h)),T=e.linear(e.unary(`exp`,e.mul(e.repeatCols(w,4,`x`),m(l.map(e=>-e),[1,4],`exponents`))),m(o.map(e=>e/u),[1,4],`coefficients`),void 0,`screening`),E=e.unary(`clamp`,e.scale(e.add(e.unary(`neg`,S),m([r.zbl_cutoff],[1],`cutoff`)),1/(r.zbl_cutoff-r.zbl_cuton),`scale`,x(h,-1)),0,1),D=e.mul(e.mul(e.unary(`square`,E),E),e.add(e.mul(e.add(e.scale(E,6),m([-15],[1],`c`)),E),m([10],[1],`c`)),`switch`),O=e.div(m(d.map((e,t)=>7.19982273921275*e*_[t]),[f.length],`ZiZj`).withUnit(b(g,h)),S,`coulomb`),k=e.mul(e.mul(O,e.sumRows(T)),D,`pair_energy`).withUnit(g);return e.segmentSum(k,p(Array.from(f,e=>s.center[e]),i,`zbl_pair`,`atom`),`per_atom`)});j=e.add(j,n,`with_repulsion`)}let P=e.scope(`energy`,()=>e.sumAll(j,`total_energy`)),F={center:Int32Array.from(u,e=>s.center[e]),neighbor:Int32Array.from(u,e=>s.neighbor[e]),shift:new Float32Array(3*u.length),label:`message-passing`};return{energy:P,perAtom:j,positions:S,rows:M,graph:F,extras:N}}};function de(e,t,n,r,i){let a=1/r,o=r=>{let o=new Float64Array(t),s=new Float64Array(t);for(let t=0;t<e.center.length;t++){let n=e.center[t],a=(e.dist[t]-(r[n]-i))/i;if(a<=0)o[n]+=1;else if(a<1){let[a,c]=l(`bump`,e.dist[t],r[n],i);o[n]+=a,s[n]-=c/i}}for(let e=0;e<t;e++){let t=r[e]*a;o[e]+=n*t**3,s[e]+=3*n*t*t*a}return[o,s]},s=new Float64Array(t),c=new Float64Array(t).fill(r),u=new Float64Array(t).fill(.5*r);for(let e=0;e<10;e++){let[e,r]=o(u);for(let i=0;i<t;i++){let t=e[i]-n;t<=0?s[i]=u[i]:c[i]=u[i];let a=u[i]-t/Math.max(r[i],1e-6);u[i]=a>=s[i]&&a<=c[i]?a:.5*(s[i]+c[i])}}let[d,f]=o(u),p=new Float64Array(t);for(let e=0;e<t;e++){let t=u[e]-(d[e]-n)/Math.max(f[e],1e-6);p[e]=Math.min(Math.max(t,r/16),r)}return{r:u,dn:f.map(e=>Math.max(e,1e-6)),atomic:p}}function fe(e,t,n){let r=e.numbers.length,i=Int32Array.from(e.numbers,e=>{let t=n[e]??-1;if(t<0)throw Error(`element Z=${e} is not supported by this model`);return t}),a=U(e,t.cutoff),o=null,s=[],c=[];if(t.num_neighbors_adaptive!=null){o=de(a,r,t.num_neighbors_adaptive,t.cutoff,t.cutoff_width_adaptive);for(let e=0;e<a.center.length;e++){let t=.5*(o.atomic[a.center[e]]+o.atomic[a.neighbor[e]]);a.dist[e]<=t&&(s.push(e),c.push(t))}}else for(let e=0;e<a.center.length;e++)s.push(e),c.push(t.cutoff);s.length;let l=Int32Array.from(s,e=>a.center[e]),u=Int32Array.from(s,e=>a.neighbor[e]),d=new Int32Array(r+1);for(let e of l)d[e+1]++;for(let e=0;e<r;e++)d[e+1]+=d[e];let f=(e,t,n,r,i)=>`${e},${t},${n},${r},${i}`,p=/* @__PURE__ */ new Map;s.forEach((e,t)=>p.set(f(a.center[e],a.neighbor[e],a.shift[3*e],a.shift[3*e+1],a.shift[3*e+2]),t));let m=Int32Array.from(s,e=>p.get(f(a.neighbor[e],a.center[e],-a.shift[3*e],-a.shift[3*e+1],-a.shift[3*e+2]))??-1);return{N:r,species:i,raw:a,adaptive:o,keep:Int32Array.from(s),center:l,neighbor:u,reverse:m,offsets:d,pairCutoff:Float64Array.from(c)}}function pe(e){let t=e.N,n=e.keep.length,r=e.raw.center.length,i=new Int32Array(r).fill(-1);e.keep.forEach((e,t)=>i[e]=t);let a=new Int32Array(t+n),o=new Int32Array(t+n).fill(-1),s=new Int32Array(t+n).fill(-1);for(let n=0,r=0;n<t;n++){a[r]=n,o[r++]=n;for(let t=e.offsets[n];t<e.offsets[n+1];t++)a[r]=n,s[r++]=t}let c=Int32Array.from([...K(t),...e.center]);return{atom:q(t),edge:J(e.center),raw_edge:{label:`candidate edge`,owner:e.raw.center,edge:i},token:{label:`token`,owner:a,atom:o,edge:s},stack:{label:`atom / edge`,owner:c,atom:Int32Array.from({length:t+n},(e,n)=>n<t?n:-1),edge:Int32Array.from({length:t+n},(e,n)=>n<t?-1:n-t)}}}var me=class{be;meta;kind=`pet`;params=/* @__PURE__ */ new Map;h;get elements(){return this.meta.atomic_types}constructor(e,t,n){this.be=e,this.meta=t,this.h=t.hypers;for(let[t,r]of n)this.params.set(t,new O(r.shape,e.upload(r.data),!1,t))}p(e){let t=this.params.get(e);if(!t)throw Error(`missing weight ${e}`);return t}has(e){return this.params.has(e)}setTrainable(e){for(let t of this.params.values())t.requiresGrad=e}async exportTensors(){let e=/* @__PURE__ */ new Map;for(let[t,n]of this.params)e.set(t,{shape:n.shape,data:await this.be.read(n.buf)});return e}prepare(e){return fe(e,this.h,this.meta.species_to_index)}get hasNC(){return!!this.meta.non_conservative?.non_conservative_force}forward(e,t,n={}){let r=this.h,i=this.be,a=n.prepared??this.prepare(t),{N:o}=a,s=a.keep.length,c=a.raw.center.length,l=r.d_pet,u=e=>{let t=e.split(`.`);return/^\d+$/.test(t[t.length-1])?t.slice(-2).join(`.`):t[t.length-1]===`energy___0`?`last_layer`:t[t.length-1]},d=(t,n)=>e.linear(t,this.p(`${n}.weight`),this.has(`${n}.bias`)?this.p(`${n}.bias`):void 0,u(n)),f=r.normalization===`RMSNorm`?`rms`:`layer`,p=(t,n)=>e.norm(f,t,this.p(`${n}.weight`),this.has(`${n}.bias`)?this.p(`${n}.bias`):null,f===`rms`?1.1920928955078125e-7:1e-5),m=(t,n)=>{let i=d(t,`${n}.w_in`);if(r.activation===`SwiGLU`){let t=i.cols/2;return d(e.mul(e.sliceCols(i,0,t,`value`),e.sigmoid(e.sliceCols(i,t,t,`gate`))),`${n}.w_out`)}return d(e.silu(i),`${n}.w_out`)},_=(t,n)=>d(e.silu(d(t,`${n}.0`)),`${n}.2`),v=(t,n,r,i)=>e.index(t,n,{out:r,src:i}),y=v(a.species,this.meta.atomic_types.length,`atom`),x=v(Int32Array.from(a.neighbor,e=>a.species[e]),this.meta.atomic_types.length,`edge`),S=e.constant(Float32Array.from(t.positions.flat()),[o,3],`positions`,`atom`).withUnit(h);S.requiresGrad=!!n.forces;let{edgeVectors:C,kept:w,pair:T}=e.scope(`geometry`,()=>{let t=e.add(e.sub(e.gather(S,v(a.raw.neighbor,o,`raw_edge`,`atom`),`r_j`),e.gather(S,v(a.raw.center,o,`raw_edge`,`atom`),`r_i`)),e.constant(Float32Array.from(a.raw.shiftVec),[c,3],`shift`,`raw_edge`));if(!a.adaptive)return t.kind=`edge`,{edgeVectors:t,kept:t,pair:e.constant(Float32Array.from(a.pairCutoff),[s],`cutoff`,`edge`)};let n=a.adaptive,i=e.scope(`adaptive_cutoff`,()=>{let i=e.rowNorm(t),s=e.constant(Float32Array.from(a.raw.center,e=>n.r[e]),[c],`r_probe`,`raw_edge`),l=e.segmentSum(e.cutoff(`bump`,i,s,r.cutoff_width_adaptive),v(a.raw.center,o,`raw_edge`,`atom`),`neighbour_count`),u=r.num_neighbors_adaptive,d=Float32Array.from(n.r,e=>u*(e/r.cutoff)**3-u),f=e.add(l,e.constant(d,[o],`baseline`,`atom`)),p=e.div(f,e.constant(Float32Array.from(n.dn),[o],`slope`,`atom`)),m=e.sub(e.constant(Float32Array.from(n.r),[o],`root`,`atom`),p);return e.unary(`clamp`,m,r.cutoff/16,r.cutoff).withUnit(h)}),l=v(a.keep,c,`edge`,`raw_edge`),u=e.scale(e.add(e.gather(i,v(a.center,o,`edge`,`atom`),`cutoff_i`),e.gather(i,v(a.neighbor,o,`edge`,`atom`),`cutoff_j`)),.5);return{edgeVectors:t,kept:e.gather(t,l,`keep`),pair:u}}),{dist:E,cf:D}=e.scope(`cutoff`,()=>{let t=e.rowNorm(w);return{dist:t,cf:e.cutoff(r.cutoff_function===`Bump`?`bump`:`cosine`,t,T,r.cutoff_width)}}),O=new Int32Array(o+s),k=new Int32Array(o),A=new Int32Array(s),j=new Int32Array(o+1),M=new Int32Array(o+1),N=new Int32Array(o+s);for(let e=0,t=0;e<o;e++){j[e]=t,k[e]=t,N[t]=e,O[t++]=e;for(let n=a.offsets[e];n<a.offsets[e+1];n++)A[n]=t,N[t]=e,O[t++]=o+n;j[e+1]=t;let n=j[e+1]-j[e];M[e+1]=M[e]+r.num_heads*n*n}let P={nSeq:o,nTok:o+s,nProb:M[o],off:e.own(i.uploadI32(j)),poff:e.own(i.uploadI32(M)),tok2seq:e.own(i.uploadI32(N))},F=v(O,o+s,`token`,`stack`),I=v(k,o+s,`atom`,`token`),L=v(A,o+s,`edge`,`token`),R=v(a.reverse,s,`edge`,`edge`),z=e.scope(`cutoff`,()=>e.unary(`logclamp`,e.gather(e.concatRows(e.constant(new Float32Array(o).fill(1),[o],`self`,`atom`),D),F,`tokens`),1e-15)),B=e.scope(`cutoff`,()=>e.concatCols([w,E],`x4`)),V=[],ee=e=>({offset:M[e],atoms:Int32Array.from({length:j[e+1]-j[e]},(t,n)=>n===0?e:a.neighbor[a.offsets[e]+n-1])}),H=(t,n,i)=>{let a=r.d_node!==l;for(let o=0;o<r.num_attention_layers;o++){let s=`gnn_layers.${t}.trans.layers.${o}`;[n,i]=e.scope(`trans.${o}`,()=>{let t=e.scope(`tokens`,()=>{let t=a?d(n,`${s}.center_contraction`):n;return e.gather(e.concatRows(t,i),F,`tokens`)}),o=t=>{let n=d(t,`${s}.attention.input_linear`),i=e.attention(n,z,P,r.num_heads,r.attention_temperature);return V.push({node:e.tape[e.tape.length-1],heads:r.num_heads,tokens:ee}),d(i,`${s}.attention.output_linear`)},c,l;if(r.transformer_type===`PostLN`){let n=e.scope(`attention`,()=>p(e.add(t,o(t)),`${s}.norm_attention`));[c,l]=e.scope(`mlp`,()=>{let t=p(e.add(n,m(n,`${s}.mlp`)),`${s}.norm_mlp`);return[e.gather(t,I,`node_token`),e.gather(t,L,`edge_tokens`)]})}else[c,l]=e.scope(`attention`,()=>{let n=o(p(t,`${s}.norm_attention`));return[e.gather(n,I,`node_token`),e.add(i,e.gather(n,L,`edge_tokens`))]}),l=e.scope(`mlp`,()=>e.add(l,m(p(l,`${s}.norm_mlp`),`${s}.mlp`)));return a?[e.scope(`node_update`,()=>{let t=e.add(n,d(c,`${s}.center_expansion`));return e.add(t,m(p(t,`${s}.norm_center_features`),`${s}.center_mlp`))}),l]:[c,l]})}return[n,i]},te=(t,n,r)=>e.scope(`gnn.${t}`,()=>{let i=`gnn_layers.${t}`,a=e.scope(`edge_tokens`,()=>{let n=d(B,`${i}.edge_embedder`),a=t===0?[n,r]:[n,e.gather(this.p(`${i}.neighbor_embedder.weight`),x,`neighbor_embedding`),r];return _(e.concatCols(a),`${i}.compress`)});return H(t,n,a)}),U=(t,n)=>e.scope(`embedding`,()=>e.gather(this.p(`${t}.weight`),n,t)),W=this.conditioning(e,t,o),G=[],K=[],q=U(`edge_embedder`,x);if(r.featurizer_type===`feedforward`){let t=U(`node_embedders.0`,y);for(let n=0;n<r.num_gnn_layers;n++){let[r,i]=te(n,t,q);W&&(r=e.add(r,W)),t=r,q=e.scope(`gnn.${n}/message_passing`,()=>{let t=e.concatCols([i,e.gather(i,R,`reverse_messages`)]),r=_(e.norm(`layer`,t,this.p(`combination_norms.${n}.weight`),this.p(`combination_norms.${n}.bias`),1e-5),`combination_mlps.${n}`);return e.add(e.add(q,i),r)})}G.push(t),K.push(q)}else for(let t=0;t<r.num_gnn_layers;t++){let[n,r]=te(t,U(`node_embedders.${t}`,y),q);W&&(n=e.add(n,W)),G.push(n),K.push(r),q=e.scope(`gnn.${t}/message_passing`,()=>e.scale(e.add(q,e.gather(r,R,`reverse_messages`)),.5))}let J=v(a.center,o,`edge`,`atom`),ne=t=>{let n=null;return G.forEach((r,i)=>{let a=(n,r)=>e.scope(`${r}_head.${i}`,()=>{let a=`${r}_heads.${t}.${i}`,o=e.silu(d(e.silu(d(n,`${a}.0`)),`${a}.2`));return d(o,`${r}_last_layers.${t}.${i}.${t}___0`)}),o=a(r,`node`),s=e.segmentSum(e.mul(a(K[i],`edge`),D),J,`sum_edges`),c=e.add(o,s);n=n?e.add(n,c):c}),n},re=e.scope(`readout`,()=>ne(`energy`)),ie;if(n.nc&&this.hasNC){let t=this.meta.non_conservative.non_conservative_force.scale;ie=e.scope(`nc_forces`,()=>{let n=e.mul(ne(`non_conservative_force`),e.constant(Float32Array.from(a.species,e=>t[e]),[o],`scale`,`atom`)).withUnit(b(g,h,-1)),r=e.index(new Int32Array(o),1,{src:`atom`}),i=e.scale(e.segmentSum(n,r,`net_force`),1/o);return e.sub(n,e.gather(i,r,`mean_force`))})}let ae=Float32Array.from(a.species,e=>this.meta.composition_energies[e]),oe=e.scope(`energy`,()=>e.add(e.scale(e.sumRows(re),this.meta.energy_scale).withUnit(g),e.constant(ae,[o],`composition`,`atom`))),se=e.scope(`energy`,()=>e.sumAll(oe,`total_energy`)),ce={center:a.center,neighbor:a.neighbor,label:`neighbour`,shift:Float32Array.from({length:3*s},(e,t)=>a.raw.shiftVec[3*a.keep[t/3|0]+t%3])};return{energy:se,ncForces:ie,perAtom:oe,positions:S,edgeVectors:C,virialVectors:C,cutoffFactors:D,attention:V,prepared:a,layout:P,graph:ce,rows:pe(a)}}conditioning(e,t,n){return this.h.system_conditioning?e.scope(`conditioning`,()=>{let r=(t.charge??0)+this.h.max_charge,i=(t.spin??1)-1,a=(t,n)=>e.gather(this.p(t),e.index(Int32Array.of(n),this.p(t).shape[0])),o=e.concatCols([a(`system_conditioning.charge_embedding.weight`,r),a(`system_conditioning.spin_multiplicity_embedding.weight`,i)]),s=e.linear(o,this.p(`system_conditioning.project.0.weight`),this.p(`system_conditioning.project.0.bias`)),c=e.linear(e.silu(s),this.p(`system_conditioning.project.2.weight`),this.p(`system_conditioning.project.2.bias`));return e.gather(c,e.index(new Int32Array(n),1),`broadcast`)}):null}};async function he(e,t,n,r,i){switch(t){case`pet`:return G(n),new me(e,n,W(r));case`ani`:return new re(e,n,W(r));case`physnet`:return new ue(e,n,W(r));case`krr`:{let t=i.models.get(n.teacher);if(!t||t.kind===`krr`)throw Error(`KRR needs a loaded teacher model (got "${n.teacher}")`);if(!n.system)throw Error(`KRR needs a structure to fit on`);let r=new ce(e,n);return await r.fit(n.system,t,n.teacherLabel??n.teacher,i.progress),r}default:throw Error(`model kind "${t}" is not available yet`)}}async function ge(e,t){let n=e.thumbs;return n?n.call(e,t):Promise.all(t.map(t=>e.thumb(t.x,t.R,t.C,t.rows,t.cols)))}const Y=e=>{let t=e.length>1?e[e.length-1]:1;return[e.reduce((e,t)=>e*t,1)/Math.max(t,1),t]};function _e(e,t){let n=/* @__PURE__ */ new Map;e.tape.forEach((e,t)=>n.set(e.out.id,t));let r=/* @__PURE__ */ new Map;for(let[e,n]of t.params)r.set(n.id,e);return e.tape.map(e=>({op:e.op,scope:e.scope,shape:e.out.shape,kind:e.out.kind,inputs:e.inputs.filter(e=>!r.has(e.id)).map(e=>n.get(e.id)??-1),params:e.inputs.filter(e=>r.has(e.id)).map(e=>r.get(e.id)),grad:!!e.out.grad,unit:e.out.unit??null}))}const ve=e=>e.map(e=>`${e.scope}|${e.op}|${e.kind}|${+!!e.grad}`).join(`;`);function ye(e,t,n){if(!e||e.owner.length!==t)return null;if(e.all){if(t<=64)return Int32Array.from({length:t},(e,t)=>t);let e=Math.max(0,Math.min(n-32,t-64));return Int32Array.from(/* @__PURE__ */ new Set([n,...Array.from({length:63},(t,n)=>e+n)]))}let r=[];return e.owner.forEach((e,t)=>{e===n&&r.length<64&&r.push(t)}),Int32Array.from(r)}async function be(e,t,n,r){let i=t.be,a=Math.min(Math.max(r.selected,0),n.positions.shape[0]-1),o=[],s=(e,t,n)=>{let[r,a]=Y(t),s=Math.min(a,128);if(!n)return{x:e,R:r,C:a,rows:Math.max(1,Math.min(r,64)),cols:s};let c=n.length,l=i.zeros(Math.max(c*a,1)),u=i.uploadI32(n);return c&&i.gather(e,u,l,c,a,!1),o.push(l,u),{x:l,R:c,C:a,rows:Math.max(c,1),cols:s}},c=e.tape.map(e=>ye(e.out.kind?n.rows[e.out.kind]:void 0,Y(e.out.shape)[0],a)),l=e.tape.map((e,t)=>s(e.out.buf,e.out.shape,c[t])),u=e.tape.map((e,t)=>e.out.grad?t:-1).filter(e=>e>=0),d=u.map(t=>s(e.tape[t].out.grad,e.tape[t].out.shape,c[t])),f=[...t.params.values()],p=f.map(e=>{let[t,n]=Y(e.shape);return{x:e.buf,R:t,C:n,rows:Math.min(t,64),cols:Math.min(n,64)}}),m=[],h=[],g=(e,t)=>{let[n,r]=Y(t),a=i.zeros(n*r),s=i.zeros(n);i.unary(`square`,e,a,n*r,0,0),i.rowSum(a,s,n,r,!1),o.push(a,s),m.push(s)};e.tape.forEach((e,t)=>{e.out.kind&&n.rows[e.out.kind]&&n.rows[e.out.kind].owner.length===Y(e.out.shape)[0]&&(g(e.out.buf,e.out.shape),h.push([t,`v`]),e.out.grad&&(g(e.out.grad,e.out.shape),h.push([t,`g`])))});let _=await ge(i,[...l,...d,...p]),v=m.length?i.readMany?await i.readMany(m):await Promise.all(m.map(e=>i.read(e))):[];o.forEach(e=>i.free(e));let y=e.tape.map(()=>null),b=e.tape.map(()=>null);h.forEach(([t,n],r)=>{let i=Y(e.tape[t].out.shape)[1];(n===`v`?y:b)[t]=v[r].map(e=>Math.sqrt(e/i))});let x=_.slice(0,l.length),S=e.tape.map(()=>null);u.forEach((e,t)=>S[e]=_[l.length+t]);let C={};f.forEach((e,t)=>C[e.name]={value:_[l.length+d.length+t],grad:null});let w=n.attention??[],T=w.map(e=>e.node.aux.probs.buf),E=T.length?i.readMany?await i.readMany(T):await Promise.all(T.map(e=>i.read(e))):[],D=w.map((t,n)=>{let{offset:r,atoms:i}=t.tokens(a),o=i.length;return{heads:t.heads,n:o,atoms:i,op:e.tape.indexOf(t.node),probs:E[n].slice(r,r+t.heads*o*o)}});return{values:x,grads:S,params:C,paramShapes:Object.fromEntries([...t.params].map(([e,t])=>[e,t.shape])),attention:D,rowSel:c,selected:a,norms:y,gradNorms:b}}const X=(e,t=[])=>self.postMessage(e,t);let Z=new u;const xe=/* @__PURE__ */ new Map;let Q=null,$=``,Se=0;async function Ce(e,t,n,r){let i=t.numbers.filter(t=>!e.elements.includes(t));if(i.length)throw Error(`this model has no parameters for element Z = ${[...new Set(i)].join(`, `)}`);r!==`conservative`&&!e.hasNC&&(r=`conservative`);let a=r!==`direct`,o=new A(Z);try{let i=performance.now(),s=e.forward(o,t,{forces:a,nc:r!==`conservative`});await Z.sync();let c=performance.now();a&&o.backward(s.energy),await Z.sync();let l=performance.now(),u=[s.energy.buf,s.perAtom.buf],d=a&&!!s.virialVectors?.grad;a&&u.push(s.positions.grad),d&&u.push(s.virialVectors.buf,s.virialVectors.grad),s.ncForces&&u.push(s.ncForces.buf);let f=Object.entries(s.extras??{});u.push(...f.map(([,e])=>e.buf));let p=Z.readMany?await Z.readMany(u):await Promise.all(u.map(e=>Z.read(e))),m=_e(o,e),h=e.kind+ve(m),g=await be(o,e,s,{selected:n}),_={id:++Se,positions:t.positions,cell:t.cell,numbers:t.numbers,mode:r,energy:p[0][0],energies:p[1],trace:{...g,shapes:o.tape.map(e=>e.out.shape),topology:h===$?void 0:m,rows:s.rows,graph:s.graph,ms:{forward:c-i,backward:a?l-c:0,capture:performance.now()-l}}};$=h,a&&(_.forces=p[2].map(e=>-e));let v=te(t);if(d&&Number.isFinite(v)){let[e,t]=p.slice(3,5),n=[[0,0,0],[0,0,0],[0,0,0]];for(let r=0;r<e.length/3;r++)for(let i=0;i<3;i++)for(let a=0;a<3;a++)n[i][a]+=e[3*r+i]*t[3*r+a]/v;_.stress=n}return s.ncForces&&(_.ncForces=p[p.length-1-f.length]),f.length&&(_.atomProps=Object.fromEntries(f.map(([e],t)=>[e,p[p.length-f.length+t]]))),_}finally{o.release()}}async function we(e){switch(e.type){case`init`:if(e.backend!==`cpu`)try{Z=await R.create()}catch(t){if(e.backend===`webgpu`)throw t}X({type:`ready`,backend:Z.name,adapter:Z instanceof R?Z.adapterName:`JavaScript (no WebGPU)`});break;case`loadModel`:{let t=await he(Z,e.kind,e.meta,e.weights,{models:xe,progress:e=>X({type:`progress`,text:e})});xe.set(e.id,t),e.activate&&(Q=t,$=``);let n=[...t.params.values()].reduce((e,t)=>e+t.size,0),r=t.report?{...e.meta,system:void 0,report:t.report,elements:t.elements}:e.meta;X({type:`model`,id:e.id,kind:t.kind,label:e.label,meta:r,nParams:n,hasNC:!!t.hasNC,activate:e.activate,elements:t.elements});break}case`use`:Q=xe.get(e.id)??Q,$=``;break;case`evaluate`:{if(!Q)throw Error(`load a model first`);let t=await Ce(Q,e.system,e.selected,e.mode),n=[];for(let e of[...t.trace.values,...t.trace.grads])e&&n.push(e.data.buffer);X({type:`pass`,pass:t},n);break}}}self.onmessage=e=>{we(e.data).catch(e=>X({type:`error`,text:e?.message??String(e)}))};

assets/index-BvLUcW0X.css

@@ -0,0 +1 @@
+:root{--bg:#0c0e13;--panel:#12151c;--line:#232835;--text:#d7dbe4;--muted:#8a93a6;--value:#f99e29;--value-neg:#3b82f5;--grad:#a855f7;--grad-pos:#22c55e;--weight:#8cd1fa;--attn:#ffdb52;--accent:#f99e29;--lightningcss-light: ;--lightningcss-dark:initial;color-scheme:dark}*{box-sizing:border-box}[hidden]{display:none!important}html,body{background:var(--bg);height:100%;color:var(--text);margin:0;font:14px/1.5 system-ui,-apple-system,Segoe UI,sans-serif}body{grid-template-rows:auto minmax(0,1fr) auto;display:grid;overflow:hidden}body>:not(header,main,footer){position:fixed}button,select,.file{font:inherit;color:var(--text);border:1px solid var(--line);cursor:pointer;background:#1a1e28;border-radius:6px;padding:4px 10px}button:hover,.file:hover{border-color:#3a4256}button.on{border-color:var(--accent);color:var(--accent)}.file input{display:none}header{border-bottom:1px solid var(--line);background:var(--panel);flex-wrap:wrap;align-items:center;gap:12px;padding:8px 14px;display:flex}header label{color:var(--muted);align-items:center;gap:6px;display:flex}.brand{letter-spacing:.02em;margin-right:8px;font-size:18px;font-weight:700}.brand span{color:var(--accent)}#status{color:var(--muted);font-variant-numeric:tabular-nums;text-align:right;margin-left:auto;font-size:13px}#status b{color:var(--text);font-weight:600}main{grid-template-columns:minmax(300px,400px) minmax(0,1fr);min-height:0;display:grid}#walk{border-right:1px solid var(--line);background:var(--panel);scroll-behavior:smooth;overflow-y:auto}#chapter-head{z-index:2;-webkit-backdrop-filter:blur(4px);backdrop-filter:blur(4px);border-bottom:1px solid var(--line);background:#12151cf5;padding:8px 20px;position:sticky;top:0}#intro{color:#c3c9d5;padding:16px 20px 4px}#intro h1{color:#eef0f4;margin:0 0 8px;font-size:22px}#intro p{margin:0 0 10px}#intro figure{margin:12px 0}#intro figcaption{color:var(--muted);margin-top:4px;font-size:12px}.keys{color:var(--muted);font-size:12px}kbd{color:var(--text);background:#1f2430;border:1px solid #343b4c;border-bottom-width:2px;border-radius:4px;padding:0 4px;font:11px ui-monospace,monospace}#article{padding:0 20px 60vh}#article h2.sec{color:var(--c);border-left:3px solid var(--c);margin:28px 0 6px;padding-left:9px;font-size:15px;line-height:1.35}#article h2.sec span{color:#7d8698;font-weight:500}#article h2.sec.bwd{--c:var(--grad)}#article .lead{color:#9aa3b5;font-size:13.5px}#article .lead p{margin:0 0 8px}#article p.step{cursor:pointer;opacity:.5;border-left:2px solid #0000;border-radius:6px;margin:4px 0;padding:6px 10px 6px 11px;transition:opacity .25s,background .25s}#article p.step:hover{opacity:.85}#article p.step.on{opacity:1;border-left-color:var(--c);background:#1b202c}#article p.step.bwd.on{border-left-color:var(--grad)}#article .dim{color:#7d8698;font-variant-numeric:tabular-nums;white-space:nowrap;font-size:12px}#article b.atom{color:#fff;background:#2a3040;border-radius:3px;padding:0 4px;font-weight:600}button.primary{background:var(--accent);color:#111;border-color:var(--accent);font-weight:600}button.primary:hover{filter:brightness(1.1)}#chapter-phase{text-transform:uppercase;letter-spacing:.08em;color:var(--value);font-size:12px}#chapter-phase.bwd{color:var(--grad)}#chapter-title{font-size:13px;font-weight:600;line-height:1.35}#article .g{color:var(--grad)}#article .tag{margin-right:2px;font-weight:650}#chapter-title i{color:#5c6578;font-style:normal}#diagram{border-top:1px solid var(--line);margin-top:14px;padding-top:10px}#diagram svg{font-family:inherit;display:block;overflow:visible}#diagram .box{cursor:pointer}#diagram .box:hover rect.b{stroke-width:1.8px}#diagram rect.b{stroke-width:1px;transition:stroke-width .15s}#diagram rect.b.on{stroke:#fff;stroke-width:2px;filter:drop-shadow(0 0 4px #ffffff59)}#diagram .t{font-size:11px;font-weight:650}#diagram .s{fill:#8a93a6;font-size:9.5px}#diagram .tg{fill:#8a93a6;text-anchor:middle;cursor:pointer;font-size:13px}#diagram .tg:hover{fill:#fff}#diagram .arrow{stroke:#7b8496;stroke-width:1.2px;fill:none}#diagram .back{stroke:#a855f7;stroke-width:1.4px;fill:none;stroke-dasharray:4 3}#diagram .back-t{fill:#a855f7;text-anchor:middle;font-size:10px}#diagram .io{fill:#1a1e28;stroke:#3a4256}#diagram .io-t{fill:#c3c9d5;text-anchor:middle;font-size:11px}#model-card{color:var(--muted);border-top:1px solid var(--line);margin:0 20px 20px;padding-top:10px;font-size:12px}#model-card code{color:var(--text)}#net{min-width:0;min-height:0;position:relative;overflow:hidden}#net canvas{display:block}#net .labels{pointer-events:none;position:absolute;inset:0}.blabel{color:#c9cfdb;white-space:nowrap;text-align:right;padding-right:4px;font-size:11px;line-height:1.2}.blabel b{font-weight:600;display:block}.blabel span{color:#6f788b;font-size:10px}.blabel.head{color:var(--attn)}.chapter-label{color:#e6e9ef;white-space:nowrap;opacity:.85;font-size:22px;font-weight:700}#legend{color:var(--muted);background:#0c0e13bf;border-radius:6px;gap:14px;padding:4px 10px;font-size:12px;display:flex;position:absolute;bottom:10px;left:12px}.sw{vertical-align:-1px;border-radius:2px;width:22px;height:10px;margin-right:5px;display:inline-block}.sw.value{background:linear-gradient(90deg, var(--value-neg), #2b303b, var(--value))}.sw.grad{background:linear-gradient(90deg, var(--grad), #2b303b, var(--grad-pos))}.sw.weight{background:linear-gradient(90deg, #5a6bd9, #2b303b, var(--weight))}.sw.attn{background:linear-gradient(90deg, #2b303b, var(--attn))}#net-buttons{gap:6px;display:flex;position:absolute;top:10px;right:12px}#tooltip{z-index:5;pointer-events:none;background:#12151cf2;border:1px solid #343b4c;border-radius:6px;max-width:340px;padding:6px 9px;font-size:12px;line-height:1.45;position:absolute}#tooltip .k{color:var(--muted)}#tooltip .v{color:var(--value);font-variant-numeric:tabular-nums}#tooltip .gv{color:var(--grad);font-variant-numeric:tabular-nums}#tooltip table.stats{border-collapse:collapse;margin-top:4px}#tooltip table.stats td{text-align:right;padding:1px 0 1px 10px}#tooltip table.stats td:first-child{text-align:left;padding-left:0}body.zen #tooltip{-webkit-backdrop-filter:blur(3px);backdrop-filter:blur(3px);background:#0c0e13d1;max-width:420px}#right{pointer-events:none;z-index:3;flex-direction:column;width:clamp(200px,25vw,340px);display:flex;position:absolute;top:46px;right:10px}#right>*{pointer-events:auto}#right-text{text-shadow:0 1px 3px #000,0 0 8px #000;pointer-events:none!important}#molecule{aspect-ratio:1;background:0 0;flex:none;width:100%;position:relative;overflow:hidden}.mol-canvas{touch-action:none;width:100%;height:100%;display:block}.mol-ui{color:var(--muted);flex-wrap:wrap;align-items:center;gap:6px;font-size:11.5px;display:flex;position:absolute;top:6px;left:6px}.mol-ui select{background:#1a1e28d9;padding:2px 4px;font-size:11.5px}.mol-ui label{cursor:pointer;background:#10131ab3;border-radius:4px;align-items:center;gap:3px;padding:1px 5px;display:flex}.mol-legend{color:var(--muted);pointer-events:none;font-size:10.5px;position:absolute;bottom:6px;left:8px;right:8px}.mol-legend canvas{border-radius:3px;width:100%;height:6px;display:block}.mol-legend div{font-variant-numeric:tabular-nums;justify-content:space-between;gap:6px;margin-top:2px;display:flex}.mol-legend b{color:#c3c9d5;text-align:center;text-overflow:ellipsis;white-space:nowrap;font-weight:500;overflow:hidden}body.weights-only #legend>span:not(:nth-child(3)){display:none}#right-text{flex-direction:column;min-width:0;display:flex}#graph-caption{color:#c8ced9;padding:4px 6px;font-size:12px}#graph-caption b{color:var(--text)}#energies{color:var(--muted);font-variant-numeric:tabular-nums;padding:2px 6px 6px;font-size:12px}#energies b{color:var(--text)}#energies .fc b{color:#a78bfa}#energies .fnc b{color:#a3e635}footer{border-top:1px solid var(--line);background:var(--panel);align-items:center;gap:8px;padding:8px 14px;display:flex}footer button{min-width:38px}.speed{color:var(--muted);align-items:center;gap:6px;font-size:13px;display:flex}#speed-val{font-variant-numeric:tabular-nums;width:64px}#scrub{flex:1;height:34px;position:relative}#timeline{cursor:pointer;width:100%;height:100%;display:block}#scrub-label{color:var(--muted);font-size:11px;position:absolute;top:-18px;right:6px}@media (width<=900px){html,body{height:auto;min-height:100%}body{font-size:13.5px;display:block;overflow:auto}header{z-index:10;gap:6px 8px;padding:6px 10px;position:sticky;top:0}header label{font-size:12px}header select{max-width:46vw}.brand{margin-right:2px;font-size:16px}#status{text-align:left;width:100%;margin-left:0;font-size:12px}main{flex-direction:column;display:flex}#net{flex:none;height:58vh;min-height:280px}#right{width:44vw;top:40px;right:6px}#right-text{display:none}#walk{border:0;border-top:1px solid var(--line);order:3;overflow:visible}#legend{flex-wrap:wrap;gap:8px;font-size:11px;right:12px}#net-buttons{top:6px;right:6px}footer{z-index:10;flex-wrap:wrap;gap:6px;padding:6px 10px;position:sticky;bottom:0}footer button{min-width:34px;padding:4px 8px}.speed{font-size:12px}#speed{width:80px}#scrub{flex-basis:100%;height:30px}#tooltip{max-width:70vw}}@media (width<=520px){header .file:has(#model-files){display:none}header label:not(.file){gap:0;font-size:0}header .file{padding:3px 8px;font-size:12px}header label select{max-width:60vw;padding:3px 6px;font-size:12px}#net{height:52vh}#graph-caption,#energies{padding:6px 8px;font-size:11.5px}}#zen-count,#zen-progress{display:none}body.zen{grid-template-rows:1fr}body.zen header,body.zen footer,body.zen #walk,body.zen #right,body.zen #legend,body.zen #net-buttons>:not(#zen){display:none!important}body.zen main{grid-template-rows:1fr!important;grid-template-columns:1fr!important}body.zen #net-buttons{opacity:.25;transition:opacity .3s}body.zen #net-buttons:hover{opacity:1}body.zen #zen-count{pointer-events:none;letter-spacing:.02em;color:#e6e9ef;font-variant-numeric:tabular-nums;flex-direction:column;font-size:44px;font-weight:200;line-height:1;display:flex;position:absolute;bottom:22px;left:28px}body.zen #zen-count small{letter-spacing:.3em;text-transform:uppercase;color:#6f788b;margin-top:8px;font-size:12px;font-weight:400}body.zen #zen-progress{background:#161a22;height:3px;display:block;position:absolute;bottom:0;left:0;right:0}body.zen #zen-progress i{width:0;position:absolute;top:0;bottom:0;left:0}body.zen #zen-progress .f{background:var(--value)}body.zen #zen-progress .b{background:var(--grad)}

index.html

@@ -1,19 +1,74 @@
 <!doctype html>
-<html>
-	<head>
-		<meta charset="utf-8" />
-		<meta name="viewport" content="width=device-width" />
-		<title>My static Space</title>
-		<link rel="stylesheet" href="style.css" />
-	</head>
-	<body>
-		<div class="card">
-			<h1>Welcome to your static Space!</h1>
-			<p>You can modify this app directly by editing <i>index.html</i> in the Files and versions tab.</p>
-			<p>
-				Also don't forget to check the
-				<a href="https://huggingface.co/docs/hub/spaces" target="_blank">Spaces documentation</a>.
-			</p>
-		</div>
-	</body>
+<html lang="en">
+  <head>
+    <meta charset="utf-8" />
+    <meta name="viewport" content="width=device-width, initial-scale=1" />
+    <title>MLIP Visualization</title>
+    <meta name="description" content="Machine-learning interatomic potentials (PET, ANI-2x, KRR/SOAP, PhysNet) running in your browser, with every forward and backward pass visualised." />
+    <script type="module" crossorigin src="./assets/index-BrmxYqWl.js"></script>
+    <link rel="stylesheet" crossorigin href="./assets/index-BvLUcW0X.css">
+  </head>
+  <body>
+    <header>
+      <div class="brand">MLIP<span>viz</span></div>
+      <label title="The kind of model">Model <select id="kind"></select></label>
+      <label title="Which trained weights (for KRR: which model it is fitted to)"><span id="variant-label">Weights</span> <select id="model"></select></label>
+      <label class="file">Open model…<input id="model-files" type="file" accept=".json,.safetensors" multiple /></label>
+      <label>Structure <select id="structure"></select></label>
+      <label class="file">Open .xyz…<input id="xyz-file" type="file" accept=".xyz,.extxyz" /></label>
+ <label id="force-mode-wrap" title="How forces are computed">Forces <select id="force-mode">
+        <option value="conservative">conservative (−∂E/∂r)</option>
+        <option value="direct">direct (non-conservative)</option>
+        <option value="both">both, compared</option></select></label>
+      <button id="rattle" title="Displace every atom a little and run a new pass">Rattle</button>
+      <div id="status"></div>
+    </header>
+    <main>
+      <aside id="walk">
+        <div id="chapter-head"><span id="chapter-phase"></span><div id="chapter-title">Loading…</div></div>
+        <div id="intro">
+          <div id="intro-text"></div>
+          <figure><div id="diagram"></div><figcaption>The architecture of the loaded model. Click a box to jump there; ± collapses it in 3D.</figcaption></figure>
+          <p class="keys"><kbd>Space</kbd> continue · <kbd>←</kbd> <kbd>→</kbd> step · drag to pan, scroll to zoom, right-drag to rotate</p>
+        </div>
+        <div id="article"></div>
+        <div id="model-card"></div>
+      </aside>
+      <section id="net">
+        <div id="tooltip" hidden></div>
+        <div id="zen-count"><span id="zen-params"></span><small>parameters</small></div>
+        <div id="zen-progress"><i class="f"></i><i class="b"></i></div>
+        <div id="legend">
+          <span><i class="sw value"></i>activation</span>
+          <span><i class="sw grad"></i>gradient ∂E/∂·</span>
+          <span><i class="sw weight"></i><span id="legend-weights">weights</span></span>
+          <span id="legend-attn"><i class="sw attn"></i>attention</span>
+        </div>
+        <div id="net-buttons">
+          <button id="follow" class="on" title="Camera follows the active op">Follow</button>
+          <button id="overview" title="Show the whole network">Overview</button>
+ <button id="weights-only" title="Show only the weights of the model">Weights</button>
+          <button id="zen" title="Just the numbers, forward and backward, on a loop (Esc to leave)">Zen</button>
+        </div>
+          <div id="right">
+            <div id="molecule"></div>
+            <div id="right-text">
+              <div id="graph-caption"></div>
+              <div id="energies"></div>
+            </div>
+          </div>
+      </section>
+
+    </main>
+    <footer>
+      <button id="restart" title="Back to the start (Home)">⏮</button>
+      <button id="prev" title="Previous op (←)">◀</button>
+      <button id="play" title="Play the whole pass">▶</button>
+      <button id="continue" class="primary" title="Play to the end of this step (Space)">Continue</button>
+      <button id="next" title="Next op (→)">▶|</button>
+      <button id="end" title="Show the whole pass (End)">⏭</button>
+      <label class="speed">Speed <input id="speed" type="range" min="0" max="1" step="0.01" value="0.35" /><span id="speed-val"></span></label>
+      <div id="scrub"><canvas id="timeline"></canvas><div id="scrub-label"></div></div>
+    </footer>
+  </body>
 </html>