Home
last modified time | relevance | path

Searched refs:needs_pytree (Results 1 – 6 of 6) sorted by relevance

/aosp_15_r20/external/pytorch/torch/distributed/tensor/
H A D_op_schema.py210 needs_pytree: bool = False variable in RuntimeSchemaInfo
247 if self.schema_info is not None and self.schema_info.needs_pytree
259 if self.schema_info is not None and self.schema_info.needs_pytree
411 and origin_schema.schema_info.needs_pytree
H A D_dispatch.py339 if runtime_schema_info is not None and runtime_schema_info.needs_pytree:
/aosp_15_r20/external/pytorch/torch/distributed/tensor/_ops/
H A D_tensor_ops.py503 @register_op_strategy(aten.stack.default, RuntimeSchemaInfo(1, needs_pytree=True))
538 @register_op_strategy(aten.cat.default, RuntimeSchemaInfo(1, needs_pytree=True))
605 @register_prop_rule(aten.index.Tensor, schema_info=RuntimeSchemaInfo(needs_pytree=True))
H A D_pointwise_ops.py665 register_op_strategy(op, schema_info=RuntimeSchemaInfo(needs_pytree=True))(
670 register_op_strategy(op, schema_info=RuntimeSchemaInfo(needs_pytree=True))(
686 register_op_strategy(op, schema_info=RuntimeSchemaInfo(needs_pytree=True))(
H A D_math_ops.py398 [aten._foreach_norm.Scalar], schema_info=RuntimeSchemaInfo(1, needs_pytree=True)
/aosp_15_r20/external/pytorch/torch/distributed/tensor/experimental/
H A D_register_sharding.py123 static_argnum, static_kwargkey or None, needs_pytree=True