unicorn who dev
Complete FireViewer Android learning SDK, source and verification records
dee7f43 verified | """Keep detector proposal TopK as a native TFLite operation, not StableHLO sort.""" | |
| import torch | |
| from litert_torch import fx_infra | |
| from litert_torch.backend.lowerings import registry, utils | |
| from litert_torch.backend.experimental.torch_tfl._lowerings import _ir_operation | |
| from litert_converter.mlir import ir | |
| def native_topk(x: torch.Tensor, k: int) -> tuple[torch.Tensor, torch.Tensor]: | |
| return torch.topk(x, k, dim=-1, largest=True, sorted=True) | |
| def fake_topk(x, k): | |
| shape = (*x.shape[:-1], k) | |
| return x.new_empty(shape), x.new_empty(shape, dtype=torch.int64) | |
| def lower_topk(ctx, x, k): | |
| values, indices = utils.node_meta_to_ir_types(ctx.node) | |
| i32 = ir.IntegerType.get_signless(32) | |
| index_type = ir.RankedTensorType.get(indices.shape, i32) | |
| result = _ir_operation('tfl.topk_v2', results=[values, index_type], | |
| operands=[x, utils.splat(k, i32, [])]) | |
| return result[0], _ir_operation('tfl.cast', results=[indices], operands=[result[1]]) | |
| def topk(x, k, dim=-1, largest=True, sorted=True): | |
| dim = dim % x.ndim | |
| moved = x.movedim(dim, -1) | |
| values, indices = torch.ops.vds_static.topk(moved if largest else -moved, k) | |
| if not largest: | |
| values = -values | |
| return values.movedim(-1, dim), indices.movedim(-1, dim) | |
| fx_infra.decomp.update_pre_convert_decomp({torch.ops.aten.topk.default: topk}) | |
| fx_infra.decomp.update_pre_lower_decomp({torch.ops.aten.topk.default: topk}) | |