开启动态shape图分档执行功能时需要调用该接口设置图被划分的档位。

def set_dim_gears(t: torch.Tensor, dim_gears: Dict[int, Union[List[int], Tuple[int]]])

参数	输入/输出	说明	是否必选
t	输入	待分档的输入Tensor。	是
dim_gears	输入	用于设置Tensor不同dim维度下的档位值。输入类型为Dict，其中key为dim维度（整型），value为档位值（整数列表或元组）。说明：支持对同一个Tensor设置一个或者多个维度的档位。若图编译、执行时Tensor的shape不在设置的档位中，会导致编译或执行报错，请合理设置档位值。不支持对同一个tensor使用该接口设置两次不一样的档位。生成的总档位数量要不大于100。首次执行时对输入tensor设置档位即可，保证首次编译时能够获取到档位，后续执行时无需设置档位，避免因为设置档位的动作造成性能下降。	是

无

该功能仅适用于整图优化场景。
set_dim_gears接口优先级高于torch.compile中的dynamic=True或False，即对于Tensor类型的输入只有set_dim_gears的Tensor以及维度会被泛化。
该功能暂不支持同时配置Dynamo导图功能、使能RefData类型转换功能。

import torch, torch_npu, torchair
input1 = torch.ones(2, 2).npu()
torchair.inference.set_dim_gears(input1, dim_gears={0:[2, 4]})