Automated g4 rollback of changelist 195091587

PiperOrigin-RevId: 195184798

Add complex128 support to FFT, FFT2D, FFT3D, IFFT, IFFT2D, and IFFT3D.

NumPy automatically upcasts to complex128 when computing FFTs, leading to issues like:
#10749

This change allows users to choose between 32-bit and 64-bit precision FFTs on CPU and GPU.

PiperOrigin-RevId: 195183206

Add a collect_trace option to run_op_benchmark for cases when callers just want
to pass RunOptions.FULL_TRACE but don't want to store trace in extras.

PiperOrigin-RevId: 195181533

BufferValue is a new base class for LogicalBuffer and HloValue. This makes it
easier to migrate from TuplePointsToAnalysis/LogicalBuffer to
HloDataflowAnalysis/HloValue. No functional changes.

PiperOrigin-RevId: 195179676

Handle negative values when slicing symbolic shapes

PiperOrigin-RevId: 195176133

Allow `Layer.add_loss` to receive non-tensor; fixes error triggered when using a weight regularizer of factor 0.

PiperOrigin-RevId: 195175909

Enable reshape of _ScopedAllocatorConcat output.

The _ScopedAllocatorConcat kernel outputs the backing tensor after performing
runtime bounds checks.  However, the shape of the backing tensor may not match
the desired output shape of the concat operation.

This change adds a "reshape" boolean attribute to _ScopedAllocatorConcat kernel.
When this attribute is set to true, the kernel outputs a reshaped backing tensor
according to the "shape" attribute.

PiperOrigin-RevId: 195169105

Add RNNEstimator which takes in arbitrary heads.

PiperOrigin-RevId: 195167853

Mark all nodes processed by AddOpsRewrite/MinBCast stages with a tag.

PiperOrigin-RevId: 195167597

Allow evaluation and prediction through warm-starting (no current checkpoint / model_dir).

PiperOrigin-RevId: 195165732

Adds the EvalListener support for run_local.

PiperOrigin-RevId: 195163507

Adding a version of rolled triangular solver code for the right-multiply case, which is used in Cholesky decomposition.  Replacing the unrolled version with a While loop drastically reduces XLA compilation times which allows much larger models to be run on TPU.

PiperOrigin-RevId: 195163298

Add prefetching to one device distribution strategy.

PiperOrigin-RevId: 195162570

Replaced calls to tensorflow::StringPiece::ToString with std::string conversions.
That is, instances of sp.ToString() are replaced with std::string(sp).

This will allow tensorflow::StringPiece::ToString to be removed, which is necessary before it can be replaced with absl::string_view.

PiperOrigin-RevId: 195162126

Optimize LogicalOr and LogicalAnd with all true or false inputs:

LogicalOr(x, true) = true
LogicalOr(x, false) = x
LogicalAnd(x, true) = x
LogicalAnd(x, false) = false

and similar if the first argument is constant.

PiperOrigin-RevId: 195161140

[XLA] BF16 propagation: do not change if propagation is confined inside a fusion.

We now use a set to track all the potential changes, and do the actual changes
on the HLOs at the end. This also makes the boolean return value (whether
anything is changed) correct.

PiperOrigin-RevId: 195160025

[TF:XLA] Bump open source llvm revision to r331338

PiperOrigin-RevId: 195158710

Add steps_per_run to LoggingTensorHook and StepCounterHook and other logging bug fixes.

PiperOrigin-RevId: 195158238

Fix support for batch_normalization with mixed precision

When the type of the input tensor `x` is not the same as the type of
the parameters `mean`, `variance`, `offset`, and `scale`, a cast is
required.

This mixed precision case occurs when using the BatchNormalization
layer with a data type of float16 or bfloat16.

PiperOrigin-RevId: 195157279

Updated ABSL to latest version in workspace.bzl.

PiperOrigin-RevId: 195144612

Fix a bug in create_python_api.py

I got an error complaining about "RuntimeError: dictionary changed size during iteration", this change fixes it.

PiperOrigin-RevId: 195144333

[XLA] Add new optimization that sinks constants into while loop bodies

Example transformation:

   state = (..., const, ...)
   while (pred(state)) {
     (..., v, ...) = state
     use(v)
     state = (..., v, ...)
   }

 =>

   state = (..., const, ...)
   while (pred(state)) {
     (..., v, ...) = state
     use(const)
     state = (..., v, ...)
   }

PiperOrigin-RevId: 195143323

Re-enabling a test.

PiperOrigin-RevId: 195142105

