Peter Hawkins 1a73743610 Move xla_bridge.constant to jax.interpreter.xla.pyval_to_ir_constant. ...

This is a more descriptive name and a better location (next to other facilities for building XLA IR).

Quite a few users of the former xla_bridge.constant() didn't need anything other than uncanonicalized array constants. Change these users to use xla_client.ops.Constant instead; no need for the fancy utility in these cases.

PiperOrigin-RevId: 404270649

2021-10-19 08:40:51 -07:00

..

image

[JAX] Canonicalize shapes in jax.image.resize().

2021-09-30 12:38:39 -07:00

lax

Move xla_bridge.constant to jax.interpreter.xla.pyval_to_ir_constant.

2021-10-19 08:40:51 -07:00

lib

Move xla_bridge.constant to jax.interpreter.xla.pyval_to_ir_constant.

2021-10-19 08:40:51 -07:00

nn

Switch internal users of jax.ops.index_... to use x.at[x].set() APIs.

2021-09-13 19:48:29 -04:00

numpy

Merge pull request #8236 from jakevdp:fix-bincount

2021-10-15 18:39:20 -07:00

ops

Add DeprecationWarnings to jax.ops.index_... operators.

2021-10-05 20:47:22 -04:00

scipy

Fix JAX functions to work if the default gather mode is set to "fill".

2021-09-30 14:21:05 -04:00

third_party

Move jax.third_party to jax._src.third_party.

2020-10-17 11:43:07 -07:00

__init__.py

Move most of the implementation of jax.scipy into jax._src.scipy.

2020-10-16 17:04:25 -04:00

abstract_arrays.py

Ensure zeros from AD are generated on device.

2021-07-26 20:40:13 +03:00

ad_checkpoint.py

Cleanup internal representation of XLA translation rules.

2021-10-16 07:53:24 -07:00

ad_util.py

Move jax.ad_util to jax._src.ad_util.

2021-06-07 14:51:34 -07:00

api_util.py

Disallow non-hashable static arguments in pmap().

2021-09-30 15:50:07 -04:00

api.py

Move xla_bridge.constant to jax.interpreter.xla.pyval_to_ir_constant.

2021-10-19 08:40:51 -07:00

cloud_tpu_init.py

Changes to make jax[tpu] work better in a docker container.

2021-07-12 17:42:46 -07:00

config.py

add a config setting to control the default PRNG implementation

2021-10-07 21:22:40 -07:00

custom_derivatives.py

Cleanup internal representation of XLA translation rules.

2021-10-16 07:53:24 -07:00

dlpack.py

Move contents of jax.lib to jax._src.lib.

2021-09-23 06:33:55 -07:00

dtypes.py

Fix inaccurate type annotations

2021-10-15 15:35:43 -07:00

errors.py

Switch internal users of jax.util.partial to use functools.partial.

2021-09-13 21:09:58 -04:00

flatten_util.py

Move implementation of jax.flatten_util to jax._src.flatten_util. Add a jax.flatten_util shim.

2021-09-08 13:54:25 -07:00

lax_reference.py

Improve implementation of cbrt() in JAX.

2021-07-22 14:01:28 -07:00

pretty_printer.py

Adds a Wadler-Lindig pretty printer.

2021-09-27 21:09:24 -04:00

prng.py

Move xla_bridge.constant to jax.interpreter.xla.pyval_to_ir_constant.

2021-10-19 08:40:51 -07:00

profiler.py

Move contents of jax.lib to jax._src.lib.

2021-09-23 06:33:55 -07:00

random.py

Cleanup internal representation of XLA translation rules.

2021-10-16 07:53:24 -07:00

source_info_util.py

Move contents of jax.lib to jax._src.lib.

2021-09-23 06:33:55 -07:00

test_util.py

Consolidate primitive and jit lowering paths.

2021-10-13 12:36:53 -04:00

traceback_util.py

Move contents of jax.lib to jax._src.lib.

2021-09-23 06:33:55 -07:00

tree_util.py

Fix jacfwd and jacrev for heterogeneous pytrees

2021-10-12 19:41:47 -04:00

util.py

Remove jax._src.util.partialmethod.

2021-10-05 12:12:41 -04:00