jax

jax copied to clipboard

Add functionality for "pure" callbacks Also avoids using CPP dispatch path when host callbacks are involved

copybara-service[bot]

[Pax] Support checkpoint policies in pipeline layer.

copybara-service[bot]

`JaxTestCase` now sets this to strict by default, so these annotations are no longer necessary: https://github.com/google/jax/blob/af18235ea30799ed2fc50e557ebba869e2bdbd41/jax/_src/test_util.py#L706-L713

jakevdp

Different runtimes for `pmap` vs `pjit`

4

Hi, I have a simple script below that compares runtimes for `pmap` vs `pjit`. I expected that the runtime for `pjit` with full data parallelism would be the same for...

wilson1yan

To use the newly added bindings for use_globsl_device_ids in AllReduce

apaszke

Allow collectives in manually sharded computations ... at least when the manual sharding applies to the whole mesh, because that's all that XLA can support right now. This is especially...

copybara-service[bot]

Bogus gradient value from value_and_grad for committed DeviceArray on a multi-GPU host

6

Hi, I'm reporting this issue with a disclaimer that it's probably not reproducible (sorry!) but hoping it's of some worth for the awesome JAX developers: We have a machine with...

yiiyama

`jax.lax.dot` with `preferred_element_type` gives wrong result

3

The bug can be reproduced by the following code ``` import jax.numpy as jnp from jax.random import PRNGKey, normal from jax.lax import dot from jax.config import config config.update("jax_enable_x64", True) A...

ChenAo-Phys

Fix a bug where `get_backend` in host callback lowering was provided with a wrong platform

1

Fix a bug where `get_backend` in host callback lowering was provided with a wrong platform

copybara-service[bot]

[call_tf] implementation of batching rule for vmap

6

Hi, I'm trying to use call_tf in combination with jacrev and hence vmap. The primitive_batcher of type call_tf_p is not implemented. I tried to modify the source code of `jax/experimental/jax2tf/call_tf.py`...

DavidDevoogdt