Saravanan666/pypi-org huggingface dataset:
Package: jax-fem
A GPU-accelerated differentiable finite element analysis package based onJAX. Used to be part of the suite of open-source python packages for Additive Manufacturing (AM) research,JAX-AM.Finite Element Method (FEM)FEM is a powerful tool, where we support the following features2D quadrilateral/triangle elements3D hexahedron/tetrahedron elementsFirst and second order elementsDirichlet/Neumann/Robin boundary conditionsLinear and nonlinear analysis includingHeat equationLinear elasticityHyperelasticityPlasticity (macro and crystal plasticity)Differentiable simulation for solving inverse/design problemswithouthuman deriving sensitivities, e.g.,Topology optimizationOptimal thermal controlIntegration with PETSc for solver choicesUpdates(Dec 11, 2023):We now support multi-physics problems in the sense that multiple variables can be solved monolithically. For example, consider runningpython -m applications.stokes.exampleWeak form is now defined through  volume integral and surface integral. We can now treat body force, "mass kernel" and "Laplace kernel" in a unified way through volume integral, and treat "Neumann B.C." and "Robin B.C." in a unified way through surface integral.Thermal profile in direct energy deposition.Linear static analysis of a bracket.Crystal plasticity: grain structure (left) and stress-xx (right).Stokes flow: velocity (left) and pressure(right).Topology optimization with differentiable simulation.InstallationClone the repository:gitclonegit@github.com:tianjuxue/jax-fem.gitcdjax-femCreate a conda environment from the givenenvironment.ymlfile:condaenvcreate-fenvironment.ymlActivate the environment and install the package:condaactivatejax-fem-env
pipinstall-e.DependenciesInstall JAXSee jax installationinstructionsQuick startCheckdemos/for a variety of FEM cases. For example, runpython-mdemos.hyperelasticity.examplefor hyperelasticity.Also,python-mtests.benchmarkswill execute a set of test cases.LicenseThis project is licensed under the GNU General Public License v3 - see theLICENSEfor details.CitationsIf you found this library useful in academic or industry work, we appreciate your support if you consider 1) starring the project on Github, and 2) citing relevant papers:@article{xue2023jax,title={JAX-FEM: A differentiable GPU-accelerated 3D finite element solver for automatic inverse design and mechanistic data science},author={Xue, Tianju and Liao, Shuheng and Gan, Zhengtao and Park, Chanwook and Xie, Xiaoyu and Liu, Wing Kam and Cao, Jian},journal={Computer Physics Communications},pages={108802},year={2023},publisher={Elsevier}}
Package: jax-fid
FID computation in Jax/FlaxThis is a port ofmseitzer/pytorch-fid, which is a port of the original FID implementation (bioinf-jku/TTUR).The parameters for theInceptionV3network are taken frommseitzer/pytorch-fid. The FID scores are almost identical (absolute difference around 1e-7).The only difference is thatmseitzer/pytorch-fidresizesthe images to 299x299 by default. In this implementation, the images are not resized by default. You can resize the images using the--img_sizeargument.InstallationYou will need Python 3.7 or later.For GPU usage, follow theJaxinstallation with CUDA.Then install:>pipinstalljax-fidFor CPU-only you can skip step 1.UsageCompute FID score>CUDA_VISIBLE_DEVICES=Npython-mjax_fid--path1/path/to/dataset1--path2/path/to/dataset2whereNis the GPU index.Pre-compute statistics for image directory>CUDA_VISIBLE_DEVICES=Npython-mjax_fid--precompute--img_dir/path/to/dataset--out_dir/path/to/statsArguments--path1- Path to image directory or .npz file containing pre-computed statistics.--path2- Path to image directory or .npz file containing pre-computed statistics.--batch_size- Batch size per device for computing the Inception activations.--img_size- Resize images to this size. The format is (height, width).--precompute- If True, pre-compute statistics for given image directory.--img_dir- Path to image directory for pre-computing statistics.--out_dir- Path where pre-computed statistics are stored.LicenseApache-2.0 License
Package: jax-finufft
JAX bindings to FINUFFTThis package provides aJAXinterface to (a
subset of) theFlatiron Institute Non-uniform Fast Fourier Transform (FINUFFT)
library. Take a look at theFINUFFT docsfor all the necessary
definitions, conventions, and more information about the algorithms and their
implementation. This package uses a low-level interface to directly expose the
FINUFFT library to JAX's XLA backend, as well as implementing differentiation
rules for the transforms.Included featuresThis library is currently CPU-only, but GPU support is in the works using thecuFINUFFT library.Type 1 and 2transforms
are supported in 1-, 2-, and 3-dimensions. All of these functions support
forward, reverse, and higher-order differentiation, as well as batching usingvmap.InstallationFor now, only a source build is supported.For building, you should only need a recent version of Python (>3.6) andFFTW. At runtime, you'll neednumpy,scipy, andjax. To set up such an environment, you can useconda(but you're welcome to
use whatever workflow works for you!):condacreate-njax-finufft-cconda-forgepython=3.9numpyscipyfftw
python-mpipinstall"jax[cpu]"Then you can install from source using (don't forget the--recursiveflag
because FINUFFT is included as a submodule):gitclone--recursivehttps://github.com/dfm/jax-finufftcdjax-finufft
python-mpipinstall.UsageThis library provides two high-level functions (and these should be all that you
generally need to interact with):nufft1andnufft2(for the two "types" of
transforms). If you're already familiar with thePython
interfaceto FINUFFT,please note that the function signatures here are different!For example, here's how you can do a 1-dimensional type 1 transform:importnumpyasnpfromjax_finufftimportnufft1M=100000N=200000x=2*np.pi*np.random.uniform(size=M)c=np.random.standard_normal(size=M)+1j*np.random.standard_normal(size=M)f=nufft1(N,c,x,eps=1e-6,iflag=1)Noting that theepsandiflagare optional, and that (for good reason, I
promise!) the order of the positional arguments is reversed from thefinufftPython package.The syntax for a 2-, or 3-dimensional transform is:f=nufft1((Nx,Ny),c,x,y)# 2Df=nufft1((Nx,Ny,Nz),c,x,y,z)# 3DThe syntax for a type 2 transform is (also allowing optionaliflagandepsparameters):c=nufft2(f,x)# 1Dc=nufft2(f,x,y)# 2Dc=nufft2(f,x,y,z)# 3DSimilar librariesfinufft: The
"official" Python bindings to FINUFFT. A good choice if you're not already
using JAX and if you don't need to differentiate through your transform.mrphys/tensorflow-nufft:
TensorFlow bindings for FINUFFT and cuFINUFFT.License & attributionThis package, developed by Dan Foreman-Mackey is licensed under the Apache
License, Version 2.0, with the following copyright:Copyright 2021 The Simons Foundation, Inc.If you use this software, please cite the primary references listed on theFINUFFT docs.
Package: jax-fixedpoint
fax: fixed-point jaxImplicit andcompetitive differentiationinJAX.Our "competitive differentiation" approach usesCompetitive Gradient Descentto solve the equality-constrained nonlinear program associated with the fixed-point problem. A standalone implementation of CGD is provided underfax/competitive/cga.pyand the equality-constrained solver derived from it can be accessed viafax.constrained.cga_lagrange_minorfax.constrained.cga_ecp. An implementation of implicit differentiation based onChristianson'stwo-phases reverse accumulation algorithm can also be obtained with the functionfax.implicit.two_phase_solver.Seefax/constrained/constrained_test.pyfor examples. Please note that the API is subject to change.ReferencesCiting competitive differentiation:@inproceedings{bacon2019optrl,
  author={Pierre-Luc Bacon, Florian Schaefer, Clement Gehring, Animashree Anandkumar, Emma Brunskill},
  title={A Lagrangian Method for Inverse Problems in Reinforcement Learning},
  booktitle={NeurIPS Optimization Foundations for Reinforcement Learning Workshop},
  year={2019},
  url={http://lis.csail.mit.edu/pubs/bacon-optrl-2019.pdf},
  keywords={Optimization, Reinforcement Learning, Lagrangian}
}Citing this repo:@misc{gehring2019fax,
  author = {Clement Gehring, Pierre-Luc Bacon, Florian Schaefer},
  title = {{FAX: differentiating fixed point problems in JAX}},
  note = {Available at: https://github.com/gehring/fax},
  year = {2019}
}
Package: jax-fixedpoint-test-manueldelverme
fax: fixed-point jaxImplicit andcompetitive differentiationinJAX.Our "competitive differentiation" approach usesCompetitive Gradient Descentto solve the equality-constrained nonlinear program associated with the fixed-point problem. A standalone implementation of CGD is provided underfax/competitive/cga.pyand the equality-constrained solver derived from it can be accessed viafax.constrained.cga_lagrange_minorfax.constrained.cga_ecp. An implementation of implicit differentiation based onChristianson'stwo-phases reverse accumulation algorithm can also be obtained with the functionfax.implicit.two_phase_solver.Seefax/constrained/constrained_test.pyfor examples. Please note that the API is subject to change.ReferencesCiting competitive differentiation:@inproceedings{bacon2019optrl,
  author={Pierre-Luc Bacon, Florian Schaefer, Clement Gehring, Animashree Anandkumar, Emma Brunskill},
  title={A Lagrangian Method for Inverse Problems in Reinforcement Learning},
  booktitle={NeurIPS Optimization Foundations for Reinforcement Learning Workshop},
  year={2019},
  url={http://lis.csail.mit.edu/pubs/bacon-optrl-2019.pdf},
  keywords={Optimization, Reinforcement Learning, Lagrangian}
}Citing this repo:@misc{gehring2019fax,
  author = {Clement Gehring, Pierre-Luc Bacon, Florian Schaefer},
  title = {{FAX: differentiating fixed point problems in JAX}},
  note = {Available at: https://github.com/gehring/fax},
  year = {2019}
}
Package: jaxfss
jaxfssJAX/Flax implementation of finite-size scaling analysisInstallationjaxfsscan be installed with pip with the following command:pip install jaxfssQuickstartCheck out thedocumentation!!Other packagesFinite-size scaling package by Gaussian process with C++https://kenjiharada.github.io/BSA/Finite-size scaling package by neural network and Gaussian process with Python (PyTorch)https://github.com/KenjiHarada/FSS-toolsCitationPlease cite this paper when you use this package for your research!![Full paper] Ryosuke Yoneda and Kenji Harada, Neural Network Approach to Scaling Analysis of Critical Phenomena,arXiv: 2209.01777.@article{yoneda2022neural,
    title={Neural Network Approach to Scaling Analysis of Critical Phenomena},
    author={Yoneda, Ryosuke and Harada, Kenji},
    url={https://arxiv.org/abs/2209.01777},
    journal={arXiv preprint arXiv:2209.01777},
    year={2022}}[Conference paper] Currently preparing!!
Package: jax-geneweaver-core
Geneweaver CoreThe Geneweaver Core Python library provides shared foundational functionality for the Geneweaver project.
It is a dependency of all other Geneweaver Python libraries, and is not intended to be used directly.InstallationThe Geneweaver Core library is available on PyPI and can be installed with pip:pipinstallgeneweaver-coreIf you are using Poetry, you can add the Geneweaver Core library to your project with:poetryaddgeneweaver-coreOverviewThis package is structured so as to share an import path with othergeneweaverpackages. This allows you to install
the package and import it without having to worry about the package name. For example, if you install thegeneweaver-corepackage, as well as thegeneweaver-dbpackage, you can import thegeneweaverpackage and access both libraries:fromgeneweaverimportcore,dbThe Geneweaver Core library provides the following functionality:geneweaver.core.schema: Pydantic schema definitionsgeneweaver.core.enum: Enumerationsgeneweaver.core.config: Configuration managementGeneweaverBaseSettings: Base Settings class to inherit from for your own package's configuration classesGeneweaverCoreSettings: Configuration class for Geneweaver Core settingsgeneweaver.core.exc: Geneweaver shared exceptionsGeneweaverException: Base exception class for GeneweaverGeneweaverError: Base error class for GeneweaverGeneweaverWarning: Base warning class for GeneweaverPlanned Functionalitygeneweaver.core.logging: Shared logging managementgeneweaver.core.utils: Shared utility functionsAcknowledgementsThis project was developed by the Computational Systems and Synthetic Biology Center at the Jackson Laboratory in
conjunction with the Baylor University Computational Biology and Bioinformatics Program.
Package: jaxhelper
JaxHelperBasic tools and helpers for JaxGetting StartedInstallationpython3-mpipinstalljaxhelperExplanationThis repository contains basic helper functions I use every day.Here are some highlights:remat: function decorator to rematerialize ("activation checkpointing") hidden states during backward passsoftmax:exp in fp32 and matmul in bf16 (-> improved convergence and speed)fewer stored intermediates yet faster gradientattention:recomputation of hidden statesmemory of O(K * N) rather than O(N^2) thanks
toSelf-attention Does Not Need O(n2) Memory. (No slowdown)
Package: jax-jumpy
Jumpy is a common backend forNumPyand optionallyJAX:If Jax is installed and jax inputs are provided then thejax.numpyfunction is runIf Jax is installed and the function is jitted then thejax.numpyfunction is runOtherwise the jumpy function returns the NumPy outputsThere are several functions (e.g.vmap,scan) that are available withjaxinstalled.Jumpy lets you write framework-agnostic code that is easy to debug by running
as raw Numpy, but is just as performant as JAX when jitted.We maintain this repository primarily so to enable writing Gymnasium and PettingZoo wrappers that can be
applied to both standard NumPy or hardware accelerated Jax based environments, however this package can be used for many more things.Installing JumpyTo install Jumpy from pypi:pip install jax-jumpy[jax]will include jax whilepip install jax-jumpywill not include jax.Alternatively, to install Jumpy from source, clone this repo,cdto it, and then:pip install .ContributingJumpy does not have a complete implementation of allnumpyorjax.numpyfunctions.
If you are missing functions then please create an issue or pull request, we will be happy to add them.In the future, we are interested in adding optional support for PyTorch and looking for pull request to complete this.
Package: jaxkern-nightly
This project has now been incorporated intoGPJax.Kernels in Jax.IntroductionJaxKern is Python library for working with kernel functions in JAX. We currently support the following kernels:StationaryRadial basis function (Squared exponential)MatérnPowered exponentialRational quadraticWhite noisePeriodicNon-stationaryLinearPolynomialNon-EuclideanGraph kernelsIn addition to this, we implement kernel approximations using theRandom Fourier featureapproach.ExampleThe following code snippet demonstrates how the first order Matérn kernel can be computed and, subsequently, approximated using random Fourier features.importjaxkernasjkimportjax.numpyasjnpimportjax.randomasjrkey=jr.PRNGKey(123)# Define the points on which we'll evaluate the kernelX=jr.uniform(key,shape=(10,1),minval=-3.,maxval=3.)Y=jr.uniform(key,shape=(20,1),minval=-3.,maxval=3.)# Instantiate the kernel and its parameterskernel=jk.Matern32()params=kernel.init_params(key)# Compute the 10x10 Gram matrixKxx=kernel.gram(params,X)# Compute the 10x20 cross-covariance matrixKxy=kernel.cross_covariance(params,X,Y)# Build a RFF approximationapprox=RFF(kernel,num_basis_fns=5)rff_params=approx.init_params(key)# Build an approximation to the Gram matrixQff=approx.gram(rff_params,X)Code StructureAll kernels are supplied with agramandcross_covariancemethod. When computing a Gram matrix, there is often some structure in the data (e.g., Markov) that can be exploited to yield a sparse matrix. To instruct JAX how to operate on this, the return type ofgramis a Linear Operator fromJaxLinOp.WithinGPJax, all kernel computations are handled using JaxKern.DocumentationA full set of documentation is a work in progress. However, many of the details in JaxKern can be found in theGPJax documentation.
Package: jaxlib
jaxlib: support library for JAXjaxlib is the support library for JAX. While JAX itself is a pure Python package,
jaxlib contains the binary (C/C++) parts of the library, including Python bindings,
the XLA compiler, the PJRT runtime, and a handful of handwritten kernels.
For more information, including installation and build instructions, refer to main
JAX README:https://github.com/google/jax/.
Package: jaxlib-rocm
jaxlib: support library for JAXjaxlib is the support library for JAX. While JAX itself is a pure Python package,
jaxlib contains the binary (C/C++) parts of the library, including Python bindings,
the XLA compiler, the PJRT runtime, and a handful of handwritten kernels.
For more information, including installation and build instructions, refer to main
JAX README:https://github.com/google/jax/.
Package: jaxlie
jaxlie[API reference][PyPI]jaxlieis a library containing implementations of Lie groups commonly used for
rigid body transformations, targeted at computer vision & robotics
applications written in JAX. Heavily inspired by the C++ librarySophus.We implement Lie groups as high-level (data)classes:GroupDescriptionParameterizationjaxlie.SO2Rotations in 2D.(real, imaginary):unit complex (∈ S1)jaxlie.SE2Proper rigid transforms in 2D.(real, imaginary, x, y):unit complex & translationjaxlie.SO3Rotations in 3D.(qw, qx, qy, qz):wxyz quaternion (∈ S3)jaxlie.SE3Proper rigid transforms in 3D.(qw, qx, qy, qz, x, y, z):wxyz quaternion & translationWhere each group supports:Forward- and reverse-mode AD-friendlyexp(),log(),adjoint(),apply(),multiply(),inverse(),identity(),from_matrix(), andas_matrix()operations. (see./examples/se3_example.py)Helpers for optimization on manifolds (see./examples/se3_optimization.py,jaxlie.manifold.*).Compatibility with standard JAX function transformations. (see./examples/vmap_example.py)(Un)flattening as pytree nodes.Serialization usingflax.We also implement various common utilities for things like uniform random
sampling (sample_uniform()) and converting from/to Euler angles (in theSO3class).Install (Python >=3.7)# Python 3.6 releases also exist, but are no longer being updated.pipinstalljaxlieIn the wildjaxfgappliesjaxlieto nonlinear least
squares problems with block-sparse structure. (for pose graph optimization,
bundle adjustment, etc)tensorf-jaxis an unofficial
implementation ofTensorial Radiance Fields (Chen et al, ECCV 2022)usingjaxlie.Miscjaxliewas originally written for our IROS 2021 paper
(link). If it's useful for you, you're
welcome to cite:@inproceedings{yi2021iros,
    author={Brent Yi and Michelle Lee and Alina Kloss and Roberto Mart\'in-Mart\'in and Jeannette Bohg},
    title = {Differentiable Factor Graph Optimization for Learning Smoothers},
    year = 2021,
    BOOKTITLE = {2021 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)}
}
Package: jaxlinop
JaxLinOpJaxLinOpis a lightweight linear operator library written injax.OverviewConsider solving a diagonal matrix $A$ against a vector $b$.importjax.numpyasjnpn=1000diag=jnp.linspace(1.0,2.0,n)A=jnp.diag(diag)b=jnp.linspace(3.0,4.0,n)# A⁻¹ bjnp.solve(A,b)Doing so is costly in large problems. Storing the matrix gives rise to memory costs of $O(n^2)$, and inverting the matrix costs $O(n^3)$ in the number of data points $n$.But hold on a second. Notice:We only have to store the diagonal entries to determine the matrix $A$. Doing so, would reduce memory costs from $O(n^2)$ to $O(n)$.To invert $A$, we only need to take the reciprocal of the diagonal, reducing inversion costs from $O(n^3)$, to $O(n)$.JaxLinOpis designed to exploit stucture of this kind.importjaxlinopA=jaxlinop.DiagonalLinearOperator(diag=diag)# A⁻¹ bA.solve(b)JaxLinOpis designed to automatically reduce cost savings in matrix addition, multiplication, computing log-determinants and more, for other matrix stuctures too!Custom Linear Operator (details to come soon)The flexible design ofJaxLinOpwill allow users to impliment their own custom linear operators.fromjaxlinopimportLinearOperatorclassMyLinearOperator(LinearOperator):def__init__(self,...)...# There will be a minimal number methods that users need to impliment for their custom operator.# For optimal efficiency, we'll make it easy for the user to add optional methods to their operator,# if they give better performance than the defaults.InstallationStable versionThe latest stable version ofjaxlinopcan be installed viapip:pipinstalljaxlinopNoteWe recommend you check your installation version:python-c'import jaxlinop; print(jaxlinop.__version__)'Development versionWarningThis version is possibly unstable and may contain bugs.Clone a copy of the repository to your local machine and run the setup configuration in development mode.gitclonehttps://github.com/JaxGaussianProcesses/JaxLinOp.gitcdjaxlinop
python-msetupdevelopNoteWe advise you create virtual environment before installing:conda create -n jaxlinop_ex python=3.10.0
conda activate jaxlinop_exand recommend you check your installation passes the supplied unit tests:python-mpytesttests/
Package: jaxlinop-nightly
JaxLinOpJaxLinOpis a lightweight linear operator library written injax.OverviewConsider solving a diagonal matrix $A$ against a vector $b$.importjax.numpyasjnpn=1000diag=jnp.linspace(1.0,2.0,n)A=jnp.diag(diag)b=jnp.linspace(3.0,4.0,n)# A⁻¹ bjnp.solve(A,b)Doing so is costly in large problems. Storing the matrix gives rise to memory costs of $O(n^2)$, and inverting the matrix costs $O(n^3)$ in the number of data points $n$.But hold on a second. Notice:We only have to store the diagonal entries to determine the matrix $A$. Doing so, would reduce memory costs from $O(n^2)$ to $O(n)$.To invert $A$, we only need to take the reciprocal of the diagonal, reducing inversion costs from $O(n^3)$, to $O(n)$.JaxLinOpis designed to exploit stucture of this kind.importjaxlinopA=jaxlinop.DiagonalLinearOperator(diag=diag)# A⁻¹ bA.solve(b)JaxLinOpis designed to automatically reduce cost savings in matrix addition, multiplication, computing log-determinants and more, for other matrix stuctures too!Custom Linear Operator (details to come soon)The flexible design ofJaxLinOpwill allow users to impliment their own custom linear operators.fromjaxlinopimportLinearOperatorclassMyLinearOperator(LinearOperator):def__init__(self,...)...# There will be a minimal number methods that users need to impliment for their custom operator.# For optimal efficiency, we'll make it easy for the user to add optional methods to their operator,# if they give better performance than the defaults.InstallationStable versionThe latest stable version ofjaxlinopcan be installed viapip:pipinstalljaxlinopNoteWe recommend you check your installation version:python-c'import jaxlinop; print(jaxlinop.__version__)'Development versionWarningThis version is possibly unstable and may contain bugs.Clone a copy of the repository to your local machine and run the setup configuration in development mode.gitclonehttps://github.com/JaxGaussianProcesses/JaxLinOp.gitcdjaxlinop
python-msetupdevelopNoteWe advise you create virtual environment before installing:conda create -n jaxlinop_ex python=3.10.0
conda activate jaxlinop_exand recommend you check your installation passes the supplied unit tests:python-mpytesttests/