Searched refs:logits_type (Results 1 – 1 of 1) sorted by relevance
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | softmax_op.cc | 203 const DataType logits_type = input_type(0); in Compile() local
214 indices, XlaHelpers::One(builder, logits_type), in Compile()
215 XlaHelpers::Zero(builder, logits_type), &labels)); in Compile()
227 xla::Broadcast(XlaHelpers::Zero(builder, logits_type), {batch_size}), in Compile()
228 xla::Broadcast(XlaHelpers::FloatLiteral(builder, logits_type, NAN), in Compile()
234 ctx, logits_type, xla_logits_type, ctx->Input(0), labels); in Compile()
|