(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]])
| 65 | |
| 66 | @register_lower_rule(mops.Dropout) |
| 67 | def dropout_lower(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]]): |
| 68 | assert len(ctx.vars_in) == 2 and len(args) == 2 and len(ctx.vars_out) == 3 |
| 69 | inp, key = args |
| 70 | random_val, new_key = rng_uint_generator(key, inp.shape, "uint32") |
| 71 | mask = random_val > np.array(ctx.op.drop_prob * np.iinfo(np.uint32).max, np.uint32) |
| 72 | multiplier = mask.astype(inp.dtype) |
| 73 | assert ctx.op.drop_prob < 1 and ctx.op.drop_prob >= 0, ctx.op.drop_prob |
| 74 | multiplier = multiplier / (1.0 - ctx.op.drop_prob) |
| 75 | out = (inp * multiplier).astype(inp.dtype) |
| 76 | mask = mask.reshape((-1,)).astype("uint8") |
| 77 | return out, mask, new_key |
| 78 | |
| 79 | |
| 80 | @register_lower_rule("DropoutBackward") |
nothing calls this directly
no test coverage detected