Saravanan666/pypi-org huggingface dataset:
Package: jaxlob
JAX-LOBJAX-LOB: A GPU-Accelerated limit order book simulator to unlock large scale reinforcement learning for tradingArchitecture of the package:gymnax_exchange: theGPUversion of rl_environmentjaxob: Jax limit order bookjaxen: Jax trading_environmentjaxrl: Jax training loopDependenciespip install jax[cuda]==0.4.11 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html jaxlib==0.4.11 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html distrax brax chex flax optax gymnax wandbInstallationTo install the latest version of Jaxlob, run:pipinstalljaxlobCitationIf you find this project useful, please cite:@misc{frey2023jaxlob,
      title={JAX-LOB: A GPU-Accelerated limit order book simulator to unlock large scale reinforcement learning for trading}, 
      author={Sascha Frey and Kang Li and Peer Nagy and Silvia Sapora and Chris Lu and Stefan Zohren and Jakob Foerster and Anisoara Calinescu},
      year={2023},
      eprint={2308.13289},
      archivePrefix={arXiv},
      primaryClass={q-fin.TR}
}ContributionAnyone who might contribute codes to this repository, please email me:kang@robots.ox.ac.uk
Package: jax-metal
jax-metalThe jax-metal package is a Metal GPU plugin to provide Metal acceleration on Mac platforms for JAX applications.InstallationThe following table tracksjax-metalversions and compatible versions ofjax,jaxlibandMacOS.jax-metalMacOSjaxlibjax0.0.5Sonoma 14.2+v0.4.20v0.4.200.0.4Sonoma 14.0+v0.4.11v0.4.110.0.3Ventura 13.4.1+, Sonoma 14.0 Betav0.4.10v0.4.11We recommend to install the binary package with venv or conda.python3 -m venv ~/jax-metal
source ~/jax-metal/bin/activate
python -m pip install -U pip
python -m pip install numpy wheel
python -m pip install jax-metalPlease refer tohttps://developer.apple.com/metal/jax/for the full setup and verification.
Package: jaxnuts
jaxnutsPython implementation of the No-U-Turn Sampler fromHoffman and Gelman(Algorithm 6) leveraging JAX.UsageImport librariesimportjaximportjax.numpyasjnpimportjax.randomasrandomfromjaxnuts.samplerimportNUTSFor low dimensional problems such as this simple example, force JAX to use the CPU (avoid GPU overhead)jax.config.update('jax_platform_name','cpu')Define a log-probability to sample fromdeflogprob(x):"""Standard normal"""return-.5*jnp.dot(x,x)Generate sampleskey=random.PRNGkey(0)sampler=NUTS(jnp.ones(2),logp=logprob,target_acceptance=.5,M_adapt=1000)key,samples,step_size=sampler.sample(1000,key)
Package: jaxonloader
jaxonloaderA dataloader, but for JAX.The idea of this package is to have a DataLoader similar to the PyTorch one. To ensure that you don't have to learn anything new to use this package, the same API is chosen here (PyTorch's API actually a very good).Unfortunately, this also means that this package doesnotfollow the functional programming paradigm, because neither does the PyTorch DataLoader API. While in that regard this DataLoader is notfunctionalper se, it still allows for reproducability since you provide a random key to shuffle the data (if you want to).At the moment, this package is not yet a 1:1 mapping from PyTorch's DataLoader, but one day, it will be! *holding up arm and clenching fist*InstallationInstall this package using pip like so:pip install jaxonloaderUsagePretty much exactly as you would use PyTorch's DataLoader. Create a dataset class by inheriting from thejaxonloaderdataset and implement the__len__and__getitem__functions. Then simply pass that to the DataLoader class as argument.On the other hand, you can also use some of the provided datasets, such as the MNIST dataset.importjaxfromjaxonloaderimportget_mnistfromjaxonloader.dataloaderimportDataLoaderkey=jax.random.PRNGKey(0)train,test=get_mnist()train_loader=DataLoader(train,batch_size=4,shuffle=False,drop_last=True,key=key,)x=next(iter(train_loader))print(x[0].shape)# (4, 784)print(x[1].shape)# (4,)Performing TransformationsAs of now, transformations are not supported :(But - since you can get a dataset from aDataFrame- you can first
transform your data and then pass it to thefrom_dataframefunction.It's not ideal, but it works for now.Other backendsOther backends are not supported and are not planned to be supported. There is already
a very good dataloader for PyTorch, and with all the support PyTorch has, it's not
needed to litter the world with yet another PyTorch dataloader. The same goes for TensorFlow as well.If you really need one, which supports all backends, check outjax-dataloaderwhich does pretty much the same thing as this package, but for all backends.Then why does this package exist?For one, I just like building things and don't really care if it's needed or not. Secondly,
I don't care about other backends (as they are already very well supported) and only want to
focus on JAX and I needed a lightweight, easy-to-handle package, which loads data in JAX.So if you're like me and just need a simple dataloader for JAX, this package is for you.
If you need a dataloader for all backends, check out the other package from the link above.
Package: jaxonnxruntime
JAX ONNX RuntimeJAX ONNX Runtime is a robust and user-friendly tool chain that enables the seamless execution of ONNX models using JAX as the backend.More specifically, this tool chain has the abilities:ONNX Model Conversion: Converts ONNX models into JAX format modules. Tested on popular large language models including GPT-2, BERT, and LLaMA.Hardware Acceleration: Enable the jit mode of the converted JAX modules, which accelerates execution on GPU and/or TPU.Compatibility with JAX ecosystem: E.g., export models by Orbax, and serve the saved models by Tensorflow Serving system.Get StartedWe follow most of the interface definitions byonnx.backendhere.Please check a brief example on model conversion and forward calling inexamples/imagenet/imagenet_main.py.Contributions and DiscussionsWe believe that collaboration is the key to building remarkable software, and we wholeheartedly welcome contributions from developers like you.
You can make a real impact and help shape the future of our project with contributions such asimplementing new operatorsand increasing support for more ML models.Our contributors will have a chance to earnGoogle Open Source Peer Bonus, so that your valuable contributions won't go unnoticed.
Your hard work will be rewarded both by the community and by Google.
Together, let's create an amazing library and foster a supportive environment for open-source enthusiasts.Thank you for taking the time to contribute! Please seethe contribution guidelines.LicenseThis project is licensed under theApache License.
Package: jaxopt
JAXoptInstallation|Documentation|Examples|Cite usHardware accelerated, batchable and differentiable optimizers inJAX.Hardware accelerated:our implementations run on GPU and TPU, in addition
to CPU.Batchable:multiple instances of the same optimization problem can be
automatically vectorized using JAX's vmap.Differentiable:optimization problem solutions can be differentiated with
respect to their inputs either implicitly or via autodiff of unrolled
algorithm iterations.InstallationTo install the latest release of JAXopt, use the following command:$pipinstalljaxoptTo install thedevelopmentversion, use the following command instead:$pipinstallgit+https://github.com/google/jaxoptAlternatively, it can be installed from sources with the following command:$pythonsetup.pyinstallCite usOur implicit differentiation framework is described in thispaper. To cite it:@article{jaxopt_implicit_diff,
  title={Efficient and Modular Implicit Differentiation},
  author={Blondel, Mathieu and Berthet, Quentin and Cuturi, Marco and Frostig, Roy 
    and Hoyer, Stephan and Llinares-L{\'o}pez, Felipe and Pedregosa, Fabian 
    and Vert, Jean-Philippe},
  journal={arXiv preprint arXiv:2105.15183},
  year={2021}
}DisclaimerJAXopt is an open source project maintained by a dedicated team in Google Research, but is not an official Google product.
Package: jaxparrow
jaxparrowjaxparrowimplements a novel approach based on a variational formulation to compute the inversion of the cyclogeostrophic balance.It leverages the power ofJAX, to efficiently solve the inversion as an optimization problem.
Given the Sea Surface Height (SSH) field of an ocean system,jaxparrowestimates the velocity field that best satisfies the cyclogeostrophic balance.InstallationThe package is Pip-installable:pipinstalljaxparrowHowever, users with access to GPUs or TPUs should first install JAX separately in order to fully benefit from its high-performance computing capacities.
SeeJAX instructions.By default,jaxparrowwill install a CPU-only version of JAX if no other version is already present in the Python environment.UsageAs a packageTwo functions are directly available fromjaxparrow:geostrophycomputes the geostrophic velocity field (returns two2darray) from:a SSH field (a2darray),the latitude and longitude at the T points (two2darray),an optional mask grid (one2darray).cyclogeostrophycomputes the cyclogeostrophic velocity field (returns two2darray) from:a SSH field (a2darray),the latitude and longitude at the T points (two2darray),an optional mask grid (one2darray).BecausejaxparrowusesC-gridsthe velocity fields are represented on two grids (U and V), and the SSH on one grid (T).In a Python script, assuming that the input grids have already been initialised / imported, it would resort to:fromjaxparrowimportcyclogeostrophy,geostrophyu_geos,v_geos=geostrophy(ssh_t=ssh,lat_t=lat,lon_t=lon,mask=mask)u_cyclo,v_cyclo=cyclogeostrophy(ssh_t=ssh,lat_t=lat,lon_t=lon,mask=mask)To vectorise the application of thegeostrophyandcyclogeostrophyfunctions across an added time dimension, one aims to utilizevmap.
However, this necessitates avoiding the use ofnp.ma.masked_array.
Hence, our functions accommodate maskarrayas parameter to effectively consider masked regions.By default, thecyclogeostrophyfunction relies on our variational method.
Itsmethodargument provides the ability to use an iterative method instead, either the one described byPenvenet al., or the one byIoannouet al..
Additional arguments also give a finer control over the three approaches hyperparameters.SeejaxparrowAPI documentationfor more details.Notebooksare available as step-by-step examples.As an executablejaxparrowis also available from the command line:jaxparrow--conf_pathconf.ymlThe YAML configuration fileconf.ymlinstruct where input netCDF files are locally stored, and how to retrieve variables and coordinates from them.
It also provides the path of the output netCDF file. Optionally, it can specify which cyclogeostrophic approach should be applied and its hyperparameters.An example configuration file detailing all the required and optional entries can be foundhere.
Package: jax-pep503
PEP 503 repository index forjax[cuda]Daily automatically rebuildsRepositoryhttps://jorenham.github.io/jax_pep503/Example: Addingjax[cuda]to aPoetryprojectAdd the repository as a secondary source to yourpyproject.toml:...[[tool.poetry.source]]name="PyPI"priority="primary"[[tool.poetry.source]]name="jorenham/jax_pep503"url="https://jorenham.github.io/jax_pep503/"priority="supplemental"[tool.poetry.dependencies]...Now you can addjax[cuda]using:poetryaddjax[cuda]See theexample projectfor more details.See alsoPEP 503 – Simple Repository APIgoogle/jaxgoogle/jax#5410 – Please provide PEP 503 compliant indices for CUDA versions of packagesjax-releases(not compatible with PEP 503)
Package: jaxprvis
JaxPr VisualizerThis project provides interactive visualizations of JAX computations.
Composite nodes can be expanded using left-click and contracted using right-click.
Space pauses the simulation.Usage Exampleimportjax.numpyasjnpfromjaximportmake_jaxprfromjax.coreimportClosedJaxprimportjaxprvisasjvisdefvisualize_loop_unroll_jaxpr(iter:int)->ClosedJaxpr:example_args=jnp.zeros((2,2))deff(x,y):z=xforiinrange(iter):z=(z+y)/xreturnzclosed_jaxpr:ClosedJaxpr=make_jaxpr(f)(example_args,example_args)jvis.visualize(closed_jaxpr)InstallationFor UsersInstall from Pypi:#Create and activate virtual environment (if not present yet)python-mvenv--upgrade-deps./venvsource./venv/bin/activate
pipinstall--upgradebuild# Install your JAX version#e.g. pip install --upgrade "jax[cpu]"pipinstall--upgradejaxprvisInstall from source:# Clonegitclonehttps://github.com/PSilvestre/jaxprviscdjaxprvis#Create and activate virtual environmentpython-mvenv--upgrade-deps./venvsource./venv/bin/activate
pipinstall--upgradebuild

