tensorflow / tensorflow Public

PiperOrigin-RevId: 420878608
Change-Id: I2d81829c56c3a24c2f547c420c25cba5b9a2b778

…tribute _xla_outside_compilation that's reachable in the call graph from a TPU cluster will be outside-compiled, just like an op that's directly in the TPU cluster.

Traversal over ops reachable from a TPU cluster is done in 3 places in this cl, so it's refactored into tpu_cluster_util.

PiperOrigin-RevId: 420875110
Change-Id: I81d51ea27d8d7d6f9322b0105445b1dc3f1abadf

PiperOrigin-RevId: 420874938
Change-Id: I3a4c640e9a0ef1962533a79171a3124768fc205c

PiperOrigin-RevId: 420871518
Change-Id: Iac67856b983427a89837a8eaf8a4a5268c182633

So far we got lucky and tensor.extract only occurs on shapes, but there's no
guarantee for that.

PiperOrigin-RevId: 420859265
Change-Id: Ib48fb3e77a8e2cad7f8f380dff8f5ee031f47775

When constructed via one path we were filling in the .custom field of nodes that weren't custom types.

Fixes google/jax#9066

PiperOrigin-RevId: 420858917
Change-Id: I5a3df78811551331e502f8c2940cd5b59dcbf2ce

The proximate reason for the upgrade is to pick up a thread-safety fix (pybind/pybind11#3237) for JAX, but we might as well just upgrade to the current version for everyone.

Remove a backport of a patch now present in the released version (pybind/pybind11#2864).

Patch out the "VERSION" file from nsync to avoid conflicts with the <version> C++ header included by pybind11.

PiperOrigin-RevId: 420856622
Change-Id: Ied007ed62b93d96cf008d6ec6336017af1fa10fd

PiperOrigin-RevId: 420853998
Change-Id: I1f1a48a2b827f817267538e49072a30d4fe16413

…omes dead due to parent computations being removed.

PiperOrigin-RevId: 420845681
Change-Id: I4f09f46c87d23489233e9c96b3f4dbf5ce6fb300

http://github.com/tensorflow/runtime/commit/8997563a6971829a5ae837177a64c4564bb94f24.

PiperOrigin-RevId: 420838341
Change-Id: I99718f44f64e0b861a149c19b2682aa0e2e87eb1

Updates LLVM usage to match
[4e77868c7c4b](llvm/llvm-project@4e77868)

PiperOrigin-RevId: 420835389
Change-Id: Ia2649089663de28782936eeeaa693feb3b3975e7

Also do minor cleanups along the way.

PiperOrigin-RevId: 420833280
Change-Id: Idb4929c39a6b9fa73cf3f50107fe0d9f5e91d434

PiperOrigin-RevId: 420830075
Change-Id: Iae45574925e1c0f1c2ae91e1adddfd0ac9be65e4

PiperOrigin-RevId: 420827501
Change-Id: I26d054e26dd095ba3f7cc55ddefe77b49a240462

PiperOrigin-RevId: 420822924
Change-Id: I868504af6fa7d0f01a849fc7821cf6b28aa2833f

PiperOrigin-RevId: 420816384
Change-Id: Ic077ff56d6a3635deb3c76ee49f1bd358c3b091a

PiperOrigin-RevId: 420815498
Change-Id: If90a8370849dce96ce2174ca3fa6e189ee913382

PiperOrigin-RevId: 420811825
Change-Id: Ib5b1c1c8a9c446b0fc9f64e7a0d4ad78ad956412

PiperOrigin-RevId: 420811652
Change-Id: I83742482770ba0bf7c3ccd57508c40fb9cdbe2f7

http://github.com/tensorflow/runtime/commit/216e913a2b523748895e97aac2df42727a760d29.

PiperOrigin-RevId: 420806092
Change-Id: I5ba2c90970cff11b60e262eb002729aa95e36777

PiperOrigin-RevId: 420805134
Change-Id: I0d6c2dd288ad9701a9308e3fc9cabafa93304449

…kend

Move cpurt-related code from tfrt/backends/cpu/jit to a separate `jitrt` backend (tfrt/backends/jitrt). Namespace renaming + types renaming will be in the next CLs.

PiperOrigin-RevId: 420803251
Change-Id: Ifd8c7e5a8c912f558f058e83921636945c6e653b

PiperOrigin-RevId: 420802498
Change-Id: I7f46df54099cb04ff8ad40b0294d77cde2ecd7fd

* For single-page changes, tag an object using `doc_controls.set_custom_page_builder_c lass(obj, cls)
* For global changes pass a dict of `{ObjectType:Type[PageInfo]}` to the DocGenerator's `page_builder_classes` argument.
* Switch generate2 to use the new customization pathways.

PiperOrigin-RevId: 420801142
Change-Id: I57d02ffa4dd439ddd63578b76109ca5794d4d8da

This can be used to override the system default device when running a
jit'd function. This is the prioritized list which determines which
settings take precedence when determining which device a jit'd
function runs on:

1. `device` or `backend` argument to jit
2. Input(s) with sticky device
3. default_device config (new!)
4. System default device

This also changes the `jit` C++ binding to take a `jit_device`
argument, which is the device determined by the `device` or `backend`
argument to jit (or None if not set), rather than determining it by
calling back into Python.

PiperOrigin-RevId: 420780860
Change-Id: I0f52fbc4fd577829997057a0b261442edd5ea62f

PiperOrigin-RevId: 420778896
Change-Id: Ic2bbb480c86a16a9986637b023386c0dc72f2908

…st-fold-helpers

PiperOrigin-RevId: 420774825
Change-Id: I2e0589dd3553bc6c5181c4e380fdf15c0145a8cf

http://github.com/tensorflow/runtime/commit/f57c209a2fe25dab87bcf1e69783dc2b8942418f.

PiperOrigin-RevId: 420756033
Change-Id: Ib1d512337c1347a9f2a9942f1a111286c0723c95

Before this change, JAX produces HLO using the XLA:Python builder APIs. After this change JAX produces MHLO using MLIR:Python APIs, and converts the MHLO to HLO for compilation with XLA. This is a lateral shift that should have little immediate impact, but unlocks a number of interesting opportunities in the future (e.g., mixing MLIR dialects within a JAX program).

[XLA:Python] Pass MLIR input as a std::string to work around pybind/pybind11#2765. A better fix would be to update pybind11 but that is hitting Windows-related hurdles; for now, just avoid relying on reference lifetime extension.

Brax: update test seeds to avoid test failures. Additional constant folding (canonicalization) in the MHLO lowering path seems to cause small numerical differences.
PiperOrigin-RevId: 420755696
Change-Id: I5e2626ea1e82c046a847300bf6bbe94208007802

Updates LLVM usage to match
[ad1b8772cf6b](llvm/llvm-project@ad1b877)

PiperOrigin-RevId: 420753473
Change-Id: I926d1abff7ae666fe207002873cf360e07b53a65

PiperOrigin-RevId: 420753231
Change-Id: Ia5fdb53c2424d255a8769179ec52b4674b0f0358

This provides c-interface related functionality that will be lowered to
match the XLA:CPU calling convention without imposing on the other mhlo IR.

PiperOrigin-RevId: 420748384
Change-Id: I3413a072a6990d7c167d384d5bfe21eecee79ccd

PiperOrigin-RevId: 420735548
Change-Id: Ic4f86961e09651c7190748251aa29a53d95880d6

PiperOrigin-RevId: 420722273
Change-Id: I8963b7195d3fe2fc49f51f982ac1eac96abf0296