link管理
链接快照平台
- 输入网页链接，自动生成快照
- 标签化管理网页链接
相关文章推荐
|
旅途中的茄子 · 更新SQL子查询 - 腾讯云开发者社区 - 腾讯云· 1 月前 · |
|
|
叛逆的镜子 · 行情接口整理(转自知乎)_新浪关于外汇接口- ...· 4 月前 · |
|
|
曾经爱过的机器猫 · 6 Tasks — The Yocto ...· 6 月前 · |
|
|
玩篮球的沙滩裤 · 刪除頑固的Eng英語美式鍵盤 – ...· 6 月前 · |
|
|
直爽的饭盒 · 联想thinkplus真无线蓝牙耳机 ...· 6 月前 · |
-
jax.debug.printandjax.debug.breakpoint -
The
checkifytransformation - JAX debugging flags
- GPU performance tips
- Persistent Compilation Cache
- Understanding Jaxprs
- External Callbacks in JAX
- Type promotion semantics
- Pytrees
- Ahead-of-time lowering and compilation
-
Exporting and serialization
- Exporting and serializing staged-out computations
- Shape polymorphism
- Interoperation with TensorFlow
- JAX Errors
- Transfer guard
-
Pallas: a JAX kernel language
- Pallas Design
- Pallas Quickstart
-
Pallas TPU
- Writing TPU kernels with Pallas
-
Pipelining and
BlockSpecs -
Advanced Tutorials
- Training a Simple Neural Network, with tensorflow/datasets Data Loading
- Training a Simple Neural Network, with PyTorch Data Loading
- Autobatching for Bayesian Inference
- Using JAX in multi-host and multi-process environments
- Distributed arrays and automatic parallelization
-
SPMD multi-device parallelism with
shard_map - Distributed data loading in a multi-host/multi-process environment
-
Named axes and easy-to-revise parallelism with
xmap - The Autodiff Cookbook
- Custom derivative rules for JAX-transformable Python functions
-
Control autodiff’s saved values with
jax.checkpoint(akajax.remat) - How JAX primitives work
- Writing custom Jaxpr interpreters in JAX
- Custom operations for GPUs with C++ and CUDA
- Generalized Convolutions in JAX
-
Developer Documentation
- Contributing to JAX
- Building from source
- Internal APIs
- Autodidax: JAX core from scratch
-
JAX Enhancement Proposals (JEPs)
- 263: JAX PRNG Design
- 2026: Custom JVP/VJP rules for JAX-transformable functions
- 4008: Custom VJP and `nondiff_argnums` update
- 4410: Omnistaging
- 9263: Typed keys & pluggable RNGs
- 9407: Design of Type Promotion Semantics for JAX
- 9419: Jax and Jaxlib versioning
- 10657: Sequencing side-effects in JAX
- 11830: `jax.remat` / `jax.checkpoint` new implementation
- 12049: Type Annotation Roadmap for JAX
- 14273: `shard_map` (`shmap`) for simple per-device code
- 15856: `jax.extend`, an extensions module
- 17111: Efficient transposition of `shard_map` (and other maps)
- 18137: Scope of JAX NumPy & SciPy Wrappers
- Investigating a regression
- Building on JAX
-
Notes
- API compatibility
- Python and NumPy version support policy
- jax.Array migration
- Asynchronous dispatch
- Concurrency
- GPU memory allocation
- Rank promotion warning
-
Public API: jax package
-
jax.numpymodule- jax.numpy.fft.fft
- jax.numpy.fft.fft2
- jax.numpy.fft.fftfreq
- jax.numpy.fft.fftn
- jax.numpy.fft.fftshift
- jax.numpy.fft.hfft
- jax.numpy.fft.ifft
- jax.numpy.fft.ifft2
- jax.numpy.fft.ifftn
- jax.numpy.fft.ifftshift
- jax.numpy.fft.ihfft
- jax.numpy.fft.irfft
- jax.numpy.fft.irfft2
- jax.numpy.fft.irfftn
- jax.numpy.fft.rfft
- jax.numpy.fft.rfft2
- jax.numpy.fft.rfftfreq
- jax.numpy.fft.rfftn
-
jax.scipymodule- jax.scipy.stats.bernoulli.logpmf
- jax.scipy.stats.bernoulli.pmf
- jax.scipy.stats.bernoulli.cdf
- jax.scipy.stats.bernoulli.ppf
-
jax.laxmodule -
jax.randommodule -
jax.shardingmodule -
jax.debugmodule -
jax.dlpackmodule -
jax.distributedmodule -
jax.dtypesmodule -
jax.flatten_utilmodule -
jax.imagemodule -
jax.nnmodule-
jax.nn.initializersmodule -
jax.opsmodule -
jax.profilermodule -
jax.stagesmodule -
jax.treemodule -
jax.tree_utilmodule -
jax.typingmodule -
jax.exportmodule -
jax.extendmodule-
jax.extend.ffimodule -
jax.extend.linear_utilmodule -
jax.extend.mlirmodule -
jax.extend.randommodule -
jax.example_librariesmodule-
jax.example_libraries.optimizersmodule -
jax.example_libraries.staxmodule -
jax.experimentalmodule-
jax.experimental.array_apimodule -
jax.experimental.checkifymodule -
jax.experimental.host_callbackmodule -
jax.experimental.mapsmodule -
jax.experimental.pjitmodule -
jax.experimental.sparsemodule- jax.experimental.sparse.BCOO
- jax.experimental.sparse.bcoo_broadcast_in_dim
- jax.experimental.sparse.bcoo_concatenate
- jax.experimental.sparse.bcoo_dot_general
- jax.experimental.sparse.bcoo_dot_general_sampled
- jax.experimental.sparse.bcoo_dynamic_slice
- jax.experimental.sparse.bcoo_extract
- jax.experimental.sparse.bcoo_fromdense
- jax.experimental.sparse.bcoo_gather
- jax.experimental.sparse.bcoo_multiply_dense
- jax.experimental.sparse.bcoo_multiply_sparse
- jax.experimental.sparse.bcoo_update_layout
- jax.experimental.sparse.bcoo_reduce_sum
- jax.experimental.sparse.bcoo_reshape
- jax.experimental.sparse.bcoo_slice
- jax.experimental.sparse.bcoo_sort_indices
- jax.experimental.sparse.bcoo_squeeze
- jax.experimental.sparse.bcoo_sum_duplicates
- jax.experimental.sparse.bcoo_todense
- jax.experimental.sparse.bcoo_transpose
-
jax.experimental.jetmodule -
jax.experimental.custom_partitioningmodule -
jax.experimental.multihost_utilsmodule -
jax.experimental.compilation_cachemodule -
jax.experimental.key_reusemodule -
jax.experimental.mesh_utilsmodule -
jax.experimental.serialize_executablemodule -
jax.experimental.shard_mapmodule -
jax.libmodule
jax.numpy.
vectorize
(
pyfunc
,
*
,
excluded
=
frozenset({})
,
signature
=
None
)
[source]
#
-
Define a vectorized function with broadcasting.
vectorize()is a convenience wrapper for defining vectorized functions with broadcasting, in the style of NumPy’s generalized universal functions . It allows for defining functions that are automatically repeated across any leading dimensions, without the implementation of the function needing to be concerned about how to handle higher dimensional inputs.jax.numpy.vectorize()has the same interface asnumpy.vectorize, but it is syntactic sugar for an auto-batching transformation (vmap()) rather than a Python loop. This should be considerably more efficient, but the implementation must be written in terms of functions that act on JAX arrays. - Parameters :
-
pyfunc – function to vectorize.
-
excluded – optional set of integers representing positional arguments for which the function will not be vectorized. These will be passed directly to
pyfuncunmodified. -
signature – optional generalized universal function signature, e.g.,
(m,n),(n)->(m)for vectorized matrix-vector multiplication. If provided,pyfuncwill be called with (and expected to return) arrays with shapes given by the size of corresponding core dimensions. By default, pyfunc is assumed to take scalars arrays as input and output. - Returns :
-
Vectorized version of the given function.
Here are a few examples of how one could write vectorized linear algebra routines using
vectorize():>>> from functools import partial>>> @partial(jnp.vectorize, signature='(k),(k)->(k)') ... def cross_product(a, b): ... assert a.shape == b.shape and a.ndim == b.ndim == 1 ... return jnp.array([a[1] * b[2] - a[2] * b[1], ... a[2] * b[0] - a[0] * b[2], ... a[0] * b[1] - a[1] * b[0]])>>> @partial(jnp.vectorize, signature='(n,m),(m)->(n)') ... def matrix_vector_product(matrix, vector): ... assert matrix.ndim == 2 and matrix.shape[1:] == vector.shape ... return matrix @ vectorThese functions are only written to handle 1D or 2D arrays (the
assertstatements will never be violated), but with vectorize they support arbitrary dimensional inputs with NumPy style broadcasting, e.g.,>>> cross_product(jnp.ones(3), jnp.ones(3)).shape >>> cross_product(jnp.ones((2, 3)), jnp.ones(3)).shape (2, 3) >>> cross_product(jnp.ones((1, 2, 3)), jnp.ones((2, 1, 3))).shape (2, 2, 3) >>> matrix_vector_product(jnp.ones(3), jnp.ones(3)) Traceback (most recent call last): ValueError: input with shape (3,) does not have enough dimensions for all core dimensions ('n', 'k') on vectorized function with excluded=frozenset() and signature='(n,k),(k)->(k)' >>> matrix_vector_product(jnp.ones((2, 3)), jnp.ones(3)).shape >>> matrix_vector_product(jnp.ones((2, 3)), jnp.ones((4, 3))).shape (4, 2)Note that this has different semantics than jnp.matmul:
>>> jnp.matmul(jnp.ones((2, 3)), jnp.ones((4, 3))) Traceback (most recent call last): TypeError: dot_general requires contracting dimensions to have the same shape, got [3] and [4].
-
-
-
-
-
推荐文章
|
|
旅途中的茄子 · 更新SQL子查询 - 腾讯云开发者社区 - 腾讯云 1 月前 |