Alvaro ca1d8f4109

Fixing weird behavior in segment_sum when num_segments is None (#4034) ...

Co-authored-by: alvarosg <alvarosg@google.com>

2020-09-11 13:51:42 -04:00

..

experimental

[jax2tf] Fix test failure on TPUs (#4247)

2020-09-10 13:55:57 +03:00

image

Fix jax.image._resize function (#3805)

2020-07-20 13:15:40 -07:00

interpreters

Extend axis env while translating the pmapped jaxpr to XLA

2020-09-11 17:56:32 +02:00

lax

Add axis_index_groups support to all_gather. (#4194)

2020-09-09 15:02:45 +03:00

lib

Add a fast C++ jit codepath. (#4089)

2020-08-19 12:39:25 -04:00

nn

Cleanup: avoid jnp.prod & np.prod on array shapes (#4086)

2020-08-18 10:17:38 -07:00

numpy

Call check_arraylike() in jax.numpy reductions (#4195)

2020-09-11 10:04:47 -07:00

ops

Fixing weird behavior in segment_sum when num_segments is None (#4034)

2020-09-11 13:51:42 -04:00

scipy

Change lax_linalg.lu to return a permutation representation of the partial pivoting information. (#4241)

2020-09-10 11:16:35 -04:00

third_party

Cleanup: avoid jnp.prod & np.prod on array shapes (#4086)

2020-08-18 10:17:38 -07:00

tools

Update JAX to avoid XLA:Python API names deprecated in jaxlib 0.1.46. (#3046)

2020-05-11 17:43:55 -04:00

__init__.py

Import profiler in jax/__init__.py (#3719)

2020-07-11 20:44:16 -07:00

abstract_arrays.py

Optimize zeros_like_shaped_array

2020-06-05 15:52:03 +00:00

ad_util.py

omnistaging, under a flag and disabled by default (#3370)

2020-07-30 12:59:36 -07:00

api_util.py

improve vmap axis spec structure mismatch errors (#3619)

2020-06-30 22:19:16 -07:00

api.py

Revert "Revert "Delete batching.last. (#4148)" (#4160)" (#4242)

2020-09-10 09:38:14 -04:00

BUILD

Add jit to jax.image.resize (#3714)

2020-07-10 10:32:13 -04:00

config.py

Use omnistaging env var even when not using absl flags for config. (#4152)

2020-08-26 14:06:27 -07:00

core.py

Extend axis env while translating the pmapped jaxpr to XLA

2020-09-11 17:56:32 +02:00

custom_derivatives.py

Use omnistaging env var even when not using absl flags for config. (#4152)

2020-08-26 14:06:27 -07:00

dlpack.py

Update JAX to avoid XLA:Python API names deprecated in jaxlib 0.1.46. (#3046)

2020-05-11 17:43:55 -04:00

dtypes.py

Avoid direct type/dtype comparisons to fix NumPy 1.19 deprecation war… (#3543)

2020-06-24 15:19:00 -04:00

flatten_util.py

deflake

2020-06-12 15:41:49 -07:00

jaxpr_util.py

move experimental.jaxpr_stats to jaxpr_util

2020-08-18 18:07:38 -07:00

lax_linalg.py

Change lax_linalg.lu to return a permutation representation of the partial pivoting information. (#4241)

2020-09-10 11:16:35 -04:00

lax_reference.py

refine population_count type check (#3887)

2020-07-28 19:46:00 -07:00

lazy.py

deflake remainder of jax (#3343)

2020-06-06 10:51:34 -07:00

linear_util.py

Interrupt lu transformation generators whenever an exception occurs

2020-09-09 20:43:05 +02:00

pprint_util.py

Attach source info to Jaxpr equations. (#3421)

2020-06-17 16:35:36 -07:00

profiler.py

Add a jax.profiler.StepTraceContext API. (#3591)

2020-07-10 17:20:06 -07:00

random.py

[jax2tf] Fix random.split when jax_exable_x64 (#4208)

2020-09-07 14:41:50 +03:00

source_info_util.py

Attach source info to Jaxpr equations. (#3421)

2020-06-17 16:35:36 -07:00

test_util.py

Cleanup: avoid jnp.prod & np.prod on array shapes (#4086)

2020-08-18 10:17:38 -07:00

traceback_util.py

check path prefixes using os.path instead of string comparisons

2020-08-18 18:08:08 -07:00

tree_util.py

Compare treedefs by num_leaves not traversal_ in tree_transpose. (#3659)

2020-07-30 13:31:17 -04:00

util.py

Add support for negative axes to vmap. (#4111)

2020-08-24 20:21:19 -04:00

version.py

updtate version and changelog for pypi (#4224)

2020-09-08 08:54:13 -07:00