Fix tsan failure in batch_dataset_op_test.
The error was being caused because we were trying to save invocation_results_` while the function call was in progress. Now we wait for all invocations to finish before saving both `invocation_results_` and `batch_results_`.

Did local A/B testing for tsan.
Before: 12/100 failed
After: All passed
PiperOrigin-RevId: 195141349

Increasing test size to reflect recent additions and prevent test timeouts.

PiperOrigin-RevId: 195138565

[XLA] Redesign: Dump HloSnapshot in local service as well. And support replaying HloSnapshot.

PiperOrigin-RevId: 195138472

Added support for packing of symbolic shapes

PiperOrigin-RevId: 195137239

docs: Link to the appropriately branched version of the live colab notebooks.

And update that link on release changes.

PiperOrigin-RevId: 195133689

Turn on two half precision tests for GPU.

PiperOrigin-RevId: 195128326

Replaced calls to tensorflow::StringPiece::ToString with std::string conversions.
That is, instances of sp.ToString() are replaced with std::string(sp).

This will allow tensorflow::StringPiece::ToString to be removed, which is necessary before it can be replaced with absl::string_view.

PiperOrigin-RevId: 195126422

Internal-only change.

PiperOrigin-RevId: 195125476

[XLA:GPU] Unroll unfused elementwise op kernels.

So far we only unrolled loop fusions, elementwise ops is a logical extension.
We don't spend a lot of time in unfused elementwise ops in benchmarks, so this
is only worth a small speedup on V100.

PiperOrigin-RevId: 195121530

Automated g4 rollback of changelist 194981511

PiperOrigin-RevId: 195120627

Renames _regression_head_with_mean_squared_error_loss to _regression_head.

PiperOrigin-RevId: 195117425

Initialize all members of CollectiveParams at construction time to avoid
warnings of uninit memory access by dynamic analysis tools.

PiperOrigin-RevId: 195117321

Fix some nits in cpu_literal_caching_test that I noticed after submission

PiperOrigin-RevId: 195114829

Internal-only change.

PiperOrigin-RevId: 195113702

Optimized the analysis of rank and size operations.

PiperOrigin-RevId: 195108832

Instantiate SwapDimension1And2InTensor3 for Eigen::half

PiperOrigin-RevId: 195107839

Add Name String to GraphOptimizationPass and Log Registered Passes

Added a name string to GraphOptimization class and set to the class name through REGISTER_OPTIMIZATION macro.
Modified RunGrouping function to log the name and phase of optimization pass that's running.
Added two additional functions to log all registered optimization passes in the order of execution.

PiperOrigin-RevId: 195106355

Automated g4 rollback of changelist 195091587

PiperOrigin-RevId: 195098224

Use experimental auto_sharding in multi worker dataset.

PiperOrigin-RevId: 195092992

ArraysExtraInfo: Add name_regexp field and regexp name matching.
PiperOrigin-RevId: 195091587

BUGFIX: Convert inputs and list of gradients into tuple if they are not instance of tuple. Otherwise this causes "unhashable keys" error when we try to hash.
Also fixed lint error.

PiperOrigin-RevId: 195061425

Add missing colocated element in test in buffer_assignment_test.
This was resulting in a flaky test because sometimes the live set would include this missing colocated element perhaps because the buffers in the allocation has some nondeterministic order (read from a map?).

PiperOrigin-RevId: 195039311

Minor refactor: establish some operator naming conventions and apply them, so that the interface is a bit more consistent.

PiperOrigin-RevId: 195034691

Make the CRF work when sequence_lengths are int32.

PiperOrigin-RevId: 195034218

[XLA:CPU] Re-use the same llvm::GlobalVariable for identical literals

This isn't necessary today, but it will be after an optimization change I'm
about to make.

LLVM has a constant merging pass too, but one of the motivations here is to
avoid the LLVM compile time overhead of having many large arrays in the IR.

PiperOrigin-RevId: 195032900

Internal change

PiperOrigin-RevId: 195028221

Internal change.

PiperOrigin-RevId: 195027918

Re-apply CL 194140820, which reverts #18251 (convolution change).

PiperOrigin-RevId: 195027049

test fix

PiperOrigin-RevId: 195023740

Sharding for tensorflow/contrib/timeseries/python/timeseries/state_space_models:structural_ensemble_test

PiperOrigin-RevId: 195019968

Fix wrongly ordered lines

PiperOrigin-RevId: 195019769

Avoid making a copy of the graph needlessly

PiperOrigin-RevId: 195017837

Adds logistic_regression_head.

PiperOrigin-RevId: 195017830

Check for overflow in shape calculation.

PiperOrigin-RevId: 195017114

Allow `warm_start_from` argument to be a SavedModel path.

PiperOrigin-RevId: 195015356

[TF:XLA] Separate on-host and on-device shape and layout in HloModule.

Previously, only one layout was stored with an HLO module. This CL allows
HLO passes to modify the on-device layouts without affecting the on-host
layout (provided by the client)

PiperOrigin-RevId: 195014875

Implementation of the fully-connected TFLite Op using the symmetric quantization.

PiperOrigin-RevId: 195010312

Go: Update generated wrapper functions for TensorFlow ops.
PiperOrigin-RevId: 195010310

Fixed some outdated comments

PiperOrigin-RevId: 195006088

Minor JNI performance improvement.

PiperOrigin-RevId: 195002949

Making ids unique in nn.embedding_lookup_sparse. This helps to reduce RPC calls for looking up the embeddings when there are repeated ids in the batch.

PiperOrigin-RevId: 195002785

Fix bug in peak buffer accounting in buffer assignment.
Buffer assignment keeps track of the set of logical buffers which are live at
the point of peak memory usage for each allocation. Previously colocated
buffers were not properly accounted for. This CL addresses this problem.

PiperOrigin-RevId: 195001567

Merge changes from github.

PiperOrigin-RevId: 194997009

Make tower-local variables non-trainable even with the default
DistributionStrategy.

PiperOrigin-RevId: 194996819

Enable checkpointless eval and predict for tf.estimator.

PiperOrigin-RevId: 194993191

Update community/swift

PiperOrigin-RevId: 194991305

Update schema.

PiperOrigin-RevId: 194989704

Relax the stringent memory allocator constraints in AssignOp if a Grappler graph analysis determines it to be safe. This will allow Assign to reuse the input buffer to initialize the variable in many cases.

PiperOrigin-RevId: 194988134

Collective Ops Part 5

Distributed-mode implementations of DeviceResolverInterface
and ParamResolverInterface.  Extend Worker interface with
new methods in support of these interfaces.

This change is part of a series of changes introducing infrastructure
for collective ops and initial implementations of reduction and broadcast.

PiperOrigin-RevId: 194984585

Open source infeed test

PiperOrigin-RevId: 194983270

Add utility to auto shard a dataset pipeline in the appropriate place by locating the file readers and sharding their input files.

PiperOrigin-RevId: 194982311

Add a pointer from Device to its owning DeviceMgr.
Allow remote function execution on TPU devices.

PiperOrigin-RevId: 194981511

Implements matrix multiply-accumulate for linear no-offset (aka symmetric) quantizer.

PiperOrigin-RevId: 194978865

Add device_util.resolve method which merges with current device as well.

PiperOrigin-RevId: 194976633

Simplified shape inference.

PiperOrigin-RevId: 194975603

Improve shape inference for tf.contrib.signal.frame.

PiperOrigin-RevId: 194972934

Boosted trees: support indicator column.

PiperOrigin-RevId: 194971229

eager: Update sample notebooks with API changes in the last few releases.

Most notably:
- Avoid using tf.contrib.eager since equivalent functionality if available outside tf.contrib
- Datasets can be directly iterated on.
- Use tf.GradientTape instead of tf.contrib.eager.implicit_gradients

PiperOrigin-RevId: 194971115

Automated g4 rollback of changelist 194917415

PiperOrigin-RevId: 194962702

Fix crash in HloGraphDumper where it crashes on tuple shaped constants

The problem is that it tries to use a special logic for 0 element constants
but the logic used to check the number of elements only supports array shapes.

PiperOrigin-RevId: 194945246

Preventing RemoveTrivialBinary from removing broadcasts.

PiperOrigin-RevId: 194937001

Protocol buffer classes now list their fields in dir(cls)

PiperOrigin-RevId: 194917415

[XLA:CPU] Open source some tests.

PiperOrigin-RevId: 194903752

Update ops-related pbtxt files.

PiperOrigin-RevId: 194883351

Move LinearOperatorKronecker and LinearOperatorBlockDiag to core.

PiperOrigin-RevId: 194881237

[XLA] Redesign: dump HloSnapshot at the point where it used to dump the SessionModule.

PiperOrigin-RevId: 194880385

Internal change.

PiperOrigin-RevId: 194877173

Remove proto header import from core/framework/tracking_allocator.h

The goal is to make kernels mostly independent of proto headers, which will let us lock down our .so import.

PiperOrigin-RevId: 194876569

Go: Update generated wrapper functions for TensorFlow ops.
PiperOrigin-RevId: 194874988

[XLA] Redesign: migrate tensorflow/compiler/tf2xla, tensorflow/compiler/aot:
- xla::ComputationBuilder -> xla::XlaBuilder
- xla::ComputationDataHandle -> xla::XlaOp
- xla::Computation -> xla::XlaComputation
- xla::CompileOnlyClient::AotComputationInstance -> xla::CompileOnlyClient::AotXlaComputationInstance
- xla::SessionModule -> xla::HloSnapshot

PiperOrigin-RevId: 194874462

[tf.data] Adding an experimental `group_by_reducer` transformation which groups elements of an input pipeline by a key, applies a reduce function to elements of each group "on-the-fly", and outputs the results once all input elements have been processed.

PiperOrigin-RevId: 194874087

Internal cleanup.

PiperOrigin-RevId: 194871141

Faster reduce_logsoftmax (specially in eager) and bugfixes in broadcast_to

PiperOrigin-RevId: 194870645

Small fix to prevent a crash if the delegate has not implemented FreeBufferHandle.

PiperOrigin-RevId: 194866595

Add MultiNodeDataset and MultiNodeIterator which are intended to work for multi-node distribution strategy.

PiperOrigin-RevId: 194862215

[XLA] Change the TF2XLA bridge to perform F16 reduction using F32 data type.

Add test cases to test that reduce sum for bfloat16 and float16 doesn't lose too
much precision.

PiperOrigin-RevId: 194862078

Improve error message for pip_smoke_test.

PiperOrigin-RevId: 194859591