google / jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

http://jax.readthedocs.io/

google/jax Issues

Implement `scipy.optimize.linear_sum_assignment`
Updated 18 days ago1
NCCL error when trying xmap on 4 T4 GPUs
Updated 19 days ago
jaxlib windows build runtime breaks due to incompatible `np.array` dtype with `ir.IndexType`
Updated 19 days ago
Unhelpful error messages related to pmap leading dimension requirements
Updated 19 days ago5
Autodiff with `numpy.linalg.qr` differs from autodiff with `scipy.linalg.qr`
Closed 19 days ago6
Extremely long compilation time for pmap on a5000 GPUs
Updated 20 days ago4
Add random.orthogonal and random.unitary
Updated 20 days ago2
Improved docs for (default) matmul precision
Updated 20 days ago
XOR for sparse matrix
Closed 21 days ago3
jnp.argsort much slower than the numpy version
Updated 21 days ago5
[Docs] Add installation instructions to documentation
Closed 21 days ago3
Backpropagating `take` fails for complex arrays
Closed 22 days ago6
[CPU] `np.linalg.eigh` outputs NaNs for matrix side size >=32767
Updated 24 days ago1
[CPU] `np.linalg.eigh` crash on a 50k x 50k matrix
Updated 24 days ago
[Docs] Add documentation URL to GitHub sidebar
Closed 24 days ago1
Potential bug in forward-mode AD with NaNs in intermediate results
Closed 25 days ago2
Control the parallelism on CPU
Updated 25 days ago
loops.Scope() arguments to iterate over
Updated a month ago1
jax.local_devices('cpu') raises strange error
Closed a month ago3
`tree_flatten` gets `None` when used in `defvjp`
Closed a month ago9
numerical instability under different shapes in `jnp.linalg.norm` when `ord="fro"`
Closed a month ago7
TPUv3 CHECK failure at all_gather_emitter.cc:1731) hlo_instruction_->operand(0)->shape().element_type() == hlo_instruction_->shape().element_type()
Closed a month ago1
Implement `jax.scipy.stats.bootstrap`
Updated a month ago
vmap + array-indexing is very slow on TPU
Updated a month ago1
Implement `at` method for `_IndexUpdateRef`
Updated a month ago8
`KeyError: <class 'bfloat16'>` after using `lax._convert_element_type`
Updated a month ago1
GCC 11 support
Closed a month ago5
Low-level error in forward-mode AD of `lax.reduce_window_sum_p`
Closed a month ago
Low-level error when calling `lax.convert_element_type_p` with uncanocinalized `new_dtype`
Closed a month ago3
unjitted jnp.where evaluates both branches
Closed a month ago1
Expose optimization path of L-BFGS optimizer
Closed a month ago3
`jnp.take` out of bound documentation incorrect
Closed a month ago1
Jax cannot use A100 GPU despite CUDA 11 version being installed
Closed a month ago2
XLA: log-indexing should be optimised to indexing-log
Updated a month ago3
Cannot change `jax_platform_name` on the fly
Updated a month ago1
AttributeError: module 'jaxlib.xla_extension' has no attribute 'CpuDevice'
Closed a month ago4
Generic Protocol for transformed function
Updated a month ago10
`TracerArrayConversionError` with numpy in the function for autodiff
Closed a month ago3
[How?] Debugging crash in jax.jit function when running under pytest-xdist
Updated a month ago4
Running multiple runs on one A100
Closed a month ago2
pre-commit mypy check seems outdated
Closed a month ago2
ImportError: cannot import name 'index' from 'jax.ops'
Closed a month ago1
JIT fails to optimise batched array indexing: huge performance difference
Closed a month ago2
`cuSparseTest.test_coo_sorted_indices_gpu_lowerings` test fails when running on a CPU-only machine
Closed a month ago4
Surprising behavior of `jax.scipy.special.expit` and `jax.scipy.special.logit`
Closed a month ago
Integrating TVM into JAX
Updated a month ago1
jax.scipy.sparse.linalg.cg inconsistent results between runs
Updated a month ago1
TypeError when using None as custom_vjp cotangent for custom pytree and asarray
Closed a month ago15
Add CustomCall support to jax2tf
Updated a month ago2
NonConcreteBooleanIndexError on call to jnp.unique
Updated a month ago5

ezoic increase your site revenue