[JAX] Add an opaque `extra_jit_context` field to the JAX C++ jit code.
This allows the JAX Python code to include extra context from, for example, the interpreter state as part of the C++ jit cache key.
PiperOrigin-RevId: 364611475
[JAX] Add an opaque `extra_jit_context` field to the JAX C++ jit code.
This allows the JAX Python code to include extra context from, for example, the interpreter state as part of the C++ jit cache key.
PiperOrigin-RevId: 364599983
This allows the JAX Python code to include extra context from, for example, the interpreter state as part of the C++ jit cache key.
PiperOrigin-RevId: 364563982
--
bf15ba5310d5f9009571928f70548bcbc7e856c3 by Matthew Johnson <mattjj@google.com>:
don't device transfer in convert_element_type
Co-authored-by: Qiao Zhang <zhangqiaorjc@google.com>
PiperOrigin-RevId: 363995032
Reimplement Sublevel to not inherit from `int`.
See docs on weakref: "CPython implementation detail: Other built-in types
such as tuple and int do not support weak references even when subclassed."
Previously any collectives buried inside control flow would fail to
compile with xmap, because it would not traverse those with its name
substitution. This adds a "catch-all" default substitution rule which
recursively applies to all jaxpr found in the params (at the top level).
Before raising an error on an unrecognized type, first check if the
object defines a __jax_array__ method. If it does, call it!
This provides a way for custom types to be auto-converted to
JAX-compatible types.
Implementing this method is not sufficient for a type to be duck-typed
enough for use with jax.numpy. But it may be necessary. That is, someone
trying to add a duck-typed array to be used with JAX identified a need
for __jax_array__ or similar. The user would still need to add lots of
other properties and methods, like dtype and shape attributes.
revives #4725 after it was rolled back. fixes#5356.
The second change in the avals-with-names stack:
- https://github.com/google/jax/pull/5524 Revise aval constructor call sites to use a new `aval.update` method
- **Add `named_shape` to `ShapedArray` and update typecompat**
- Propagate presence of name (mapped) vs absence (replicated) in abstract eval based on existing batching rules
- Make `mapped_aval`, `unmapped_aval`, and their xmap equivalents swap positional and named axes (rather than just creating and deleting positional ones)
- Enable `lax.full` to create values with named axes
- Ensure `grad` and `jacfwd`/`jacrev` consistently act elementwise over named axes (by e.g. using a seed with named axes in `grad`, and prohibiting collectives if TAP isn't too unhappy) and align `vmap(transpose)` with `transpose(vmap)` by moving the `psum` in `transpose(psum)` into `backward_pass`
- Add `axis_name` kwarg to grad to indicate operating collectively over one or more named axes
PiperOrigin-RevId: 355880632
... as well as in a few random functions that it needs (`uniform`,
`normal` and `truncated_normal`). The interface itself doesn't change to
much with the exception of the `shape` arguments of all those functions
now accepting `jax.core.NamedShape` (I didn't move it to be part of the
API just yet, but we can do that any time), which makes it possible to
generate sharded random arrays (in particular the random bits are
different on different shards). I also haven't updated the docstrings,
because I don't know if we're ready to go fully public with this
feature.
Previously, a few places in our code assumed that all collectives (i.e.
primitives that operate over named axes) keep all of their axes in the
`axis_name` attribute. This was fine for a few simple use cases, but we
are now considering allowing named axes in many more primitives which
can have semantically different attributes where axis names can appear.
This change adds to the error message when we hit an escaped tracer. In
particular, it adds source info for the function that was transformed.
This change currently only applies to escaped `DynamicJaxprTracer`s
(arising from `jit`, `pmap`, `scan`, and other staging functions) and
not other traces. A natural follow-up would be to attach this
information to other traces.
Co-authored-by: Lena Martens <lenamartens@google.com>
This extra source info is still only on jaxpr staging tracers, but those
seem to be the most common culprits. I moved the `_line_info` attribute
to the base Tracer class in core.py in anticipation of populating it for
more traces than just DynamicJaxprTrace, but I'll leave that extension
to follow-up.
I adapted the main escaped tracer error messages in core.py, and also
slightly generalized and debugged source_info_util functions (thanks for
explaining the path prefix bug, @froystig !).
--
7342318774c6f1195f0e238f1209425109ea8944 by Matthew Johnson <mattjj@google.com>:
check for __jax_array__ method for conversion
--
6742016382b0511f5ac9ec21f67d2122a9f37cb7 by Matthew Johnson <mattjj@google.com>:
fix typo
--
5eb36855e53d8d4e81e281d08dc9264d2671f21f by Matthew Johnson <mattjj@google.com>:
ensure some jnp funs duck-type with __jax_array__
PiperOrigin-RevId: 347763582
AD didn't use `HashableFunction` enough, tripping up the compilation
cache. I've also used the occasion to make function hashing a little
safer by including the Python bytecode of the wrapped function as part
of the key.
Specifically we:
1. remove the need for split_axis rules in batching.py, and instead just
rely on collective rules (namely to handle vectorizing over a single
named axis even if the collective is applied over multiple named axes)
2. simplify BatchTrace.process_primitive so that we don't pass tracers
into rules and rely on a subtle recursion
This change breaks all_to_all when used with multiple axis names, and in
particular it breaks all_to_all given the current gmap/xmap lowering
strategy of substituting multiple axis names in place of single axis
names. We believe we can replicate the previous logic with the new rule
organization, but we're leaving that for follow-up work because it's
tricky, and because we might end up changing lowering strategies not to
require axis substitution in the same way.
Previously `pmap` didn't have the `out_axes` parameter (unlike `vmap`),
but its semantics would match the specification of `out_axes=0` (i.e.
all outputs should be stacked along the first axis). This patch makes it
possible to specify non-zero values for out_axes, but more importantly
it lays down the groundwork for `xmap` which will have to use some
extremely similar (if not the same) code paths.
One thing to note is that when I started this implementation I was also
planning to add support for `out_axes=None`, which would allow us to
stop using the `unbroadcast` hack, and most of the code is written with
that in mind. Unfortunately it turned out that the correct
implementation of the transpose rule for maps that do allow unmapped
outputs would require me to pretty much simulate what avals-with-names
is supposed to achieve. Technically replicated outputs should work
today, for as long as the user does not do reverse-mode AD of `pmap`.
But I decided that it's better to just disable them altogether until we
can get the full and correct behavior.
* Implementation details *
This patch is significantly more involved than the one that implemented
general `in_axes` support. That previous one at least had the foundation
of `mapped_invars` which already behaved pretty similarly to general
`in_axes`. From a quick glance one might think that `out_axes` should
behave similarly to `in_axes`, but it turns out that this is not the
case, at least not if we're interested in keeping those primitives
final-style.
** Thunking **
The biggest difficulty with handling `out_axes` in final style
primitives is that we want to treat them as a prefix of the output
pytree, but we don't know the structure of the output pytree until the
user function is evaluated! And the user function is not evaluated until
we've applied all transforms and reached the impl rule! The solution to
this problem is "straightforward": instead of putting `out_axes` as a
primitive parameter, we bundle an `out_axes_thunk` which can only be
called successfully after the wrapped function has been executed. The
thunk returns a list of flat `out_axes`, expanded to the output pytree.
However, the thunking presents us with two problems:
*** Transformations ***
Each transformation that modifies the number of outputs needs to ensure
that the thunk is updated to reflect the new values. To make things
worse a lot of the transforms can learn the number of added outputs
_only after the wrapped function is evaluated_, which leads to the
following "time travel" pattern that can be found in most `Trace`s:
```py
@lu.transformation_with_aux
def compute_output_statistic(*args, **kwargs):
outputs = yield args, kwargs
yield outputs, compute_statistic(outputs)
wrapped_fun, output_statistic = compute_output_statistic(wrapped_fun)
def new_out_axes_thunk():
old_out_axes = params['out_axes_thunk']()
return compute_new_out_axes(old_out_axes(), output_statistic())
primitive.bind(wrapped_fun, dict(params, out_axes_thunk=new_out_axes_thunk))
```
The reason why we have to structure the code this way is that we can
only specify a new `out_axes_thunk` before we bind the primitive, but we
need the outputs of bind to know how to update the `out_axes_thunk`. To
make things worse, the implementation of `bind` is allowed to make a
call to `out_axes_thunk` _immediately after `wrapped_fun` is evaluated_.
This means that we cannot compute the output statistic in the
implementation of the transformation, but we have to use an extra
`lu.transformation_with_aux` for that (this populates the statistic
store immediately after `wrapped_fun` is evaluated).
The `compute_statistic` function depends on the transform in question.
E.g. in the JVP trace it counts the number of non-zero tangent results.
The situation is of course further complicated when we take
`post_process_map` into account. The new `process_env_traces` now always
sets up this funny time travel trampoline just in case it ends up being
necessary, and `post_process_map` is now expected to return `(outputs,
(todo, out_axes_transform))` instead of just `(outputs, todo)`.
*** Compilation cache ***
Because the `out_axes_thunk`s are now arguments to a _global_
compilation cache (in the form of `lu.cache` decorator on
`parallel_callable`), we have to ensure that they implement `hash` and
`==`. This is what forces us to add some slightly weird helpers such as
`_hashable_function` and `_ignore_elem_list`. The code that uses those
makes an assumption that the output pytree depends deterministically on
the identity of the wrapped function, which I think is in line with
general JAX assumptions. Otherwise the cache would depend on the
identity of the thunk, which changes with every function invocation.
Relaxing the global constraint on the cache (e.g. allowing each
`pmap(f)` instance to have a separate cache) would make this easier too.
* Why final style? *
Now, making the primitives initial-style would remove the necessity for
thunking, because we could have obtained the output pytree right when
the function is wrapped. I assumed there is a good argument for making
`pmap` pretend that it's a final-style primitive, but I'm not sure why
that is? I hope it's something better than just avoiding a single jaxpr
tracing.
This should allow us to try out xmap not only in a simulation (i.e.
faking the devices using vmap, which we still support), but also on real
hardware.
Limitations:
- No compilation caching yet
- Nested xmaps not supported yet
- Transforms (AD, vmap, etc.) of xmaps not supported yet
Benefits:
- An xmap over multiple mesh axes already implements a more efficient
lowering than the one used for nested pmaps.
The `resources` context-manager is now called `fake_resources`, while
real meshes can be defined in a specific context using the
`mesh(devices, axis_names)` manager. `devices` is supposed to be an
`ndarray` of JAX device objects (e.g. obtained from `jax.devices()`),
while `axis_names` should be a tuple of length matching the rank of
`devices` and specifying mesh axis names.
For concrete examples see the changes in `gmap_tests.py`.
In principle the current version of the code should also work in a
multi-host setting, but I haven't tested it just yet.
... and in map primitives in general (which is why the patch touches
most traces).
This also fixes a bug in the transpose rule for map primitives, which
would fail to adjust the aval associated with zeros returned from the
map body.
