Searched refs:rng_op (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | decompose_resource_ops.cc | 106 auto rng_op = cast<RngReadAndSkipOp>(op); in matchAndRewrite() local
109 if (!matchPattern(rng_op.alg(), m_Constant(&alg_constant))) { in matchAndRewrite()
132 if (op_type != rng_op.getType()) { in matchAndRewrite()
136 if (!HasResourceSubtype(rng_op.resource())) { in matchAndRewrite()
144 if (res_type != GetResourceSubtype(rng_op.resource())) { in matchAndRewrite()
152 rewriter.create<ReadVariableOp>(loc, res_type, rng_op.resource()); in matchAndRewrite()
170 rewriter.create<MulOp>(loc, u64_scalar, step_size, rng_op.delta()); in matchAndRewrite()
191 rewriter.create<AssignVariableOp>(loc, rng_op.resource(), new_state); in matchAndRewrite()
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/ |
| D | hlo_function_importer.cc | 607 auto rng_op = Cast<HloRngBitGeneratorInstruction>(instruction); in ImportInstructionImpl() local
610 func_builder->getI32IntegerAttr(rng_op->algorithm()), operands[0]); in ImportInstructionImpl()
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | ir_emitter_unnested.cc | 2462 auto rng_op = in EmitRngGetAndUpdateState() local
2469 BuildKernelThunkForMlir(rng_op, rng_op.state(), mlir_input.thunk_info, in EmitRngGetAndUpdateState()
2474 llvm_ir::RngGetAndUpdateState(rng_op.delta(), module_, &b_); in EmitRngGetAndUpdateState()
2476 const Shape shape = TypeToShape(rng_op.state().getType()); in EmitRngGetAndUpdateState()
|