google / jax

52 Pull requests merged by 15 people

• Short-circuit references to jax.core via jax.abstract_arrays.

  #4949 merged Nov 19, 2020

• Fix build.py to work on Linux once again.

  #4965 merged Nov 19, 2020

• Make dtype an explicit argument to nan-reductions

  #4938 merged Nov 19, 2020

• Deprecation: Raise a TypeError when indexing with a non-tuple sequence

  #4867 merged Nov 19, 2020

• [JAX] Move contents of jax.random into jax._src.random.

  #4964 merged Nov 19, 2020

• Add .bazelversion file.

  #4962 merged Nov 19, 2020

• Build on Windows

  #4843 merged Nov 19, 2020

• Remove jax.experimental.optix.

  #4961 merged Nov 19, 2020

• [JAX] Delete jax.lax_linalg.

  #4958 merged Nov 19, 2020

• Temporarily disable JAX tests that are failing due to LLVM commit https://github.com/llvm/llvm-project/commit/fcad8d3635cff61a2749dcef94c0d51fa1e3e413

  #4960 merged Nov 19, 2020

• [JAX] Move pprint_util into jax._src.

  #4956 merged Nov 19, 2020

• [JAX] Delete jax.source_info_util.

  #4955 merged Nov 19, 2020

• [JAX] Move implementation of jax.scipy.sparse.linalg into jax._src.

  #4957 merged Nov 19, 2020

• Add multi-host support to sharded_jit.

  #4746 merged Nov 19, 2020

• Add initial changelog for jax v0.2.7

  #4947 merged Nov 19, 2020

• update version & CHANGELOG for jax v0.2.6

  #4946 merged Nov 18, 2020

• [jax2tf] Added tests for the conversion of transpose.

  #4926 merged Nov 18, 2020

• Add support for XLA variadic reduce

  #4868 merged Nov 18, 2020

• Remove noisy warnings for var() of int

  #4940 merged Nov 18, 2020

• [masking] Ensure that Poly operations are sound

  #4878 merged Nov 18, 2020

• Fix mypy error caused by cyclic import dependency.

  #4944 merged Nov 18, 2020

• Use the jax module to determine the jax source directory base, rather…

  #4943 merged Nov 18, 2020

• fix grad-of-switch weak type mismatch test

  #4916 merged Nov 18, 2020

• Support shape parameter in jax.numpy.*_like()

  #4932 merged Nov 18, 2020

• lax_numpy: fix handling of out keyword in wrapped methods

  #4929 merged Nov 18, 2020

• Preserve name scope of named_call primitive in jax2tf

  #4649 merged Nov 17, 2020

• Add a named_call public function in lax

  #4796 merged Nov 17, 2020

• [jax2tf] Add harness to test conversion of zeros_like.

  #4927 merged Nov 17, 2020

• Filter out non-integer index_dtypes in argmin/max.

  #4921 merged Nov 17, 2020

• fix mypy error

  #4936 merged Nov 17, 2020

• [jax2tf] Add testing for conversion of comparison operators.

  #4912 merged Nov 17, 2020

• Update bazel version to 3.1.0.

  #4934 merged Nov 17, 2020

• Accept `None` as a trivial LazyExpr.

  #4933 merged Nov 17, 2020

• Relax test tolerance on polymul tests.

  #4930 merged Nov 17, 2020

• Update in1d() TODO

  #4928 merged Nov 17, 2020

• [jax2tf] Add testing for argmax and argmin.

  #4883 merged Nov 17, 2020

• [jax2tf] Add tests for the conversion of concatenate.

  #4901 merged Nov 17, 2020

• Change the DeviceArray error from ValueError to RuntimeError when del…

  #4913 merged Nov 17, 2020

• [jax2tf] Add testing for select_p conversion.

  #4925 merged Nov 17, 2020

• [jax2tf] Fix conversion of clamp and add testing code.

  #4906 merged Nov 17, 2020

