Home
last modified time | relevance | path

Searched refs:RngOutput (Results 1 – 6 of 6) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dprng.h27 struct RngOutput { struct
44 using BitGeneratorTy = std::function<RngOutput(XlaOp key, XlaOp initial_state, argument
50 RngOutput ThreeFryBitGenerator(XlaOp key, XlaOp initial_state,
61 RngOutput PhiloxBitGenerator(XlaOp key, XlaOp initial_state,
70 RngOutput UniformFloatingPointDistribution(XlaOp key, XlaOp initial_state,
77 RngOutput UniformIntDistribution(XlaOp key, XlaOp initial_state,
84 RngOutput NormalFloatingPointDistribution(XlaOp key, XlaOp initial_state,
Dprng.cc243 RngOutput ThreeFryRngBit32(XlaOp key, XlaOp initial_state, const Shape& shape) { in ThreeFryRngBit32()
255 RngOutput ThreeFryRngBit64(XlaOp key, XlaOp initial_state, const Shape& shape) { in ThreeFryRngBit64()
407 RngOutput PhiloxRngBit32(XlaOp op_key, XlaOp initial_state, in PhiloxRngBit32()
434 RngOutput PhiloxRngBit64(XlaOp op_key, XlaOp initial_state, in PhiloxRngBit64()
532 RngOutput ThreeFryBitGenerator(XlaOp key, XlaOp initial_state, in ThreeFryBitGenerator()
553 RngOutput PhiloxBitGenerator(XlaOp key, XlaOp initial_state, in PhiloxBitGenerator()
581 RngOutput UniformFloatingPointDistribution(XlaOp key, XlaOp initial_state, in UniformFloatingPointDistribution()
585 RngOutput bits_state = bit_generator(key, initial_state, shape); in UniformFloatingPointDistribution()
592 RngOutput UniformIntDistribution(XlaOp key, XlaOp initial_state, in UniformIntDistribution()
595 RngOutput bits_state = bit_generator(key, initial_state, shape); in UniformIntDistribution()
[all …]
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dstateful_random_ops.cc51 return xla::RngOutput{data, new_state}; in BitGen()
61 return xla::RngOutput{data, new_state}; in BitGen()
66 xla::RngOutput StatefulRngUniform(Algorithm alg, xla::XlaOp key, in StatefulRngUniform()
92 xla::RngOutput StatefulRngUniformFullInt(Algorithm alg, xla::XlaOp key, in StatefulRngUniformFullInt()
96 xla::RngOutput output = BitGen(alg)(key, initial_state, shape); in StatefulRngUniformFullInt()
115 using SamplerReturnType = xla::StatusOr<xla::RngOutput>;
203 xla::RngOutput value_state = status_or_value.ConsumeValueOrDie(); in CompileImpl()
231 xla::RngOutput uniform_state = StatefulRngUniform( in Compile()
274 xla::RngOutput value_state = xla::NormalFloatingPointDistribution( in Compile()
315 xla::RngOutput uniform_result = StatefulRngUniform( in Compile()
Dstateless_random_ops_v2.cc70 xla::RngOutput BitGenerator(xla::RandomAlgorithm const& alg, xla::XlaOp key, in BitGenerator()
78 return xla::RngOutput{/*value=*/xla::GetTupleElement(result, 1), in BitGenerator()
111 xla::RngOutput StatelessRngUniformV2(xla::RandomAlgorithm const& alg, in StatelessRngUniformV2()
145 xla::RngOutput StatelessRngUniformFullInt(xla::RandomAlgorithm const& alg, in StatelessRngUniformFullInt()
151 xla::RngOutput output = BitGenerator(alg, key, counter, shape); in StatelessRngUniformFullInt()
158 return xla::RngOutput{BitcastConvertType(output.value, type), in StatelessRngUniformFullInt()
Dstateless_random_ops.cc51 return xla::RngOutput{/*value=*/xla::GetTupleElement(result, 1), in GetBitGeneratorForDevice()
59 return xla::RngOutput{/*value=*/xla::GetTupleElement(result, 1), in GetBitGeneratorForDevice()
124 xla::RngOutput output = in StatelessRngUniformFullInt()
/external/tensorflow/tensorflow/compiler/xla/service/
Drng_bit_generator_expander.cc68 RngOutput output; in GetGeneratorComputation()