HomeSort by relevance Sort by last modified time
    Searched defs:branch_index (Results 1 - 5 of 5) sorted by null

  /external/tensorflow/tensorflow/compiler/xla/service/
conditional_simplifier.cc 36 // computation. If the given conditional has a constant branch_index, tries to
54 int branch_index = 0; local
57 VLOG(2) << "Not attempting to remove conditional as its branch_index is "
64 branch_index = conditional->operand(0)->literal().Get<bool>({}) ? 0 : 1;
66 branch_index = conditional->operand(0)->literal().Get<int32>({});
67 if (branch_index < 0 || branch_index >= conditional->branch_count()) {
68 branch_index = conditional->branch_count() - 1;
75 conditional->shape(), {conditional->mutable_operand(branch_index + 1)},
76 conditional->branch_computation(branch_index)));
    [all...]
hlo_evaluator.cc 1264 int branch_index; local
    [all...]
  /external/tensorflow/tensorflow/compiler/xla/service/gpu/
conditional_thunk.cc 64 int32 branch_index = -1; local
71 stream->ThenMemcpy(&branch_index, branch_index_address, sizeof(int32));
77 "Failed to retrieve branch_index value on stream %p: %s.", stream,
81 branch_index = pred ? 0 : 1;
83 // Handle default scenario for branch_index not in [0, num_branches).
84 if (branch_index < 0 || branch_index >= hlo_instruction()->branch_count()) {
85 branch_index = hlo_instruction()->branch_count() - 1;
89 // Execute the branch computation corresponding to the value of branch_index.
91 TF_RETURN_IF_ERROR(branch_thunks_[branch_index]->ExecuteOnStream
    [all...]
  /external/tensorflow/tensorflow/core/kernels/
functional_ops.cc 243 const Tensor& branch_index = ctx->input(0); variable
244 OP_REQUIRES_ASYNC(ctx, TensorShapeUtils::IsScalar(branch_index.shape()),
245 errors::InvalidArgument("branch_index must be scalar"),
247 int32 branch = branch_index.scalar<int32>()();
317 Name("Case").Device(DEVICE_GPU).HostMemory("branch_index"), CaseOp);
  /external/tensorflow/tensorflow/compiler/xla/tests/
conditional_test.cc 197 XlaOp branch_index; local
199 &builder, &branch_index);
211 Conditional(branch_index, branches_p, operands);
240 XlaOp branch_index; local
242 &builder, &branch_index);
264 Conditional(branch_index, branches_p, operands);
415 XlaOp branch_index; local
417 CreateR0Parameter<int32>(bi, 0, "pred", &builder, &branch_index);
436 Conditional(branch_index, branches_p,
    [all...]

Completed in 224 milliseconds