• Fix dtypes.is_weakly_typed

  #4914 merged Nov 17, 2020

• Fix select_and_gather_add_transpose rule.

  #4894 merged Nov 16, 2020

• Lower complex128 scatter-add to two float64 scatter-adds on GPU.

  #4907 merged Nov 16, 2020

• [jax2tf] Fix mypy error in harnesses.

  #4908 merged Nov 16, 2020

• Update XLA.

  #4910 merged Nov 16, 2020

• use `typematch` rather than `==` in cond transpose assertion

  #4896 merged Nov 16, 2020

• [jax2tf] Add a harness for broadcast_p and fix conversion.

  #4874 merged Nov 15, 2020

• Fix doc for index operations

  #4803 merged Nov 13, 2020

• Rm redundant pre-omnistaging attribute annotations from core.TraceStack

  #4854 merged Nov 13, 2020

• jaxlib: Fix Python 3.9 build and drop CUDA 10.0 support

  #4886 merged Nov 13, 2020

• fix typo in autodiff cookbook

  #4890 merged Nov 12, 2020

• README etc. updates for new jaxlib release

  #4887 merged Nov 12, 2020

15 Pull requests opened by 11 people

• Cloned from CL 333372978 by 'g4 patch'.

  #4892 opened Nov 13, 2020

• More stable implementation of np.hypot

  #4893 opened Nov 13, 2020

• Add logging for pmap input (re)sharding

  #4895 opened Nov 13, 2020

• ci-build.yaml: mypy jax failing on the current version of Python

  #4897 opened Nov 15, 2020

• Add support for out_axes in pmap

  #4904 opened Nov 16, 2020

• Extract cloud TPU Colab set-up to a re-usable helper.

  #4911 opened Nov 16, 2020

• update jax repository citation

  #4919 opened Nov 17, 2020

• Improve pmap caching

  #4923 opened Nov 17, 2020

• Add derivative for eigenvalues

  #4941 opened Nov 18, 2020

• Add vector-valued gradient example to quickstart

  #4948 opened Nov 19, 2020

• Remove devices from xla.AxisEnv

  #4953 opened Nov 19, 2020

• Add XLA lowering for xmap

  #4954 opened Nov 19, 2020

• [jax2tf] Fix and test conversion of round.

  #4959 opened Nov 19, 2020

• [jax2tf] Added testing code for div/rem

  #4963 opened Nov 19, 2020

• jax.numpy: match numpy arguments for reductions

  #4969 opened Nov 19, 2020

17 Issues closed by 8 people

• Add support for Jax on Windows

  #438 closed Nov 19, 2020

• Omnistaging slows execution time by a factor of 100

  #4902 closed Nov 18, 2020

• Not able to calculate logpdf using jax.scipy.stats.multivariate_normal

  #4931 closed Nov 17, 2020

• A nasty bug it seems

  #4909 closed Nov 17, 2020

• grad of function containing skipped singularities is nan

  #4915 closed Nov 17, 2020

• Slow gradient calculation when indexing a complex-valued object inside a loss.

  #4905 closed Nov 16, 2020

• Cannot print a doubly-reshaped ShardedDeviceArray

  #3597 closed Nov 16, 2020

• np.cint/np.int32 type confusion

  #4903 closed Nov 16, 2020

• weak_types: error in grad of switch

  #4696 closed Nov 16, 2020

• Missing libcublasLt.so on CUDA 10.0

  #4705 closed Nov 15, 2020

• psum inconsistent across vmap/pmap context

  #4327 closed Nov 14, 2020

• CUDNN_STATUS_INTERNAL_ERROR

  #4884 closed Nov 13, 2020

• illegal memory access while two GPUs is used

  #3518 closed Nov 13, 2020

• Support for variable length in jitted scans?

  #4473 closed Nov 13, 2020

• Behler Parinello network in example Colab notebook

  #4881 closed Nov 12, 2020

• Jax slow on cuda 11.0

  #4889 closed Nov 12, 2020

