absl-py==2.1.0 \
    --hash=sha256:526a04eadab8b4ee719ce68f204172ead1027549089702d99b9059f129ff1308 \
    --hash=sha256:7820790efbb316739cde8b4e19357243fc3608a152024288513dd968d7d959ff
    # via
    #   keras-nightly
    #   tb-nightly
astor==0.7.1 \
    --hash=sha256:95c30d87a6c2cf89aa628b87398466840f0ad8652f88eb173125a6df8533fb8d \
    --hash=sha256:fb503b9e2fdd05609fbf557b916b4a7824171203701660f0c55bbf5a7a68713e
    # via -r requirements.in

    return absl::OkStatus();
  }
  Status SetAttrFloat(const char* attr_name, float value) override {
    op_->node_builder.Attr(attr_name, value);
    return absl::OkStatus();
  }
  Status SetAttrBool(const char* attr_name, bool value) override {
    op_->node_builder.Attr(attr_name, value);
    return absl::OkStatus();
  }

  void Prune() ABSL_LOCKS_EXCLUDED(mu_);

  bool BlockNotStale(const std::shared_ptr<Block>& block)
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);

  /// Look up a Key in the block cache.
  std::shared_ptr<Block> Lookup(const Key& key) ABSL_LOCKS_EXCLUDED(mu_);

  void MaybeFetch(const Key& key, const std::shared_ptr<Block>& block,
                  TF_Status* status) ABSL_LOCKS_EXCLUDED(mu_);

        "//tensorflow/core/distributed_runtime/rpc/eager:grpc_eager_client",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/time",
        "@local_xla//xla/tsl/c:tsl_status_internal",
        "@local_xla//xla/tsl/distributed_runtime/coordination:coordination_service_agent",

                 absl::Span<AbstractTensorHandle* const> grad_outputs,
                 absl::Span<AbstractTensorHandle*> grad_inputs) override {
    // TODO(b/161805092): Support broadcasting.

    DCHECK(grad_outputs[0]);
    grad_inputs[0] = grad_outputs[0];
    grad_inputs[1] = grad_outputs[0];

    grad_inputs[0]->Ref();
    grad_inputs[1]->Ref();
    return absl::OkStatus();
  }

  // target element.
  Status ComputeGradient(
      AbstractContext* ctx, absl::Span<AbstractTensorHandle* const> targets,
      absl::Span<AbstractTensorHandle* const> sources,
      absl::Span<AbstractTensorHandle* const> output_gradients,
      absl::Span<AbstractTensorHandle*> result);
};

}  // namespace gradients
}  // namespace tensorflow

        input->DataType(), shape, &handle));
    params->emplace_back(handle);
  }
  return absl::OkStatus();
}

// Runs `model` maybe wrapped in a function.
Status RunModel(Model model, AbstractContext* ctx,
                absl::Span<AbstractTensorHandle* const> inputs,
                absl::Span<AbstractTensorHandle*> outputs, bool use_function) {
  if (use_function) {
    const char* fn_name = "test_fn";

                       absl::Span<AbstractTensorHandle* const> inputs,
                       absl::Span<AbstractTensorHandle*> outputs) {
  PartialTensorShape shape;
  TF_RETURN_IF_ERROR(inputs[0]->Shape(&shape));
  if (shape.dims() != 0) {
    return errors::InvalidArgument(
        "Tensor expected to have scalar shape found rank: ", shape.dims());
  }
  return absl::OkStatus();
}

Status MatMulModel(AbstractContext* ctx,
                   absl::Span<AbstractTensorHandle* const> inputs,
                   absl::Span<AbstractTensorHandle*> outputs) {
  return ops::MatMul(ctx, inputs[0], inputs[1], &outputs[0],
                     /*transpose_a=*/false,
                     /*transpose_b=*/false, "MatMul");
}

Status MulModel(AbstractContext* ctx,
                absl::Span<AbstractTensorHandle* const> inputs,

typedef struct GCSFile {
  const std::string path;
  const bool is_cache_enable;
  const uint64_t buffer_size;
  ReadFn read_fn;
  absl::Mutex buffer_mutex;
  uint64_t buffer_start ABSL_GUARDED_BY(buffer_mutex);
  bool buffer_end_is_past_eof ABSL_GUARDED_BY(buffer_mutex);
  std::string buffer ABSL_GUARDED_BY(buffer_mutex);

  GCSFile(std::string path, bool is_cache_enable, uint64_t buffer_size,
          ReadFn read_fn)