python-mbuild
pipinstall-e.For DevelopersSimilar to the above, but instead install the dev requirements.
Run the following to install:# Clonegitclonehttps://github.com/PSilvestre/jaxprviscdCORAL#Create and activate virtual environmentpython-mvenv--upgrade-deps./venvsource./venv/bin/activate
pipinstall--upgradebuild

python-mbuild
pipinstall-e".[dev]"#Run the testspytest
Package: jaxqualin
A python package for extracting quasinormal modes from black-hole ringdown simulations.Key Features•Installation•Usage•Paper Results•Coming Soon•How to Cite•LicenseKey FeaturesFit a ringdown waveform with quasinormal modes (QNMs) of fixed or free frequenciesNonlinear least-squares fitting with automatic differentiation viaJaxFitAgnostic identification of QNMs within the waveformSaving and reusing results withpickleEasy visualization of resultsCall hyperfit polynomials to approximate QNM amplitudes in the ringdown of binary black hole (BBH) mergersInstallationpipinstalljaxqualinUsageBasic usage examples can be found under the Examples tab on thepackage homepage.NoteWe did not extensively test and do not recommend runningjaxqualinon a GPUPaper ResultsInteractive plots of the methods paper results can be found under the Results tab on thepackage homepage.Coming SoonFull APISupport for real (Schwarzshild) ringdown waveformsFitting for the mass and spin of the remnantHow to CitePlease cite the methods paper if you used our package to produce results in your publication.
Here is the BibTeX entry:@misc{cheung2023extracting,
      title={Extracting linear and nonlinear quasinormal
      modes from black hole merger simulations}, 
      author={Mark Ho-Yeuk Cheung and Emanuele Berti and 
      Vishal Baibhav and Roberto Cotesta},
      year={2023},
      eprint={2310.04489},
      archivePrefix={arXiv},
      primaryClass={gr-qc}
}LicenseMITGitHub@mhycheung
Package: jax-quant-finance
JAX-Quant-Finance: High Performance Quantative Finance on JAXIntroductionInspired byTF Quant Finance, we developed a quantative finance library based onJAX. This library provides high-performance components leveraging the hardware
acceleration, parallel scientific computing and automatic differentiation of JAX.We can:run financial workloads on CPU/GPU/TPU with XLA accelerationcalculate mathematical derivative of financial models, i.e. Greeksdistribute workloads on multiple devices and machinesexamplesdirectory contains several demonstrations of using the JAX-Quant-Finance.InstallJAX installationYou must first followJAX's installation guideto install JAX based on your device architecture (CPU/GPU/TPU).JAX-Quant-Financepip install jax-quant-finance --upgrade64-bit precisionTo enable 64-bit precision, set the respective JAX flagbeforeimportingjax_quant_finance(see the JAXguide), for example:fromjax.configimportconfigconfig.update("jax_enable_x64",True)
Package: jax-random-projections
JAX Random Projection TransformersUsing JAX to speed up sklearn's random projection transformersInstallationNote: Installation with pip will install the CPU-only version of JAXTo use a GPU followJAX's installation guidebefore installingjax-random_projections.pip install jax-random_projectionsUsagefromjax_random_projections.sparseimportSparseRandomProjectionJAXtransfomer=SparseRandomProjectionJAX()transfomer.fit_transform(X)For the API documentation, refer tosklearn's SparseRandomProjection documentation.
The only difference is thatjax-random_projectionscurrently only supportsxla.DeviceArrayand doesn't supportdense_output=Falseandyforfit()This library currently only includes theSparseRandomProjectionbut a future release will also includeGaussianRandomProjection.jax-random_projectionsalso includesSparseRandomProjectionJAXCachedwhich uses a lru cache (maxsize=5) to speed up repeated calls by caching the random matrix for data with the same input dimension.
Package: jax-resnet
JAX ResNet - Implementations and Checkpoints for ResNet VariantsA Flax (Linen) implementation of ResNet (He et al. 2015), Wide ResNet
(Zagoruyko & Komodakis 2016), ResNeXt (Xie et al. 2017), ResNet-D (He et al.
2020), and ResNeSt (Zhang et al. 2020). The code is modular so you can mix and
match the various stem, residual, and bottleneck implementations.InstallationYou can install this package from PyPI:pipinstalljax-resnetOr directly from GitHub:pipinstall--upgradegit+https://github.com/n2cholas/jax-resnet.gitUsageSee the bottom ofjax-resnet/resnet.pyfor the available aliases/options for
the ResNet variants (all models are inFlax)Pretrained checkpoints fromtorch.hubare available for the
following networks:ResNet [18, 34, 50, 101, 152]WideResNet [50, 101]ResNeXt [50, 101]ResNeSt [50-Fast, 50, 101, 200, 269]The models aretestedto have the same intermediate activations and outputs as thetorch.hubimplementations, except ResNeSt-50 Fast, whose activations don't match exactly
but the final accuracy does.A pretrained checkpoint for ResNetD-50 is available fromfast.ai.
The activations do not match exactly, but the final accuracy matches.importjax.numpyasjnpfromjax_resnetimportpretrained_resnestResNeSt50,variables=pretrained_resnest(50)model=ResNeSt50()out=model.apply(variables,jnp.ones((32,224,224,3)),# ImageNet sized inputs.mutable=False)# Ensure `batch_stats` aren't updated.You must install PyTorch yourself
(instructions) to use these
functions.Transfer LearningTo extract a subset of the model, you can useSequential(model.layers[start:end]).Theslice_variablesfunction (found in incommon.py)
allows you to extract the corresponding subset of the variables dict. Check out
that docstring for more information.Checkpoint AccuraciesThe top 1 and top 5 accuracies reported below are on the ImageNet2012
validation split.  The data was preprocessed as in the officialPyTorch
example.ModelSizeTop 1Top 5ResNet1869.75%89.06%3473.29%91.42%5076.13%92.86%10177.37%93.53%15278.30%94.04%Wide ResNet5078.48%94.08%10178.88%94.29%ResNeXt5077.60%93.70%10179.30%94.51%ResNet-D5077.57%93.85%The ResNeSt validation data was preprocessed as inzhang1989/ResNeSt.ModelSizeCrop SizeTop 1Top 5ResNeSt-Fast5022480.53%95.34%ResNeSt5022481.05%95.42%10125682.82%96.32%20032083.84%96.86%26941684.53%96.98%ReferencesDeep Residual Learning for Image Recognition. Kaiming He, Xiangyu Zhang,
Shaoqing Ren, Jian Sun.arXiv 2015.Wide Residual Networks. Sergey Zagoruyko, Nikos Komodakis.BMVC
2016Aggregated Residual Transformations for Deep Neural Networks. Saining Xie,
Ross Girshick, Piotr Dollár, Zhuowen Tu, Kaiming He.CVPR
2017.Bag of Tricks for Image Classification with Convolutional Neural Networks.
Tong He, Zhi Zhang, Hang Zhang, Zhongyue Zhang, Junyuan Xie, Mu Li.CVPR
2019.ResNeSt: Split-Attention Networks. Hang Zhang, Chongruo Wu, Zhongyue Zhang,
Yi Zhu, Zhi Zhang, Haibin Lin, Yue Sun, Tong He, Jonas Mueller, R. Manmatha,
Mu Li, Alexander Smola.arXiv 2020.
Package: jax-russell
jax_russelljax-rusellis a package that implements financial option formulas, and leverages Jax's autodifferentiation to support calculating "the greeks."Formulas are taken from Espen Haug'sThe Complete Guide to Option Pricing Formulas, unless otherwise noted.Pre-alpha, API unstable.Documentation:https://SeanEaster.github.io/jax_russellGitHub:https://github.com/SeanEaster/jax_russellPyPI:https://pypi.org/project/jax_russell/Free software: GPL-3.0-onlyFeaturesStandard tree methods, like Cox-Ross-Rubinstein and Rendleman Bartter, for American and European optionsGeneralized Black-Scholes-MertonFirst- and second-order risk measures ("the greeks") via auto-differentiationSupport for options on stocks, stocks with continuous dividend, futures and margined futures (see the section "Usage")PlannedMore comprehensive testing (greeks against first- and second-differences)Skewed and leptokurtic methodsCreditsThis package was created withCookiecutterand thewaynerv/cookiecutter-pypackageproject template.
Package: jax-scipy-spatial
jax-scipy-spatialThis package implementsscipy.spatialAPI for JAX.Currently the following items are implemented:scipy.spatial.distancescipy.spatial.transform.Rotationscipy.spatial.transform.SlerpInstallpip install .Usageimport jax_scipy_spatial.transform as jtr

rotation = jtr.Rotation.from_euler('xyz', jnp.array([0., 0., 180.]), degrees=True)
print(rotation)DocumentationPlease refer to scipy documentation.ContributingTo run unit tests on your local machine:tox
Package: jax-sgmc
Modular Stochastic Gradient MCMC for JaxIntroduction|Implemented Solvers|Features|Installation|ContributingIntroductionJaxSGMCbrings Stochastic Gradient Markov chain Monte Carlo (SGMCMC)
samplers to JAX. Inspired byoptax,JaxSGMCis built on a modular concept to increase reusability and
accelerate research of new SGMCMC solvers. Additionally,JaxSGMCaims to
promote probabilistic machine learning by removing obstacles in switching
from stochastic optimizers to SGMCMC samplers.Quickstart with solvers fromalias.pyTo get started quickly using SGMCMC samplers,JaxSGMCprovides some popular
pre-built samplers inalias.py:SGLD (rms-prop):https://arxiv.org/abs/1512.07666SGHMC:https://arxiv.org/abs/1402.4102reSGLD:https://arxiv.org/abs/2008.05367v3SGGMC:https://arxiv.org/abs/2102.01691AMAGOLD:https://arxiv.org/abs/2003.00193OBABO:https://arxiv.org/abs/2102.01691FeaturesModular SGMCMC solversJaxSGMCaims to increase reusability of SGMCMC components via a toolbox of
helper functions and a modular concept:In the simplest case of employing a pre-built sampler fromalias.py, the user only needs to provide the computational
model, consisting of functions for Prior and Likelihood.
Schedulers allow to change sampler properies over the course of the training.
Advanced users may build custom samplers from given components.Data Input / Output underjitJaxSGMCprovides a toolbox to pass reference data to the computation
and save collected samples from the Markov chain.By combining different data loader / collector classes and general wrappers it
is possible to read data from and save samples to different data types via the
mechanisms of JAX's Host-Callback module.
It is therefore also possible to access datasets bigger than the device memory.Saving Data:HDF5Numpy.npzLoading Data:HDF5Numpy arraysTensorflow datasetsComputing the stochastic potentialStochastic Gradient MCMC requires the evaluation of a potential function for a
batch of data.JaxSGMCallows to compute this potential from likelihoods accepting only
single observations and batches them automatically with sequential, parallel or
vectorized execution.
Moreover,JaxSGMCsupports passing a model state between the evaluations of
the likelihood function, which is saved corresponding to the samples, speeding
up postprocessing.InstallationBasic SetupJaxSGMCcan be installed via pip:pipinstalljax-sgmc--upgradeThe above command installsJax for CPU. To runJaxSGMC on the GPU,
the GPU version of JAX has to be installed.
Further information can be found here:Jax Installation InstructionsAdditional PackagesSome parts ofJaxSGMCrequire additional packages:Data Loading with tensorflow:pipinstalljax-sgmc[tensorflow]--upgradeSaving Samples in the HDF5-Format:pipinstalljax-sgmc[hdf5]--upgradeInstallation from SourceFor development purposes,JaxSGMCcan be installed from source in
editable mode:gitclonegit@github.com:tummfm/jax-sgmc.git
pipinstall-e.[test,docs]This command additionally installs the requirements to run the tests:pytesttestsAnd to build the documentation (e.g. in html):make-CdocshtmlContributingContributions are always welcome! Please open a pull request to discuss the code
additions.