Skip to content

Instantly share code, notes, and snippets.

@apivovarov

apivovarov/jax_exp_accuracy.py

Last active January 29, 2025 02:44
Show Gist options
  • Save apivovarov/8a623503812a3aebf45b7ab7a8c74321 to your computer and use it in GitHub Desktop.
Save apivovarov/8a623503812a3aebf45b7ab7a8c74321 to your computer and use it in GitHub Desktop.
Download ZIP
jax exp accuracy
Raw
jax_exp_accuracy.py
# EXP
import jax
import jax.numpy as jnp
import numpy as np
print("JAX devices:", jax.devices())
a_np=np.array([[ 4603795724020441080, 4600877961149805680, 13824041971732116924, 13819622319934898176, 4591988256500146912, 13827780906010235892, 4592430392547641040, 13825287425120520724, 13828493035404110830, 4606660095605074408, 4602255938349038484, 4606604461173478372, 13827339776963088236, 4604220211627411202, 13824977737624476168, 4600871068755542572, 4606725295095174736, 4590846528761481504, 13822202805976138512, 13815414477567454288, 4606044597329528470, 4600076899447978620, 4591604889013648656, 4602697894782769988, 13827420058876989058, 4604756101012177124, 4605317067340288710, 13816787606007860464, 4602913948182442482, 4606332577827281932, 4604773657456759696, 13827381634999381006], [13826069199251957758, 4605825876037599016, 4599813349846477404, 4606034169064504034, 13810066424407939872, 4599928033056128564, 4604540450308053286, 4604372500785701914, 4596525362274821144, 13829135723316964212, 13793157593565504512, 13830164911641200584, 4602356688433217876, 13830210289030481844, 13825308288440654400, 4601870525974233060, 13830003905246268438, 4602330803283680356, 4600447725252968480, 4605100252105226294, 13828963658025299726, 13827343825652027438, 4603134915437094254, 13814046872961454352, 13810042193260139552, 4604167307366376072, 4606033062320603078, 4605597141026472258, 4597880993938036656, 4606067446121546990, 4589983112770652368, 4600405314893865540], [ 4606069237487458980, 4599235680658581428, 4599757359845027396, 13817863721271044472, 4603603891133792882, 4606791741599233806, 13822936828116972048, 13822072658456387652, 4595819460881724152, 13828708105407324654, 13813837329113632608, 13829768141884429998, 13814422052130623952, 13827666734994999012, 4606802658294129364, 13826236629650219638, 4598021171576849472, 13827256533478973132, 13830183862603772116, 4601945077705622892, 4602294726917956076, 13815156423273663184, 4607071637750782598, 4587511187726335392, 4598468157934128900, 13819451793331020968, 13827718925525307736, 4604088216257528092, 13828227428997654058, 4603483051428073770, 4606096207024034564, 4600980610063051148], [ 4601230477731166404, 4604008352368757090, 4604390030445437916, 4594762325462626056, 4605038345926119878, 13828469803655542236, 4600025567291157000, 4600276942943729836, 13802345749547869952, 4595774937787746912, 13822874157556610396, 13813727490755034080, 13829694714460857804, 13824013472505861416, 13813833100957238560, 4592370918022224592, 4576128359434768896, 4584739952264901312, 13827050192719802336, 13823402776330588056, 4594485711061461200, 4601430502824236388, 13820400350963075816, 4600543338738204420, 4600944155842297916, 13810628166570744160, 4602338930620146136, 13827178186363813712, 4606492425347911640, 4604508454304196248, 13818018033364639888, 13818847237051969304], [ 4598775694471616500, 13824050006322250272, 4601097954115765536, 4606870729910027376, 13829713443471052428, 4601456116703906472, 4606894991607668196, 13824386931545514196, 4602442342223123452, 4602816957507909990, 13823169394811731084, 4606529927856107862, 13819994423349447624, 4595106427327398064, 4605448450766388494, 4606725485912082418, 4605364022706235704, 4599958951757948032, 13825318172286297784, 4594319557628457984, 4582754203195657920, 13829275908975294170, 4606545196521135896, 4604845744874278768, 4605155961798503560, 4603131156721214224, 4604577983025610870, 4602184342661448320, 13827132672501306316, 13800517232603109248, 4600948485287822840, 4604439060058118208], [13814998160532354704, 13826369050796962918, 13830314111043305598, 4600703112645622820, 4602552646999566028, 13829902010764224818, 4597190514496879824, 13829654099519760780, 4596638846048137288, 13828011317286425114, 13825881464508500500, 4593912419592410184, 4604685199422644060, 4601547610506334120, 4598743216769112776, 4564198280070704128, 4606619906716367384, 13827122551284837260, 13830344064600032924, 13827223491115217294, 4598184474246303684, 4597467533865064760, 4604837746843504850, 4606045460718303606, 4599677036395660192, 4603797961413859530, 13818086713933828352, 4602748669717744762, 4605387673361392792, 4594608421134352072, 13827156874512558628, 13828334314505626982], [ 4598648981497804284, 4599380246629421884, 4595093503582194952, 4605554454121431858, 4601708514465329436, 4586130705501702848, 13824375487321887344, 4572187693808658688, 13817653932520747824, 4601033266723201652, 4604309157893069358, 4602105563009406616, 4604014371286715578, 4595187971352460736, 13828096405065023618, 13827611085967992938, 13826558993514498426, 4602277681696862444, 4604669938793348108, 4606514461478761660, 4604654538585365360, 4596939700715512544, 4596300572651445216, 13823511910032977608, 13815046396610859600, 13821402551356072328, 13826416391207812932, 13825994791218906832, 4604970694731977738, 13827448379653660376, 13827512256625724538, 13821933217534114420], [ 4595410578515141864, 4604484924865375132, 4605938825390536998, 4605736246004472636, 13827664593540227612, 4584498938805509312, 13826691435908409326, 4597298364381022496, 4603581386963987742, 4605441449232694538, 13824977595640565532, 13827176842674474242, 13822735473129753188, 13826376549463804784, 13826219908730269944, 4607030675595392936, 13823051575749199936, 13826603986878959814, 4599726254202165888, 4595322716191354800, 4598918353554796780, 13822927425404816932, 4606936946412423078, 13815922227547223616, 13829068890972578258, 4600675571510228616, 13822480940617032764, 13821093970725301768, 13826371092010142984, 13828836993993604376, 13824154349596740468, 4605073883519762422], [ 4597790853067374624, 13827128658574098568, 4607058728554971688, 13822872060822718352, 13817364403348219784, 13822242074772197056, 4607011722896010804, 13817973325833960688, 4601071460549787140, 13829913306165079114, 4605696340988493788, 4606675598028598930, 4605311386035035380, 4600121167488162116, 4595799625794445488, 4601807272245421896, 4601521498589815612, 4589956919156483616, 4603949613415033678, 13828114330944066828, 4593820240114039016, 13814464445640417344, 13827313918532335630, 4601595749379474576, 4590465879142131920, 4602655078289600772, 13826764025306954942, 4601744120493935408, 4604821024518012658, 13822787649143835416, 13830066551825615060, 4603134225411046000], [13829797479618874640, 4599080764521696448, 13813286416073869952, 4600634754215217124, 13826441350000159612, 4599713432858968312, 13825539358513717464, 13830509724349521222, 13826296784698656592, 4605746896207520950, 4593456347733033472, 4586082880962772256, 4593114126009835248, 4588813487871728000, 4606366944753840954, 4601010705935644276, 4602826708006143978, 13826101506407483360, 13829548723434399350, 13827699649427818020, 13818078400565509376, 4562227919181106176, 4589782307663625344, 4595863911324012008, 13823185472853447852, 13816252718612168768, 13819259727733384640, 4606412799250035876, 13822341271517232008, 13829124853293015326, 4603396954619615706, 4605827450748928406], [13824361604490940300, 4606509502907677974, 4597033114796508904, 4606138868563719022, 4602827211852458298, 4598537667496368464, 13822833028664509600, 13829786749304216120, 4597928693822168888, 4606786548333833568, 4606370230766918186, 13830173673279769966, 4606045024564455004, 13830087137371444222, 4596435441960200120, 13822875513618111152, 4600966867964770564, 4598654146326483444, 13820358755909241176, 13828763671765082738, 13812276515785431744, 13827511094232466068, 4606279892026114656, 13827918464394599764, 4598139723930698880, 13820705273680128000, 4606794920415080750, 13817523721503884240, 13821625319113881376, 4606160199796480608, 13828370059464226006, 13827965949281133140], [13827350051511694096, 13826409052498138118, 4595886904175347024, 13828649610691833950, 4602595895942315468, 13828729478807338970, 4578754876964353792, 13826408137069078552, 13826960077934259882, 4602455525848531940, 13826732134355433600, 4606385241263091898, 4596096599213088624, 4604726518555485716, 4602829540611371830, 13827205335565849856, 4601850122417718100, 13805041226329658880, 13828763632347793646, 4605442870520280624, 13828536391761766878, 4598384386768935080, 13824447645049952592, 13818685464119265632, 4606434119283409290, 13829953547990568290, 4607007744174527636, 13828994606884507034, 4590035730983730784, 4604748192581170310, 13829698192112334364, 13826966426978918854], [ 4597375778465647552, 13825226195371127492, 4605779970772346182, 13803705964890846144, 13817460253145661840, 4606030411012317350, 4598191400502193048, 4602018284280306488, 4606272961749682218, 13825765149177151512, 4606220808118095980, 13821143873657994400, 4587200854127043968, 4599827819400569624, 13827316945383071974, 13828010038123056558, 4605345816419273812, 13825603163040171524, 4605177138683912794, 13817400500731381912, 4600206869326129092, 13823076823279106516, 4599539237559954796, 13826330705296270856, 13823704793283742948, 4598525298883554028, 4601378803759234660, 13829276101192054906, 13825498679416813352, 4601960586883307312, 4606341279270291834, 13827259727672490642], [13823390342457362676, 13822705598117802584, 4594977933325672512, 4603582720686752828, 13823353302791900384, 13820636021176768648, 13830422508140317718, 4605413382184390074, 13829300785588869364, 13828953307191198974, 13828734165124638914, 13818032488230618384, 4605253256072618704, 13822875086716164988, 4585239050737005984, 4586619146740168160, 4582140748479122432, 13829851538403537656, 4605484914595179424, 13827634630251743924, 13826390908235684562, 4601799935361737476, 4604856091846312174, 4602189154335101328, 13829571084223794386, 4606215096455777348, 4605442040061978074, 4603014548484753084, 13822061359683202288, 13824785390130814292, 13829667667249188898, 13827929186386009414], [13811301031146389024, 4592727667324059968, 13829981587211486000, 13827339239218798088, 13824433021354464900, 4584872547342638912, 13829637474082733366, 13825196104284875504, 4602606480305496056, 13817301051380356464, 4604152381758005362, 4600575325971102488, 4602855494905722306, 4606136016115141458, 13829343974101818966, 4605829495941978576, 4604363265371646854, 4596483975940276360, 13821881246532053912, 4585160257677603552, 13821231856241683152, 4604568679713834410, 13815523848754968864, 4603585114039418346, 13822342801526128000, 13802122475453229184, 4601810027080836288, 4585282030329371744, 13828968520427722448, 4603040137713849970, 13810288388946617056, 13828203221509754130], [13826363579981018372, 4599007150830121312, 13828488439636147528, 4598450122335897520, 4603386552613546948, 4602581073801930924, 13828633524431857392, 13829271959491709840, 4606611872538316172, 13829245005333174886, 13828118879371842362, 13829399398989415226, 4604678590118518216, 4603796570712166848, 13826198046339709484, 13828121938807854156, 4598343153127498464, 13830229956820322812, 13808709816927633472, 4598391814008576996, 4606465271393917880, 4603974617152073300, 4603261580571233358, 13827247368746533956, 4603688386898021500, 13825560942976837248, 13824313959171209636, 4600614879020851176, 4603641975005399798, 4600218650761528140, 4599784952887303568, 13818405975641911024], [ 4605612239042285316, 4600039721461448708, 13825346410079006056, 4596917915412914376, 4600112659146771848, 13828206399875626464, 13790123117307047936, 4597732116509437536, 4604512708358410164, 4598462945067138140, 4589651540957683088, 4594445788861244400, 13830342876983831032, 4597252659099917032, 13827235370151328606, 4602342769476007608, 13827684316698735846, 13829066764832076736, 13814834852646376704, 13826201814106590694, 4593956157755261896, 4592753887595507600, 4601827033326654744, 4603451549915856162, 4603162779400997504, 4590881538601268944, 13817218819310326976, 4607088924855081342, 13821685726313436848, 13807868847633436096, 4600695555874173836, 4593937702605285680], [ 4573891753914012416, 4605938249261469410, 13821888954293202196, 4606885437905660284, 4594061936328899760, 13827279946319086098, 13821325577479382992, 4604681763033365924, 4603677365304872344, 4601356414184823860, 4604655997868769522, 4604825088955762884, 13810983562238053632, 4605815137594718404, 4604171496536443660, 4604046012561580130, 13829621479666380910, 4606419632062216574, 13820827646844046808, 13828107019930370228, 13826998646985346772, 13820801729634469632, 13821571944961727444, 4606464266826868266, 4605736974265070438, 13818003308715760720, 13826905699105857216, 13822135746536733648, 13827422529626017080, 13825812246970389096, 13823048884627647260, 13824595666540459204], [13821284955441578904, 4605789312186202342, 4603087713246031628, 13827714855245235504, 4604934950390606334, 4605700556889974596, 13826454778298482232, 4602442398752805572, 4604670794157451636, 13822231718428262892, 4604575364577443986, 4580679336982058112, 4606292235448492280, 4596343906754331872, 13808795225070352544, 4605692118558183396, 13827904420948239420, 4605736019079089052, 4594430269974191088, 13829834441561563308, 13822112939801637872, 4604842353554126784, 4605403196103055852, 4598809139088583352, 13809543303875893216, 13827291977364931750, 4585836052239282176, 13826685498830966672, 13826442159504713722, 13821365975444423800, 4600828650787735136, 13829282489050525420], [ 4604388113456732480, 13813580378264085968, 4590480202189059072, 4590154399749471360, 4605715789962498800, 13826644119529976070, 4606203184857529246, 13830470050224051846, 4598743793805616336, 4603952938085372782, 13821856176809588816, 4590633348437182176, 13828103383689392620, 13829283213427913652, 4591883264566843136, 4604637362989104462, 13818871061693557480, 4600439987716015744, 13829803514246718678, 13828306632362244146, 4590368542766764048, 13828473258024943582, 4591880947653568800, 4605051255683289248, 4605331092864637464, 4604196797364628718, 13817060811065169472, 4602897540321436514, 13824271278019616924, 4603077999984485944, 13827066362417887834, 13830030801454551330], [ 4592785118289075712, 13813362754511117712, 13805990854259005248, 4606135658069089112, 13826872054518257004, 4585255151109184928, 4606602926821398644, 4599359585788885256, 13827998984123649620, 13826914388851675012, 4593456934796315440, 4603738104647781228, 13821983242189830272, 13828888619646712942, 4603169750557961282, 4600727010315840376, 4603743073292582116, 13821612440878203640, 4591627632234258352, 13829707725531871996, 13829873123405356812, 4603990050058738654, 13829372619554149196, 4604746812525920850, 4587837723271550144, 4601767479307688716, 4601494269520866980, 4597908995631984064, 13826098253708538682, 4590214124395790944, 13827483823916400242, 13827206786863140384], [13823473488207390176, 13829947207917023874, 4604320747241337590, 4593313913509743136, 4598210203692464628, 4588162846156940288, 4598661731112058412, 4587869922036140640, 4600781803829224524, 4603671781467218004, 13815115148178145360, 13823109770103490448, 4599671062064447776, 4599295474894228004, 4604395625432390598, 4606132215601285672, 4602915654281961448, 13829408526531088834, 13827835111865501982, 13827212002131558418, 13823385470155218276, 13827184050425459310, 4605493478174995776, 4606991081533929884, 4606782517912439844, 4603721871923059516, 13827528038698615136, 4594925169376484632, 4584784442724067520, 4596111942045000072, 4587421782494661408, 4607172199958650916], [ 4606530566480325750, 4606642367241173290, 13823902562352857416, 4602865226128778540, 4600773664775532516, 13828929119815519002, 4603901224777558788, 13822079020747324184, 13799077880255699072, 13825074653812349808, 4602959270544561078, 13826896302162170078, 4604269702570047664, 13829960153255510218, 4577812294173671168, 13825023522688250644, 13821061250284651336, 4606947518195859006, 4602293568544429648, 4590851523353922352, 13827987666395123838, 13830242716809084044, 13822367718395149772, 4591435446504927520, 4603114133435597934, 4599275860743619272, 13827630417612341608, 13826761392166915100, 4579997621808854400, 4604598489609596116, 4598910656211662916, 13830029263340265832], [13823951995268147460, 4606090413700838244, 4599066885440916956, 4601323093615202364, 4605594129314913266, 13825716672759946020, 4605598725760284238, 13828941894591413366, 4605965985127719836, 4589169163878222016, 13827974502037598360, 13827528884031558548, 13822983819911232292, 13827384932967361752, 4596709222463756992, 13826782453164688736, 13822292630333655556, 13823839308471092696, 4603631663371585516, 13829008498812155382, 4604216256451623310, 13824494665945570792, 13829996249968295310, 4599016474099342712, 13830473917281861522, 4604876766141488270, 13826311418154106960, 13829839053274851836, 4604456545709788456, 4599842609709678320, 4598341885004143140, 13830040175371067988], [13826634642183815090, 13825673461761231760, 13824694095929871184, 4604319004540172654, 4605781912435491780, 13825931414691390732, 4604973237555882864, 13826349807543230354, 13827926927204115918, 4598388604214887940, 4595665709755694504, 13823904403504205172, 13829713344349134654, 13830230102914843048, 13828563207587320778, 13817937054773091480, 4596711004396909928, 4599712083518687428, 4597125307657486032, 13823229922977203488, 13819774860399680776, 13822339761169602308, 4600424975778042188, 13827346129199905294, 13829017570534387650, 13827791376333847652, 4605912959577794518, 13826902042713787904, 13828600511705889012, 13823570186887684676, 4603914345067716778, 4604516575244318188], [ 4606465189716778014, 4605219436693836440, 4587287666941726336, 13825758125205510324, 4597673550400602240, 13822187882108253380, 13818300549807974952, 4598185195991147216, 13828832150512763938, 4584483828076051648, 4598216848371118272, 4603302910276122738, 4606948527068982034, 13828208319050544062, 13825916894451042408, 4604831516743965462, 13818835352935327504, 4603705574385747628, 4603639001271468548, 4606859836920329974, 13828143875049351520, 4605664315974489080, 4586091937526692704, 13818373294449695216, 13830097769252602140, 4602662435294468356, 13827657963312361744, 13818658529323138496, 4604787944551745406, 4602329255037708300, 13809180948048181920, 13829972952679735888], [ 4597306778128720896, 4605059434215367440, 4606576840464207008, 4606580577761967542, 4603096164299275930, 13824048749159250832, 4598677566266687044, 13827363506750922278, 13817790763175868056, 13822376167169058876, 13820667924533147360, 13827170520321529150, 4594983662329875408, 4592585228595054048, 4605073398940806008, 13828510527156728762, 4606098661634054786, 4598360988089532292, 4606892917602039660, 13828815464911728236, 13828808458143743794, 13805815289859451904, 13824050307816259744, 4575255827443260928, 4600059909612764712, 13829180802015019072, 13826944400708324200, 4602833783987394186, 13825680649512090692, 4604421030649657990, 13825680112725804436, 13811367111542485632], [ 4589184785480759888, 4598160581572324040, 4604904377282054552, 4605136571095672544, 4600664415246273476, 13827698296772047398, 13830175884881840966, 4597897077985872408, 13826366582388253844, 13829356173149566154, 4603892687792269484, 4589639283633715184, 13826784745608138984, 4603388824972869338, 13829464492996151142, 13823009214961614516, 13827934599603971302, 13825620360864290020, 4595201548339702984, 4589830888416825248, 13809137287945081376, 4586629608752727872, 4604691283118729840, 4600746991054966852, 4602479205347651768, 13827116619080336626, 4598847353612514356, 13829712920772045168, 13828358432047759564, 4599553344443316236, 4602214525037607064, 13816717826287506096], [13827128607360352030, 13828566285068914320, 4606063942917048822, 4589602994860620896, 4603066327379693470, 13823776690584227220, 13826477663670958032, 4603709333610444870, 13826725277132772272, 4604815930874755456, 13821851405924289620, 13829685705885881416, 4604752940056263848, 13796152185494760960, 13828052948681441430, 13821739558829997924, 13823976575410343740, 13826281635251123382, 13822415337891362572, 13829428847504462620, 13821203133759330296, 13828060781778267698, 4598444115440733060, 4594770842279076904, 4604096072333508660, 4602470431819067312, 13810229023042185728, 4603758373669374368, 4603302479351971492, 4603137362624875984, 13820843505190957680, 4604529522416517974], [13824704293832132456, 4592720084910122416, 4598471786184901812, 13819876150431221800, 4604315920394282966, 13818586597039796000, 13828017199891221460, 4584026992193217792, 13812032772899228672, 4604451579158471198, 4601041259460289964, 13828500195172425808, 4585529586678645344, 13829228068237402974, 13811720085791470688, 4605643980974233122, 4603474247917650682, 13826355839360426436, 13829812123963631818, 4602817359967559298, 13812675105710851392, 13824266341364617244, 4595727129509946736, 13826515390070723150, 13830247860579143076, 4600367769480537424, 4602067529828331492, 13829566786958136062, 13822996701408594820, 4604862410893211164, 4580225276029599552, 4581955997758299648], [ 4599720170269365092, 4599662254860764940, 4606022344514755732, 13820689110104882696, 13824369498717683980, 4605389073745381624, 4593814351775093424, 13821596713128288720, 4600804151026064352, 13829437010125199530, 4605421142407749348, 4604639235479373848, 13826389186826358352, 13826149950664365824, 13791685724643785216, 13824671606270886384, 4602024866919069300, 13828876150768707492, 4604711855354075408, 13822528610263594128, 4605686799563135118, 13829987348568804778, 13814688066147913232, 13827799898969498050, 4605579053700597788, 4601305339046909084, 4605268542757940886, 4603713025298812448, 13815188852885995072, 13830195326255187552, 13826141278360555510, 4599930864279491804], [ 4604366102265413836, 4599160439286445028, 4606088600042946502, 13802387233379528704, 13827586810284901092, 4603812300036649912, 4587047099129876192, 13821946982021374348, 13802850569746354176, 4600943074496299472, 13815509655977406208, 4603676816397965464, 13829296400899384558, 4597199418449700664, 4605747326993451250, 4600465065165879524, 13819216058782629208, 4604794293834232192, 4598736979390752400, 13816943217311715040, 13826706863000297334, 13827956869437138620, 4598345147039549292, 13827081606756224866, 4586525758048966944, 13826152458017154536, 4603735337166011110, 4603629543890339778, 4601270976909107572, 4604205181816499426, 13826817754951648212, 4597427886650109624]], dtype=np.uint64).view(np.float64)
c_np = np.exp(a_np)
for dtype in [jnp.float32, jnp.float16, jnp.bfloat16]:
a = jnp.array(a_np, dtype=dtype)
@jax.jit
def myfunc(a):
result = jnp.exp(a)
return result
y = myfunc(a)
y_np = np.asarray(y)
arr1 = c_np
arr2 = y_np
abs_diff = np.abs(arr1 - arr2)
max_abs_diff = np.max(abs_diff)
# Relative difference
epsilon = 1e-12 # To avoid division by zero
rel_diff = abs_diff / (np.abs(arr1) + epsilon)
max_rel_diff = np.max(rel_diff)
# Output results
print("dtype:", arr2.dtype)
print("Max Absolute Difference: %.3e" % max_abs_diff)
print("Max Relative Difference: %.3e" % max_rel_diff)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment