PiperOrigin-RevId: 564578779

PiperOrigin-RevId: 564577405

These patterns are used only by TFLite which now has direct support for stablehlo.scatter.

PiperOrigin-RevId: 564569917

PiperOrigin-RevId: 564559805

Sometimes scales are zero in the all zero case. This is obviously
silly but it triggers a failure in the pipeline. Make scale the
minimum scale value as whatevs.

PiperOrigin-RevId: 564556649

PiperOrigin-RevId: 564551916

Add the custom call xla.gpu.init_cuBLAS to the beginning of the entry function of the module, when the module contains a gemm.

PiperOrigin-RevId: 564550744

The user will provide inputs via AHWB. Behind the scenes we bind the AHWB's to SSBO. Both cl & gl inference are supported.

PiperOrigin-RevId: 564547199

PiperOrigin-RevId: 564546615

PiperOrigin-RevId: 564546472

We are not going to pursue direct lowering from LMHLO to API calls, though we might consider "reference implementation API" in the future, but it will go through intermediate "XLA runtime dialect". For now remove dead code.

PiperOrigin-RevId: 564538039

Leaky Relu, Elu, Relu, Relu1, Relu0To1, Relu6, Tanh, Sigmoid, and HardSwish

PiperOrigin-RevId: 564537536

PiperOrigin-RevId: 564529533

PiperOrigin-RevId: 564528951

Needed for sharding int4 models.

PiperOrigin-RevId: 564528408

PiperOrigin-RevId: 564526562

PiperOrigin-RevId: 564526383

PiperOrigin-RevId: 564518057

PiperOrigin-RevId: 564516421

Rearrange the code that unroll the last while-loop iteration to avoid cloning
the instruction being pipelined, which is unused in the iteration. This is to
prepare for the next change to support backward pipelining recv-done without
generating the unused recv and recv-done pair.

PiperOrigin-RevId: 564515742

1. Adding bidirectional propagation paths among nodes within the same shard_group
2. Enforce the shardings from the same shard_as group to take the same sharding in the end.

PiperOrigin-RevId: 564504486

PiperOrigin-RevId: 564500466

PiperOrigin-RevId: 564494782

According to
https://www.tensorflow.org/api_docs/python/tf/math/unsorted_segment_sum:

> If the given segment ID i is negative, the value is dropped and will not be
> added to the sum of the segment.

PiperOrigin-RevId: 564485570

PiperOrigin-RevId: 564472190

Leaky Relu, Elu, Relu, Relu1, Relu0To1, Relu6, Tanh, Sigmoid, and HardSwish

PiperOrigin-RevId: 564472058

[SPMD] Create a placeholder sharding of type UNKNOWN with the lowest precedence that can be overwriten by any other shardings. (2/N)

UNKNOWN type is not supported in subgroup sharding.

PiperOrigin-RevId: 564456004

This way libtpu can use symbols from it.

PiperOrigin-RevId: 564454296

PiperOrigin-RevId: 564441419

Also
- Unify naming to ${scope}_${thing}_dim_idx.
- Make logic for split_lhs_nc a bit easier to follow (at least in my opinion), by
  decoupling it from the batch_size computation.

PiperOrigin-RevId: 564435916

PiperOrigin-RevId: 564434713

PiperOrigin-RevId: 564428072

[XLA] Use max size of buffer in heap simulator and buffer assignment and make size DCHECKs into CHECKs

PiperOrigin-RevId: 564421004

PiperOrigin-RevId: 564420754

PiperOrigin-RevId: 564415721

Sometimes, there's an exorbitant number of edges between two same notes A and B
(e.g. passing parameters through). After ~64 it becomes unreadable, so better
to skip others.

PiperOrigin-RevId: 564379502

different flag default.

PiperOrigin-RevId: 564374361

Updates LLVM usage to match
[e592c2dcf5b7](https://github.com/llvm/llvm-project/commit/e592c2dcf5b7)

PiperOrigin-RevId: 564346651

[Rolling forward]

PiperOrigin-RevId: 564345870

This makes it harder to accidentally use -1.

PiperOrigin-RevId: 564334287