(
key, oshape, odtype="uint32", algorithm=RandomAlgorithm.RNG_THREE_FRY
)
| 26 | |
| 27 | |
| 28 | def rng_uint_generator( |
| 29 | key, oshape, odtype="uint32", algorithm=RandomAlgorithm.RNG_THREE_FRY |
| 30 | ): |
| 31 | |
| 32 | assert np.dtype(odtype) in { |
| 33 | np.dtype("uint8"), |
| 34 | np.dtype("uint16"), |
| 35 | np.dtype("uint32"), |
| 36 | np.dtype("uint64"), |
| 37 | }, f"only unsigned int supported, got {odtype}({type(odtype)})" |
| 38 | assert algorithm == RandomAlgorithm.RNG_THREE_FRY, "only ThreeFry supported now" |
| 39 | assert _shape_equal(key.shape, (2, 2)), f"key shape error, {key.shape}" |
| 40 | assert key.dtype == "int32", f"key dtype error, {key.dtype}" |
| 41 | |
| 42 | # bitcast (2x2,i32) -> (2,u64) |
| 43 | org_key_shape, org_key_dtype = key.shape, key.dtype |
| 44 | key = key.bitcast((2,), "uint64") |
| 45 | |
| 46 | if odtype == "uint32" or odtype == "uint64": |
| 47 | rng_odtype = odtype |
| 48 | else: |
| 49 | rng_odtype = "uint32" |
| 50 | |
| 51 | algorithm_attr = _rng_algorithm(algorithm) |
| 52 | new_key, out_vals = hlo.RngBitGeneratorOp( |
| 53 | ir_utils.make_ir_type_according_meta(key.shape, key.dtype), |
| 54 | ir_utils.make_ir_type_according_meta(oshape, rng_odtype), |
| 55 | algorithm_attr, |
| 56 | key.tensor, |
| 57 | ).results |
| 58 | new_key, out_vals = HLOTensor(new_key), HLOTensor(out_vals) |
| 59 | new_key = new_key.bitcast(org_key_shape, org_key_dtype) |
| 60 | |
| 61 | if rng_odtype != odtype: |
| 62 | out_vals = out_vals.astype(odtype) |
| 63 | return out_vals, new_key |
| 64 | |
| 65 | |
| 66 | @register_lower_rule(mops.Dropout) |
no test coverage detected