--
4fcdadbfb3f4c484fd4432203cf13b88782b9311 by Matthew Johnson <mattjj@google.com>:
add jax.ensure_compile_time_eval to public api
aka jax.core.eval_context
COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/7987 from google:issue7535 4fcdadbfb3f4c484fd4432203cf13b88782b9311
PiperOrigin-RevId: 420928687
The pretty-printing changes a few months ago defined variable names
based on the state in JaxprPpContext instances. But that meant incorrect
variable names could be printed in jaxpr type checking error messages.
This commit correctly threads through the context so as to provide
error messages with coherent variable names.
Final-style higher-order primitives, like call_p, xla_call_p (underlying
jit), xla_pmap_p (underlying pmap), and xmap_p (underlying xmap) have
slightly different bind signatures (while tracing) from their signatures
when they appear in jaxprs. In particular, their trace-time binds are
parameterized by a Python callable (or really a lu.WrappedFun)
representing the function to be applied, while in jaxpr eqns they are
parameterized by a jaxpr representing the same.
As a result, to round-trip from jaxpr to Python traceable, in
core.eval_jaxpr we have to convert from one parameter signature to the
other. (Basically we had to take the jaxpr and turn it into a Python
callable, via lu.wrap_init(partial(core.eval_jaxpr, call_jaxpr, ...)).)
However due to historical path dependence these conversion mechanisms
were all slightly distinct and kind of a mess. There was a case analysis
for call_jaxpr and map_jaxpr in core.eval_jaxpr_eqn (a helper function
created only because of this complexity), and there was a separate table
only used for the xmap rule.
In this PR we uniformized things! We basically only have a table (to
simplify core.eval_jaxpr), but instead of having it as a table we just
attached the rules to the different primitive classes (CallPrimitive,
MapPrimitive, and XmapPrimitive) to make things less error-prone (we
have a few different CallPrimitive instantiations, like call_p,
xla_call_p, named_call_p, and remat_call_p, and this way we don't have
to remember to populate the table separately for each).
This was actually a warmup simplification before we attempt to simplify
custom derivatives (to unify custom_jvp_call_p and
custom_jvp_call_jaxpr_p).
Co-authored-by: Roy Frostig <frostig@google.com>
Instead, just give AxisPrimitive its own bind function. This way the
logic is nicely separated by concerns. In addition, this factorization
will let us more easily experiment with other ways to find the top trace
(e.g. for assert_p in checkify).
A callback under ad_checkpoint.checkpoint will be invoked
twice when taking the gradient: once during the forward pass
and once again during the backward pass when the residuals
for the forward pass are rematerialized.
mro() has a different signature on metaclasses, but __mro__ is a cached tuple property that appears to have the same signature everywhere. As far as I can tell, it always exists.
PiperOrigin-RevId: 416410647
This is a strictly mechanical change that moves abstract value canonicalization out of the core.AbstractValue subclasses and into their callers. This makes it safe to manipulate non-canonical abstract values even inside an -x32 context.
The callers to which canonicalization was added were:
a) all callers of `ConcreteArray` inside the JAX Tree.
b) all callers of `ShapedArray` and `UnshapedArray` that were found to be passing non-canonical dtypes during a global presubmit. These were identified by adding an assertion that the dtype is in fact canonical and fixing all the resulting test failures.
PiperOrigin-RevId: 414704700
* jax._src.device_array, which contains the definition of DeviceArray.
* jax.interpreters.xla, which contains code for lowering jaxprs into XLA computations.
* jax._src.dispatch, which contains code for executing primitives and jit-compiled functions (xla_call_p's impl logic).
The purpose of splitting up this file is that I would like to treat jax.interpreters.mlir lowering as an alternative to jax.interpreters.xla, but we wish to share the device_array and computation dispatch pieces. Currently jax.interpreters.mlir duplicates most of the dispatch logic. (That refactoring is for a future change; this change just moves the existing code around.)
PiperOrigin-RevId: 411565432
also:
* fix jit invariance bug around weak types
* elide trivial broadcasts
This started as an attempt to simplify some jaxpr pretty-prints, by (1)
eliding some convert_element_type applications that I thought were
unnecessary and (2) eliding some trivial broadcasts.
But it turned out that we were actually pruning more
convert_element_types than we should! In particular, see
test_weak_type_jit_invariance; that test fails on the main branch even
if we add the fixes in DynamicJaxprTrace.new_const, because [this
logic](b53a174042/jax/interpreters/partial_eval.py (L1225))
was not paying attention to weak types and hence clobbered them.
In addition to fixing those bugs that turned up (the changes in
DynamicJaxprTrace, and in what is now _convert_elt_type_fwd_rule), this
PR generalizes the jaxpr simplification machinery so as not to be a
couple special cases on convert_element_type_p. Insetad, we have tables
of rules! How we love them.
These rule signatures should let us add simplifications like forwarding
variables through calls and other higher-order primitives. That's all
future work though.
Bug: 8367
Small refactoring to jax.image.resize to make it compatible with
shape polymorphismin jax2tf. In the process added also support for
jnp.arange([dim_poly]). Note that the underlying lax.iota already
supported shape polymorphism.
--
b40245e38d7837a7777735ad60f3b5b1ac2d499d by Sharad Vikram <sharad.vikram@gmail.com>:
Use `SourceInfo` named tuple to keep track of source information
PiperOrigin-RevId: 406293469
Previously jax.core.DropVar was a singleton value (jax.core.dropvar) whose type was always jax.core.AbstractUnit. However, this type is misleading: a DropVar is an equation output, and typically we would expect it to have an array type. In particular, the unit type confuses new-style translation rules that expect to use the output aval on an equation as part of the lowering logic.
Instead, change DropVar to be a non-singleton subclass of Var instead with a flexible choice of aval.
PiperOrigin-RevId: 404071001
The number of buffers used to represent an abstract value is a property specific to a particular representation of that abstract value. Currently the only representation is an XLA representation, but that may change in the future. Instead, callers who want to know how XLA would represent an aval should ask the XLA module instead. In this case, we call len(xla.aval_to_xla_shapes(...)) instead.
It was confusing to overload, since we sometimes think of avals like
shapes paired with dtypes, and in that case len(aval) should perhaps be
like len(aval.shape). The only place where this behavior was relied on
was sparse/ops.py.
* trailing-whitespace
* dangerous-default-value. None of these appear to be bugs in practice, but the potential for accidentally mutating the default value is there, and the cost of avoiding the problem is small.
* invalid-envvar-default. Pass strings as getenv() defaults.
* unnecessary-semicolon. Use tuples instead for this one-liner.
* invalid-hash-returned. Raise an exception rather than asserting false.
* pointless-string-statement. Use comments instead.
* unreachable. Use @unittest.skip() decorator rather than raising as first line in test.
* logging-not-lazy. Make the logging lazy.
* bad-format-string-type. Use f-string instead.
* subprocess-run-check. Pass check=...
PiperOrigin-RevId: 400858477
All control-flow primitives are `AxisPrimitive`s now, which means that we're doing
lots of those used-names traversals during dispatch and they can be expensive!
This adds caching to try to lower the cost of that change.
PiperOrigin-RevId: 395921887
Even though `vmap` and `pmap` don't use avals with names, the batching infrastructure
is used to implement xmap and pjit. So while we keep the introduction of names carefully
scoped, forgetting to remove them at the right points leads to extremely confusing errors.
PiperOrigin-RevId: 395423006
I've also updated the docs for ``jax.ops`` to note that ``at[].set()``
is guaranteed to be performed in-place under JIT. Someone who knows XLA
well should double check that fact!
The string representation of ConcreteArray did not include the data type of the
wrapped value. This makes it harder to spot the reason for errors arising from
inconsistent values (issue #5364). This commit adds the data type to the string
representation of ConcreteArray.
