Out(NodeWith(
                     Op("Slice"),
                     Name("begin/static_shaped_slice/static_shaped_slice"))),
                 _)));
}

// New constants being created need to have control dependencies copied to
// ensure correct control flow analysis in TF V2.
TEST(SliceToDynamicSliceRewriteTest, WithControlDepsToConstant) {
  Scope root = Scope::NewRootScope()

      tensorflow::GetTypeFromTFTensorShape({1}, shape_dtype);
  auto expanded_start_index = rewriter->create<TF::ExpandDimsOp>(
      loc, vector_type, start_index, scalar_zero);
  auto start_position = rewriter->create<TF::ConcatOp>(
      loc, position_type, scalar_zero,
      ArrayRef<Value>({expanded_start_index, partial_start_position}));

  // Create the slice size tensor. This is done by concatenating `size` and
  // `partial_size`.

  void MarkAsResult(AbstractTensorHandle* gradient) const override;

  void DeleteGradient(AbstractTensorHandle* gradient) const override;

 private:
  // The context where the aggregation op `Add` is to be created.
  AbstractContext* ctx_;
};

// Returns the number of elements in the gradient tensor.
int64_t TapeVSpace::NumElements(AbstractTensorHandle* tensor) const {

  }
  void runOnOperation() override;
};

// Creates a call to function `func` in region `caller_region`. Use `args` as
// the call arguments, and terminate the region with a yield. The arguments are
// cast to the required type before the call. `use_region_args` control whether
// the input arguments are used as is (for IfOp) or block arguments of the same
// type as the input arguments are created and then used as call arguments (for
// While).

      builder.clone(*op, mapping);
    }
    // Creates the ReturnOp so that the per-host function returns the
    // correct values of the cloned operations.
    llvm::SmallVector<Value, 4> results_after_mapping;
    for (Value result : metadata.results) {
      results_after_mapping.push_back(mapping.lookupOrDefault(result));
    }
    builder.create<func::ReturnOp>(loc, results_after_mapping);

                options.shape_determination_fns) {
  if (options.shape_determination_fns.empty()) {
    LOG(ERROR) << "shape_representation_fns must be non-empty.";
  }
  VLOG(1) << "Created PJRT base device " << options.compilation_device_name
          << " device_name: " << name();
}

/*static*/ absl::StatusOr<const PjRtBaseDevice::Metadata*>
PjRtBaseDevice::GetMetadataFromDevice(DeviceBase* device) {

          !IsQintValueDefinedByIntToQintCast(op->getOperand(i))) {
        new_operands.push_back(rewriter.create<TF::CastOp>(
            op->getLoc(), orig_op_type, operands[i]));
      } else {
        new_operands.push_back(operands[i]);
      }
    }

    // Create a new UQ op.
    OperationState state(op->getLoc(), op->getName().getStringRef(),

// Caches a XlaDeviceAllocator per <backend, device ordinal> pair. A
// XlaDeviceAllocator is created on demand and is associated with a
// XlaDevice. It outlives the device itself (for instance, the buffer
// backing a tensor holds a pointer to the allocator for book-keeping,
// and this buffer can outlast the device).
class XlaDeviceAllocatorState {
 public:
  // Creates or returns a cached XlaDeviceAllocator for a given
  // backend and device_ordinal.

  llvm::SmallVector<mlir::TF::SplitOp, 4> split_ops_for_tiled_input;
  split_ops_for_tiled_input.reserve(
      input_sharding.tile_assignment_devices_size());

  // Creates a tree of split nodes for sharding tiled inputs. Splits nodes
  // are created such that input data is sharded in row major order.
  // Split nodes at ith depth from the original input node represent nodes
  // that split the input data at i-th dimension.

class TPUAnnotateDynamicShapeInputsPass
    : public impl::TPUAnnotateDynamicShapeInputsPassBase<
          TPUAnnotateDynamicShapeInputsPass> {
  void runOnOperation() override;
};

// Finds op that created a given value. If the value is a BlockArgument, this
// returns the owner of the Block.
Operation* GetOpOfValue(Value value) {
  if (auto block_arg = mlir::dyn_cast<BlockArgument>(value))