| 70 | } |
| 71 | |
| 72 | StatusOr<HloComputation*> GetComputationForRng(HloInstruction* rng) { |
| 73 | XlaBuilder builder("rng"); |
| 74 | const Shape u64_shape = ShapeUtil::MakeShape(xla::U64, {}); |
| 75 | const Shape u128_shape = ShapeUtil::MakeShape(xla::U64, {2}); |
| 76 | const Shape& result_shape = rng->shape(); |
| 77 | |
| 78 | XlaOp key = Parameter(&builder, 0, u64_shape, "key"); |
| 79 | XlaOp state = Parameter(&builder, 1, u128_shape, "state"); |
| 80 | XlaOp a_or_mean = |
| 81 | Parameter(&builder, 2, rng->operand(0)->shape(), "a_or_mean"); |
| 82 | XlaOp b_or_sigma = |
| 83 | Parameter(&builder, 3, rng->operand(1)->shape(), "b_or_sigma"); |
| 84 | |
| 85 | auto generator = [](xla::XlaOp key, xla::XlaOp state, |
| 86 | const xla::Shape& shape) { |
| 87 | return PhiloxBitGenerator(key, state, shape); |
| 88 | }; |
| 89 | |
| 90 | XlaOp result; |
| 91 | if (rng->random_distribution() == RNG_NORMAL) { |
| 92 | result = |
| 93 | NormalFloatingPointDistribution(key, state, generator, result_shape) |
| 94 | .value; |
| 95 | // Transform standard normal distribution to normal distribution with the |
| 96 | // given mean and standard deviation. |
| 97 | result = a_or_mean + (b_or_sigma * result); |
| 98 | } else { |
| 99 | CHECK_EQ(rng->random_distribution(), RNG_UNIFORM); |
| 100 | if (primitive_util::IsFloatingPointType(result_shape.element_type())) { |
| 101 | result = UniformFloatingPointDistribution( |
| 102 | key, state, generator, a_or_mean, b_or_sigma, result_shape) |
| 103 | .value; |
| 104 | } else { |
| 105 | result = UniformIntDistribution(key, state, generator, a_or_mean, |
| 106 | b_or_sigma, result_shape) |
| 107 | .value; |
| 108 | } |
| 109 | } |
| 110 | |
| 111 | TF_ASSIGN_OR_RETURN(XlaComputation xla_computation, builder.Build()); |
| 112 | |
| 113 | TF_ASSIGN_OR_RETURN(ProgramShape program_shape, |
| 114 | xla_computation.GetProgramShape()); |
| 115 | HloModuleConfig config(program_shape); |
| 116 | TF_ASSIGN_OR_RETURN(auto new_module, HloModule::CreateFromProto( |
| 117 | xla_computation.proto(), config)); |
| 118 | HloModule* module = rng->parent()->parent(); |
| 119 | HloCloneContext context(module); |
| 120 | return module->DeepCloneComputation(new_module->entry_computation(), |
| 121 | &context); |
| 122 | } |
| 123 | |
| 124 | } // namespace |
| 125 |
no test coverage detected