• `triangular_solve` does not work well with batching

  #4888 closed Nov 12, 2020

14 Issues opened by 12 people

• Scanning over reduction rather than reducing over scan is a powerful optimisation - but hard to implement

  #4968 opened Nov 19, 2020

• Add "float64 only" flag

  #4967 opened Nov 19, 2020

• Add "strict dtype checking" flag

  #4966 opened Nov 19, 2020

• [jax2tf] Incorrect translation for jax.lax.round

  #4952 opened Nov 19, 2020

• NCCL out of memory on dual 2080ti training

  #4950 opened Nov 19, 2020

• [jax2tf] jax.random.randint bug

  #4942 opened Nov 18, 2020

• expit outputs 0 when maybe it shouldn't - but it's probably fine

  #4939 opened Nov 17, 2020

• Accelerating jacobian with static_argnums

  #4937 opened Nov 17, 2020

• Very slow compilation of gradient (esp. on CPU)

  #4924 opened Nov 17, 2020

• GPU not used for processing

  #4922 opened Nov 17, 2020

• Unimplemented: DNN library is not found.

  #4920 opened Nov 17, 2020

• `nondiff_argnums` for `custom_vjp` primitives?

  #4918 opened Nov 17, 2020

• jacrev and nan

  #4900 opened Nov 16, 2020

• change grad complex conjugation convention?

  #4891 opened Nov 12, 2020

21 Unresolved conversations

Sometimes conversations happen on old items that aren’t yet closed. Here is a list of all the Issues and Pull Requests with unresolved conversations.

• Internal runtime error?

  #4869 commented on Nov 17, 2020 • 7 new comments

• [RFC] Support complex dtype in random.normal

  #4805 commented on Nov 13, 2020 • 5 new comments

• autodiff support for jax.numpy.linalg.eig

  #2748 commented on Nov 18, 2020 • 3 new comments

• jit fails for large uint64 inputs

  #4809 commented on Nov 13, 2020 • 2 new comments

• Automatically doing TensorFloat32 Math on Ampere GPUs

  #4873 commented on Nov 15, 2020 • 2 new comments

• Using vmap inside defjvp

  #3201 commented on Nov 17, 2020 • 2 new comments

• Handle degenerate eigenvalues in JVP rule for eigh

  #1665 commented on Nov 18, 2020 • 2 new comments

• Add typing annotations and expose py.typed

  #4711 commented on Nov 13, 2020 • 2 new comments

• [Documentation? + Feature Req] Pytrees do not appear to work with `jax.experimental.loops`

  #4600 commented on Nov 13, 2020 • 1 new comment

• No tools in profiler

  #4427 commented on Nov 13, 2020 • 1 new comment

• Feature request: add jax[cudaversion] to pypi

  #2453 commented on Nov 13, 2020 • 1 new comment

• jax2tf.convert: handy when used as a decorator, but can't take arguments at the moment

  #4859 commented on Nov 13, 2020 • 1 new comment

• feature request: sparse jacobian and sparse hessians

  #1032 commented on Nov 16, 2020 • 1 new comment

• Hot wiring xla'd kernels

  #4866 commented on Nov 17, 2020 • 1 new comment

• Performance of index_update on GPU with complex double data type

  #4115 commented on Nov 17, 2020 • 1 new comment

• XLA profiling on TPU

  #4754 commented on Nov 17, 2020 • 1 new comment

• optimize device_put of ShardedDeviceArray

  #4885 commented on Nov 17, 2020 • 1 new comment

• Should JAX deprecate indexing with lists?

  #4564 commented on Nov 19, 2020 • 1 new comment

• Support DCE on `make_jaxpr`

  #4818 commented on Nov 19, 2020 • 1 new comment

• DeviceArray.__iter__ returns DeviceArrays, without host sync

  #3821 commented on Nov 17, 2020 • 1 new comment

• experiments in the implementation of vmap collectives

  #4835 commented on Nov 16, 2020 • 0 new comments