Package: jaxfi
jaxfi(JAX Friendly Interface) - JAX with a PyTorch-like interfaceFriendly Interface to JAX, that behaves similar to PyTorch while maintaining compatibility.jaxfi(JAX Friendly Interface) - JAX with a PyTorch-like interfaceWorking with CPU and GPUJAX modules are accessible directly🔪 The Sharp Bits 🔪NotesInstallationChangelogNews: Better, improved interface!from jaxfi import jaxmis all you need!Creates a JAX-like module that behaves very similarly to PyTorch, so>>>fromjaxfiimportjaxmjaxm.norm===torch.normjaxm.rand===torch.randjaxm.cat===torch.catjaxm.manual_seed===torch.manual_seedMake sure to import this module before anything that might importjax(e.g.,jaxopt).# DOfromjaxfiimportjaxmimportjaxopt# DON'T!!!importjaxoptfromjaxfiimportjaxmWorking with CPU and GPUJAX has automatic device placement in functions, so omit thedeviceargument
when creating arrays in functions, i.e., in functions, specify only the dtype.Placing arrays on GPU and CPU is easy, either specify device/dtype directly or
usejaxm.toto move the array to a specific device/dtype.>>>jaxm.rand(2,device="cuda")>>>jaxm.rand(2,device="gpu",dtype=jaxm.float64)>>>jaxm.rand(2,device="cpu")>>>jaxm.to(jaxm.zeros(2),"cuda")Arrays are created on the CPU by default, but that can be changed usingjaxm.set_default_dtype(jaxm.float32)jaxm.set_default_device("gpu")jaxm.get_default_device()jaxm.get_default_dtype()*Defaultdtyperefers to CPU default dtype, default GPU dtype is alwaysfloat32, butfloat64arrays can be created on the GPU by specifying thedtype explicitly or by usingjaxm.to.jaxmbehaves like numpy (jax.numpy). Some methods are
patched directly from jax.jaxm.grad===jax.gradjaxm.jacobian===jax.jacobianjaxm.hessian===jax.hessianjaxm.jit===jax.jitjaxm.vmap===jax.vmapJAX modules are accessible directlyFinally, jax-backed modules are available directly in jaxm>>>jaxm.jax>>>jaxm.numpy>>>jaxm.random>>>jaxm.scipy>>>jaxm.lax🔪 The Sharp Bits 🔪Random numbers are implemented using a global random key (which can also be
manually set using e.g.,jaxm.manual_seed(2023)). However, that means parallelized
routines will generate the same random numbers.# DON'T DO THISjaxm.jax.vmap(lambda_:jaxm.randn(10))(jaxm.arange(10))# every row of random numbers is the same!# DO THIS INSTEADn=10random_keys=jaxm.make_random_keys(n)jaxm.jax.vmap(lambdakey,idx:jaxm.randn(10,key=key))(random_keys,jaxm.arange(n))jit-ted functions will also return  the same random numbers every time# DON'T DO THISf=jaxm.jit(lambdax:x*jaxm.randn(3))f(1)# [-1.12918106, -2.04245763, -0.40538156]f(1)# [-1.12918106, -2.04245763, -0.40538156]f(1)# [-1.12918106, -2.04245763, -0.40538156]# DO THISf=jaxm.jit(lambdax,key=None:x*jaxm.randn(3,key=key))f(1)# [-1.12918106, -2.04245763, -0.40538156]f(1,jaxm.make_random_key())# [-2.58426713,  0.90726101,  2.1546499 ]# jaxm.make_random_keys(n) is also availableNotesI'm not affiliated withJAXorPyTorchin any way.Installation$pipinstalljaxfiThe package name recently change fromjfitojaxfi, PyPI hosts it asjaxfi.Alternatively, to install from source, issue$pipinstall.from the project root, or simply run$pipinstallgit+https://github.com/rdyro/jaxfi-JAXFriendlyInterface.gitChangelogversion 0.6.3jaxm.tonow also moves numpy, not just jax, arrays to a device and dtypeexperimentalauto_pmapfunction available, automatically assigning first
batch dimension to multiple devices, e.g., dividing 16 tasks into 6 CPUsversion 0.6.0official name change fromjfitojaxfiversion 0.5.0settled on the default numpy module copy behavioromitdevicewhen creating arrays in functions - this now works correctlyintroduced more tests