Searched refs:shifted_logits (Results 1 – 3 of 3) sorted by relevance
63 auto shifted_logits = xla::Sub(logits, logits_max, batch_dims); in Compile() local64 auto exp_shifted = xla::Exp(shifted_logits); in Compile()78 ? xla::Sub(shifted_logits, xla::Log(sum), batch_dims) in Compile()106 auto shifted_logits = xla::Sub(logits, logits_max, {kBatchDim}); in CrossEntropyWithLogits() local109 auto exp_shifted_logits = xla::Exp(shifted_logits); in CrossEntropyWithLogits()127 auto sub = xla::Sub(shifted_logits, log_sum_exp, {kBatchDim}); in CrossEntropyWithLogits()
66 auto shifted_logits = (logits - logits.maximum(along_class) in Compute() local73 softmax.device(d) = shifted_logits; in Compute()86 softmax.device(d) = shifted_logits.exp(); in Compute()
495 shifted_logits = array_ops.concat([504 logits=shifted_logits,